def _check_auth_available(self):
"""
Verifies that information for making API requests to Databricks is available to MLflow,
raising an exception if not.
"""
rest_utils.get_databricks_http_request_kwargs_or_fail(
self.databricks_profile)
def _check_databricks_auth_available():
try:
process.exec_cmd(["databricks", "--version"])
except process.ShellCommandException:
raise ExecutionException(
"Could not find Databricks CLI on PATH. Please install and configure the Databricks "
"CLI as described in https://github.com/databricks/databricks-cli")
# Verify that we can get Databricks auth
rest_utils.get_databricks_http_request_kwargs_or_fail()
def test_databricks_params_throws_errors(get_config_for_profile):
# No hostname
get_config_for_profile.return_value = \
DatabricksConfig(None, "user", "pass", None, insecure=False)
with pytest.raises(Exception):
rest_utils.get_databricks_http_request_kwargs_or_fail()
# No authentication
get_config_for_profile.return_value = \
DatabricksConfig("host", None, None, None, insecure=False)
with pytest.raises(Exception):
rest_utils.get_databricks_http_request_kwargs_or_fail()
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 test_databricks_params_throws_errors(ProfileConfigProvider):
# No hostname
mock_provider = mock.MagicMock()
mock_provider.get_config.return_value = \
DatabricksConfig(None, "user", "pass", None, insecure=True)
ProfileConfigProvider.return_value = mock_provider
with pytest.raises(Exception):
rest_utils.get_databricks_http_request_kwargs_or_fail()
# No authentication
mock_provider = mock.MagicMock()
mock_provider.get_config.return_value = \
DatabricksConfig("host", None, None, None, insecure=True)
ProfileConfigProvider.return_value = mock_provider
with pytest.raises(Exception):
rest_utils.get_databricks_http_request_kwargs_or_fail()
def test_databricks_params_custom_profile(ProfileConfigProvider):
mock_provider = mock.MagicMock()
mock_provider.get_config.return_value = \
DatabricksConfig("host", "user", "pass", None, insecure=True)
ProfileConfigProvider.return_value = mock_provider
params = rest_utils.get_databricks_http_request_kwargs_or_fail("profile")
assert params['verify'] is False
ProfileConfigProvider.assert_called_with("profile")
def _get_databricks_rest_store(store_uri):
parsed_uri = urllib.parse.urlparse(store_uri)
profile = None
if parsed_uri.scheme == 'databricks':
profile = parsed_uri.hostname
http_request_kwargs = rest_utils.get_databricks_http_request_kwargs_or_fail(profile)
return RestStore(http_request_kwargs)
def test_databricks_params_user_password(get_config_for_profile):
get_config_for_profile.return_value = \
DatabricksConfig("host", "user", "pass", None, insecure=False)
params = rest_utils.get_databricks_http_request_kwargs_or_fail()
assert params == {
'hostname': 'host',
'headers': {
'Authorization': 'Basic dXNlcjpwYXNz'
},
'secure_verify': True,
}
def test_databricks_params_token(get_config):
get_config.return_value = \
DatabricksConfig("host", None, None, "mytoken", insecure=False)
params = rest_utils.get_databricks_http_request_kwargs_or_fail()
assert params == {
'hostname': 'host',
'headers': {
'Authorization': 'Basic dG9rZW46bXl0b2tlbg=='
},
'verify': True,
}
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 _print_description_and_log_tags(self):
eprint("=== Launched MLflow run as Databricks job run with ID %s. Getting run status "
"page URL... ===" % self._databricks_run_id)
run_info = self._job_runner.jobs_runs_get(self._databricks_run_id)
jobs_page_url = run_info["run_page_url"]
eprint("=== Check the run's status at %s ===" % jobs_page_url)
tracking.get_service().set_tag(self._mlflow_run_id,
MLFLOW_DATABRICKS_RUN_URL, jobs_page_url)
tracking.get_service().set_tag(self._mlflow_run_id,
MLFLOW_DATABRICKS_SHELL_JOB_RUN_ID, self._databricks_run_id)
tracking.get_service().set_tag(self._mlflow_run_id,
MLFLOW_DATABRICKS_WEBAPP_URL,
get_databricks_http_request_kwargs_or_fail(
profile=self._job_runner.databricks_profile)['hostname'])
job_id = run_info.get('job_id')
# In some releases of Databricks we do not return the job ID. We start including it in DB
# releases 2.80 and above.
if job_id is not None:
tracking.get_service().set_tag(self._mlflow_run_id,
MLFLOW_DATABRICKS_SHELL_JOB_ID, job_id)
def test_databricks_params_custom_profile(get_config_for_profile):
get_config_for_profile.return_value = \
DatabricksConfig("host", "user", "pass", None, insecure=True)
params = rest_utils.get_databricks_http_request_kwargs_or_fail("profile")
assert params['secure_verify'] is False
get_config_for_profile.assert_called_with("profile")
def test_databricks_params_no_verify(get_config_for_profile):
get_config_for_profile.return_value = \
DatabricksConfig("host", "user", "pass", None, insecure=True)
params = rest_utils.get_databricks_http_request_kwargs_or_fail()
assert params['verify'] is False
