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