def _get_run_link(self, tracking_uri, run_id):
# if using the default Databricks tracking URI and in a notebook, we can automatically
# figure out the run-link.
if is_databricks_default_tracking_uri(tracking_uri) and (
is_in_databricks_notebook() or is_in_databricks_job()):
# use DBUtils to determine workspace information.
workspace_host, workspace_id = get_workspace_info_from_dbutils()
else:
# in this scenario, we're not able to automatically extract the workspace ID
# to proceed, and users will need to pass in a databricks profile with the scheme:
# databricks://scope:prefix and store the host and workspace-ID as a secret in the
# Databricks Secret Manager with scope=<scope> and key=<prefix>-workspaceid.
workspace_host, workspace_id = get_workspace_info_from_databricks_secrets(
tracking_uri)
if not workspace_id:
print(
"No workspace ID specified; if your Databricks workspaces share the same"
" host URL, you may want to specify the workspace ID (along with the host"
" information in the secret manager) for run lineage tracking. For more"
" details on how to specify this information in the secret manager,"
" please refer to the model registry documentation.")
# retrieve experiment ID of the run for the URL
experiment_id = self.get_run(run_id).info.experiment_id
if workspace_host and run_id and experiment_id:
return construct_run_url(workspace_host, experiment_id, run_id,
workspace_id)
def create_model_version(self,
name,
source,
run_id,
tags=None,
run_link=None):
"""
Create a new model version from given source or run ID.
:param name: Name ID for containing registered model.
:param source: Source path where the MLflow model is stored.
:param run_id: Run ID from MLflow tracking server that generated the model
:param tags: A dictionary of key-value pairs that are converted into
:py:class:`mlflow.entities.model_registry.ModelVersionTag` objects.
:param run_link: Link to the run from an MLflow tracking server that generated this model.
:return: Single :py:class:`mlflow.entities.model_registry.ModelVersion` object created by
backend.
"""
tracking_uri = self._tracking_client.tracking_uri
# for Databricks backends, we support automatically populating the run link field
if is_databricks_uri(
tracking_uri
) and tracking_uri != self._registry_uri and not run_link:
# if using the default Databricks tracking URI and in a notebook, we can automatically
# figure out the run-link.
if is_databricks_default_tracking_uri(
tracking_uri) and is_in_databricks_notebook():
# use DBUtils to determine workspace information.
workspace_host, workspace_id = get_workspace_info_from_dbutils(
)
else:
# in this scenario, we're not able to automatically extract the workspace ID
# to proceed, and users will need to pass in a databricks profile with the scheme:
# databricks://scope/prefix and store the host and workspace-ID as a secret in the
# Databricks Secret Manager with scope=<scope> and key=<prefix>-workspaceid.
workspace_host, workspace_id = \
get_workspace_info_from_databricks_secrets(tracking_uri)
if not workspace_id:
print(
"No workspace ID specified; if your Databricks workspaces share the same"
" host URL, you may want to specify the workspace ID (along with the host"
" information in the secret manager) for run lineage tracking. For more"
" details on how to specify this information in the secret manager,"
" please refer to the model registry documentation.")
# retrieve experiment ID of the run for the URL
experiment_id = self.get_run(run_id).info.experiment_id
if workspace_host and run_id and experiment_id:
run_link = construct_run_url(workspace_host, experiment_id,
run_id, workspace_id)
return self._get_registry_client().create_model_version(
name=name,
source=source,
run_id=run_id,
tags=tags,
run_link=run_link)
def test_get_workspace_info_from_databricks_secrets():
mock_dbutils = mock.MagicMock()
mock_dbutils.secrets.get.return_value = "workspace-placeholder-info"
with mock.patch("mlflow.utils.databricks_utils._get_dbutils", return_value=mock_dbutils):
workspace_host, workspace_id = get_workspace_info_from_databricks_secrets(
"databricks://profile:prefix"
)
mock_dbutils.secrets.get.assert_any_call(key="prefix-host", scope="profile")
mock_dbutils.secrets.get.assert_any_call(key="prefix-workspace-id", scope="profile")
assert workspace_host == "workspace-placeholder-info"
assert workspace_id == "workspace-placeholder-info"
def test_get_workspace_info_from_databricks_secrets():
mock_dbutils = mock.MagicMock()
mock_dbutils.secrets.get.return_value = 'workspace-placeholder-info'
with mock.patch("mlflow.utils.databricks_utils._get_dbutils",
return_value=mock_dbutils):
workspace_host, workspace_id = \
get_workspace_info_from_databricks_secrets('databricks://profile/prefix')
mock_dbutils.secrets.get.assert_any_call(key='prefix-host',
scope='profile')
mock_dbutils.secrets.get.assert_any_call(key='prefix-workspace-id',
scope='profile')
assert workspace_host == 'workspace-placeholder-info'
assert workspace_id == 'workspace-placeholder-info'
