def _copy_artifact(local_file, artifact_uri, artifact_path=None):
basename = os.path.basename(local_file)
if artifact_path:
http_endpoint = _get_dbfs_endpoint(
artifact_uri, posixpath.join(artifact_path, basename))
else:
http_endpoint = _get_dbfs_endpoint(artifact_uri, basename)
host_creds = get_databricks_host_creds('registry')
print("Copying file to " + http_endpoint + " in registry workspace")
try:
if os.stat(local_file).st_size == 0:
# The API frontend doesn't like it when we post empty files to it using
# `requests.request`, potentially due to the bug described in
# https://github.com/requests/requests/issues/4215
http_request_safe(host_creds,
endpoint=http_endpoint,
method='POST',
data="",
allow_redirects=False)
else:
with open(local_file, 'rb') as f:
http_request_safe(host_creds,
endpoint=http_endpoint,
method='POST',
data=f,
allow_redirects=False)
except MlflowException as e:
# Note: instead of catching the error here, we could check for the existence of file before trying the copy.
if "File already exists" in e.message:
print("File already exists - continuing to the next file.")
import time
else:
throw(e)
def test_http_request_wrapper(request):
host_only = MlflowHostCreds("http://my-host", ignore_tls_verification=True)
response = mock.MagicMock()
response.status_code = 200
response.text = "{}"
request.return_value = response
http_request_safe(host_only, "/my/endpoint")
request.assert_called_with(
url="http://my-host/my/endpoint",
verify=False,
headers=_DEFAULT_HEADERS,
)
response.text = "non json"
request.return_value = response
http_request_safe(host_only, "/my/endpoint")
request.assert_called_with(
url="http://my-host/my/endpoint",
verify=False,
headers=_DEFAULT_HEADERS,
)
response.status_code = 400
response.text = ""
request.return_value = response
with pytest.raises(MlflowException, match="Response body"):
http_request_safe(host_only, "/my/endpoint")
response.text = (
'{"error_code": "RESOURCE_DOES_NOT_EXIST", "message": "Node type not supported"}'
)
request.return_value = response
with pytest.raises(
RestException,
match="RESOURCE_DOES_NOT_EXIST: Node type not supported"):
http_request_safe(host_only, "/my/endpoint")
def _databricks_api_request(self, endpoint, method, **kwargs):
host_creds = databricks_utils.get_databricks_host_creds(
self.databricks_profile)
return rest_utils.http_request_safe(host_creds=host_creds,
endpoint=endpoint,
method=method,
**kwargs)
def _call_endpoint(self, api, json_body):
endpoint, method = _METHOD_TO_INFO[api]
response_proto = api.Response()
# Convert json string to json dictionary, to pass to requests
if json_body:
json_body = json.loads(json_body)
host_creds = self.get_host_creds()
if method == 'GET':
response = http_request_safe(
host_creds=host_creds, endpoint=endpoint, method=method, params=json_body)
else:
response = http_request_safe(
host_creds=host_creds, endpoint=endpoint, method=method, json=json_body)
js_dict = json.loads(response.text)
parse_dict(js_dict=js_dict, message=response_proto)
return response_proto
def _databricks_api_request(self, endpoint, **kwargs):
host_creds = self.get_host_creds()
return http_request_safe(host_creds=host_creds, endpoint=endpoint, **kwargs)
