def test_http_request_request_headers_user_agent(request):
"""This test requires the package in tests/resources/mlflow-test-plugin to be installed"""
from mlflow_test_plugin.request_header_provider import PluginRequestHeaderProvider
# The test plugin's request header provider always returns False from in_context to avoid
# polluting request headers in developers' environments. The following mock overrides this to
# perform the integration test.
with mock.patch.object(PluginRequestHeaderProvider,
"in_context",
return_value=True), mock.patch.object(
PluginRequestHeaderProvider,
"request_headers",
return_value={_USER_AGENT: "test_user_agent"},
):
host_only = MlflowHostCreds("http://my-host",
server_cert_path="/some/path")
expected_headers = {
_USER_AGENT:
"{} {}".format(
DefaultRequestHeaderProvider().request_headers()[_USER_AGENT],
"test_user_agent")
}
response = mock.MagicMock()
response.status_code = 200
request.return_value = response
http_request(host_only, "/my/endpoint", "GET")
request.assert_called_with(
"GET",
"http://my-host/my/endpoint",
verify="/some/path",
headers=expected_headers,
timeout=120,
)
def test_http_request_request_headers(request):
"""This test requires the package in tests/resources/mlflow-test-plugin to be installed"""
from mlflow_test_plugin.request_header_provider import PluginRequestHeaderProvider
# The test plugin's request header provider always returns False from in_context to avoid
# polluting request headers in developers' environments. The following mock overrides this to
# perform the integration test.
with mock.patch.object(PluginRequestHeaderProvider,
"in_context",
return_value=True):
host_only = MlflowHostCreds("http://my-host",
server_cert_path="/some/path")
response = mock.MagicMock()
response.status_code = 200
request.return_value = response
http_request(host_only, "/my/endpoint")
request.assert_called_with(
url="http://my-host/my/endpoint",
verify="/some/path",
headers={
**_DEFAULT_HEADERS, "test": "header"
},
)
def mlflow_call_endpoint(endpoint, method, body = '{}'):
client = MlflowClient()
host_creds = client._tracking_client.store.get_host_creds()
if method == 'GET':
response = http_request(host_creds=host_creds, endpoint=f"/api/2.0/mlflow/{endpoint}", method=method, params=json.loads(body))
else:
response = http_request(host_creds=host_creds, endpoint=f"/api/2.0/mlflow/{endpoint}", method=method, json=json.loads(body))
return response.json()
def test_http_request_hostonly(request):
host_only = MlflowHostCreds("http://my-host")
response = mock.MagicMock()
response.status_code = 200
request.return_value = response
http_request(host_only, "/my/endpoint")
request.assert_called_with(
url="http://my-host/my/endpoint", verify=True, headers=_DEFAULT_HEADERS,
)
def test_http_request_cleans_hostname(request):
# Add a trailing slash, should be removed.
host_only = MlflowHostCreds("http://my-host/")
response = mock.MagicMock()
response.status_code = 200
request.return_value = response
http_request(host_only, "/my/endpoint")
request.assert_called_with(
url="http://my-host/my/endpoint", verify=True, headers=_DEFAULT_HEADERS,
)
def log_artifact(self, local_file, artifact_path=None):
basename = os.path.basename(local_file)
if artifact_path == '':
raise IllegalArtifactPathError('artifact_path cannot be the empty string.')
if artifact_path:
http_endpoint = self._get_dbfs_endpoint(os.path.join(artifact_path, basename))
else:
http_endpoint = self._get_dbfs_endpoint(os.path.basename(local_file))
with open(local_file, 'rb') as f:
http_request(endpoint=http_endpoint, method='POST', data=f, **self.http_request_kwargs)
def test_429_retries(request):
host_only = MlflowHostCreds("http://my-host", ignore_tls_verification=True)
class MockedResponse(object):
def __init__(self, status_code):
self.status_code = status_code
self.text = "mocked text"
request.side_effect = [MockedResponse(x) for x in (429, 200)]
assert http_request(host_only, "/my/endpoint",
max_rate_limit_interval=0).status_code == 429
request.side_effect = [MockedResponse(x) for x in (429, 200)]
assert http_request(host_only, "/my/endpoint",
max_rate_limit_interval=1).status_code == 200
request.side_effect = [MockedResponse(x) for x in (429, 429, 200)]
assert http_request(host_only, "/my/endpoint",
max_rate_limit_interval=1).status_code == 429
request.side_effect = [MockedResponse(x) for x in (429, 429, 200)]
assert http_request(host_only, "/my/endpoint",
max_rate_limit_interval=2).status_code == 200
request.side_effect = [MockedResponse(x) for x in (429, 429, 200)]
assert http_request(host_only, "/my/endpoint",
max_rate_limit_interval=3).status_code == 200
# Test that any non 429 code is returned
request.side_effect = [MockedResponse(x) for x in (429, 404, 429, 200)]
assert http_request(host_only, "/my/endpoint").status_code == 404
# Test that retries work as expected
request.side_effect = [MockedResponse(x) for x in (429, 503, 429, 200)]
with pytest.raises(MlflowException, match="failed to return code 200"):
http_request(host_only, "/my/endpoint", retries=1)
request.side_effect = [MockedResponse(x) for x in (429, 503, 429, 200)]
assert http_request(host_only, "/my/endpoint",
retries=2).status_code == 200
def test_http_request_with_token(request):
host_only = MlflowHostCreds("http://my-host", token="my-token")
response = mock.MagicMock()
response.status_code = 200
request.return_value = response
http_request(host_only, "/my/endpoint")
headers = dict(_DEFAULT_HEADERS)
headers["Authorization"] = "Bearer my-token"
request.assert_called_with(
url="http://my-host/my/endpoint", verify=True, headers=headers,
)
def test_http_request_with_basic_auth(request):
host_only = MlflowHostCreds("http://my-host", username="******", password="******")
response = mock.MagicMock()
response.status_code = 200
request.return_value = response
http_request(host_only, "/my/endpoint")
headers = dict(_DEFAULT_HEADERS)
headers["Authorization"] = "Basic dXNlcjpwYXNz"
request.assert_called_with(
url="http://my-host/my/endpoint", verify=True, headers=headers,
)
def test_http_request_with_insecure(request):
host_only = MlflowHostCreds("http://my-host", ignore_tls_verification=True)
response = mock.MagicMock()
response.status_code = 200
request.return_value = response
http_request(host_only, '/my/endpoint')
request.assert_called_with(
url='http://my-host/my/endpoint',
verify=False,
headers=_DEFAULT_HEADERS,
)
def test_http_request_server_cert_path(request):
host_only = MlflowHostCreds("http://my-host",
server_cert_path='/some/path')
response = mock.MagicMock()
response.status_code = 200
request.return_value = response
http_request(host_only, '/my/endpoint')
request.assert_called_with(
url='http://my-host/my/endpoint',
verify='/some/path',
headers=_DEFAULT_HEADERS,
)
def _upload_to_dbfs(self, src_path, dbfs_fuse_uri):
"""
Uploads the file at `src_path` to the specified DBFS URI within the Databricks workspace
corresponding to the default Databricks CLI profile.
"""
eprint("=== Uploading project to DBFS path %s ===" % dbfs_fuse_uri)
http_endpoint = dbfs_fuse_uri
http_request_kwargs = \
rest_utils.get_databricks_http_request_kwargs_or_fail(self.databricks_profile)
with open(src_path, 'rb') as f:
rest_utils.http_request(
endpoint=http_endpoint, method='POST', data=f, **http_request_kwargs)
def test_http_request_with_insecure(request):
host_only = MlflowHostCreds("http://my-host", ignore_tls_verification=True)
response = mock.MagicMock()
response.status_code = 200
request.return_value = response
http_request(host_only, "/my/endpoint", "GET")
request.assert_called_with(
"GET",
"http://my-host/my/endpoint",
verify=False,
headers=DefaultRequestHeaderProvider().request_headers(),
timeout=120,
)
def test_http_request_server_cert_path(request):
host_only = MlflowHostCreds("http://my-host", server_cert_path="/some/path")
response = mock.MagicMock()
response.status_code = 200
request.return_value = response
http_request(host_only, "/my/endpoint", "GET")
request.assert_called_with(
"GET",
"http://my-host/my/endpoint",
verify="/some/path",
headers=_DEFAULT_HEADERS,
timeout=120,
)
def test_http_request_with_token(request):
host_only = MlflowHostCreds("http://my-host", token='my-token')
response = mock.MagicMock()
response.status_code = 200
request.return_value = response
http_request(host_only, '/my/endpoint')
request.assert_called_with(
url='http://my-host/my/endpoint',
verify=True,
headers={
'Authorization': 'Bearer my-token'
},
)
def test_http_request_with_basic_auth(request):
host_only = MlflowHostCreds("http://my-host", username='******', password='******')
response = mock.MagicMock()
response.status_code = 200
request.return_value = response
http_request(host_only, '/my/endpoint')
headers = dict(_DEFAULT_HEADERS)
headers['Authorization'] = 'Basic dXNlcjpwYXNz'
request.assert_called_with(
url='http://my-host/my/endpoint',
verify=True,
headers=headers,
)
def test_http_request_cleans_hostname(request):
# Add a trailing slash, should be removed.
host_only = MlflowHostCreds("http://my-host/")
response = mock.MagicMock()
response.status_code = 200
request.return_value = response
http_request(host_only, "/my/endpoint", "GET")
request.assert_called_with(
"GET",
"http://my-host/my/endpoint",
verify=True,
headers=DefaultRequestHeaderProvider().request_headers(),
timeout=120,
)
def log_artifacts(self, local_dir, artifact_path=None):
if artifact_path:
root_http_endpoint = self._get_dbfs_endpoint(artifact_path)
else:
root_http_endpoint = self._get_dbfs_endpoint(os.path.basename(local_dir))
for (dirpath, _, filenames) in os.walk(local_dir):
dir_http_endpoint = root_http_endpoint
if dirpath != local_dir:
rel_path = get_relative_path(local_dir, dirpath)
dir_http_endpoint = build_path(root_http_endpoint, rel_path)
for name in filenames:
endpoint = build_path(dir_http_endpoint, name)
with open(build_path(dirpath, name), 'r') as f:
http_request(endpoint=endpoint, method='POST', data=f,
**self.http_request_kwargs)
def test_http_request_client_cert_path(request):
host_only = MlflowHostCreds("http://my-host",
client_cert_path="/some/path")
response = mock.MagicMock()
response.status_code = 200
request.return_value = response
http_request(host_only, "/my/endpoint", "GET")
request.assert_called_with(
"GET",
"http://my-host/my/endpoint",
verify=True,
cert="/some/path",
headers=DefaultRequestHeaderProvider().request_headers(),
timeout=120,
)
def test_http_request_with_token(request):
host_only = MlflowHostCreds("http://my-host", token="my-token")
response = mock.MagicMock()
response.status_code = 200
request.return_value = response
http_request(host_only, "/my/endpoint", "GET")
headers = DefaultRequestHeaderProvider().request_headers()
headers["Authorization"] = "Bearer my-token"
request.assert_called_with(
"GET",
"http://my-host/my/endpoint",
verify=True,
headers=headers,
timeout=120,
)
def log_artifacts(self, local_dir, artifact_path=None):
if artifact_path:
root_http_endpoint = self._get_dbfs_endpoint(artifact_path)
else:
root_http_endpoint = self._get_dbfs_endpoint(
os.path.basename(local_dir))
for (dirpath, _, filenames) in os.walk(local_dir):
dir_http_endpoint = root_http_endpoint
if dirpath != local_dir:
rel_path = get_relative_path(local_dir, dirpath)
dir_http_endpoint = build_path(root_http_endpoint, rel_path)
for name in filenames:
endpoint = build_path(dir_http_endpoint, name)
with open(build_path(dirpath, name), 'rb') as f:
response = http_request(endpoint=endpoint,
method='POST',
data=f,
allow_redirects=False,
**self.http_request_kwargs)
if response.status_code == 409:
raise MlflowException(
'File already exists at {} and can\'t be overwritten.'.
format(endpoint))
elif response.status_code != 200:
raise MlflowException(
'log_artifacts to "{}" returned a non-200 status code.'
.format(endpoint))
def databricks_api_request(self, endpoint, method, **kwargs):
request_params = rest_utils.get_databricks_http_request_kwargs_or_fail(
self.databricks_profile)
request_params.update(kwargs)
response = rest_utils.http_request(
endpoint=endpoint, method=method, **request_params)
return json.loads(response.text)
def _databricks_api_request(self, endpoint, method, **kwargs):
host_creds = databricks_utils.get_databricks_host_creds(
self.databricks_profile)
return rest_utils.http_request(host_creds=host_creds,
endpoint=endpoint,
method=method,
**kwargs)
def _dbfs_path_exists(self, dbfs_path):
"""
Return True if the passed-in path exists in DBFS for the workspace corresponding to the
default Databricks CLI profile. The path is expected to be a relative path to the DBFS root
directory, e.g. 'path/to/file'.
"""
host_creds = databricks_utils.get_databricks_host_creds(self.databricks_profile_uri)
response = rest_utils.http_request(
host_creds=host_creds,
endpoint="/api/2.0/dbfs/get-status",
method="GET",
json={"path": "/%s" % dbfs_path},
)
try:
json_response_obj = json.loads(response.text)
except Exception:
raise MlflowException(
"API request to check existence of file at DBFS path %s failed with status code "
"%s. Response body: %s" % (dbfs_path, response.status_code, response.text)
)
# If request fails with a RESOURCE_DOES_NOT_EXIST error, the file does not exist on DBFS
error_code_field = "error_code"
if error_code_field in json_response_obj:
if json_response_obj[error_code_field] == "RESOURCE_DOES_NOT_EXIST":
return False
raise ExecutionException(
"Got unexpected error response when checking whether file %s "
"exists in DBFS: %s" % (dbfs_path, json_response_obj)
)
return True
def log_artifact(self, local_file, artifact_path=None):
basename = os.path.basename(local_file)
if artifact_path == '':
raise IllegalArtifactPathError(
'artifact_path cannot be the empty string.')
if artifact_path:
http_endpoint = self._get_dbfs_endpoint(
os.path.join(artifact_path, basename))
else:
http_endpoint = self._get_dbfs_endpoint(
os.path.basename(local_file))
with open(local_file, 'rb') as f:
response = http_request(endpoint=http_endpoint,
method='POST',
data=f,
allow_redirects=False,
**self.http_request_kwargs)
if response.status_code == 409:
raise MlflowException(
'File already exists at {} and can\'t be overwritten.'.
format(http_endpoint))
elif response.status_code != 200:
raise MlflowException(
'log_artifact to "{}" returned a non-200 status code.'.
format(http_endpoint))
def databricks_api_request(self, endpoint, method, **kwargs):
host_creds = databricks_utils.get_databricks_host_creds(
self.databricks_profile)
response = rest_utils.http_request(host_creds=host_creds,
endpoint=endpoint,
method=method,
**kwargs)
return json.loads(response.text)
def test_http_request_with_basic_auth(request):
host_only = MlflowHostCreds("http://my-host",
username="******",
password="******")
response = mock.MagicMock()
response.status_code = 200
request.return_value = response
http_request(host_only, "/my/endpoint", "GET")
headers = DefaultRequestHeaderProvider().request_headers()
headers["Authorization"] = "Basic dXNlcjpwYXNz"
request.assert_called_with(
"GET",
"http://my-host/my/endpoint",
verify=True,
headers=headers,
timeout=120,
)
def _dbfs_list_api(json, http_request_kwargs):
"""
Pulled out to make it easier to mock.
"""
return http_request(endpoint=LIST_API_ENDPOINT,
method='GET',
json=json,
**http_request_kwargs)
def _dbfs_is_dir(dbfs_path, http_request_kwargs):
response = http_request(endpoint=GET_STATUS_ENDPOINT, method='GET',
json={'path': dbfs_path}, **http_request_kwargs)
json_response = json.loads(response.text)
try:
return json_response['is_dir']
except KeyError:
raise Exception('DBFS path %s does not exist' % dbfs_path)
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(
host_creds=host_creds, endpoint=endpoint, method=method, params=json_body)
else:
response = http_request(
host_creds=host_creds, endpoint=endpoint, method=method, json=json_body)
response = self._verify_rest_response(response, endpoint)
js_dict = json.loads(response.text)
parse_dict(js_dict=js_dict, message=response_proto)
return response_proto
