def test_create_model_version_run_link_with_configured_profile(
mock_registry_store):
experiment_id = "test-exp-id"
hostname = "https://workspace.databricks.com/"
workspace_id = "10002"
run_id = "runid"
workspace_url = construct_run_url(hostname, experiment_id, run_id,
workspace_id)
get_run_mock = mock.MagicMock()
get_run_mock.return_value = Run(
RunInfo(run_id, experiment_id, "userid", "status", 0, 1, None), None)
with mock.patch(
"mlflow.tracking.client.is_in_databricks_notebook",
return_value=False
), mock.patch(
"mlflow.tracking.client.get_workspace_info_from_databricks_secrets",
return_value=(hostname, workspace_id),
):
client = MlflowClient(tracking_uri="databricks",
registry_uri="otherplace")
client.get_run = get_run_mock
mock_registry_store.create_model_version.return_value = ModelVersion(
"name",
1,
0,
1,
source="source",
run_id=run_id,
run_link=workspace_url)
model_version = client.create_model_version("name", "source", "runid")
assert model_version.run_link == workspace_url
# verify that the client generated the right URL
mock_registry_store.create_model_version.assert_called_once_with(
"name", "source", "runid", [], workspace_url, None)
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 test_create_model_version_run_link_in_notebook_with_default_profile(
mock_registry_store):
experiment_id = 'test-exp-id'
hostname = 'https://workspace.databricks.com/'
workspace_id = '10002'
run_id = 'runid'
workspace_url = construct_run_url(hostname, experiment_id, run_id,
workspace_id)
get_run_mock = mock.MagicMock()
get_run_mock.return_value = Run(
RunInfo(run_id, experiment_id, 'userid', 'status', 0, 1, None), None)
with mock.patch('mlflow.tracking.client.is_in_databricks_notebook',
return_value=True), \
mock.patch('mlflow.tracking.client.get_workspace_info_from_dbutils',
return_value=(hostname, workspace_id)):
client = MlflowClient(tracking_uri='databricks',
registry_uri='otherplace')
client.get_run = get_run_mock
mock_registry_store.create_model_version.return_value = \
ModelVersion('name', 1, 0, 1, source='source', run_id=run_id, run_link=workspace_url)
model_version = client.create_model_version('name', 'source', 'runid')
assert (model_version.run_link == workspace_url)
# verify that the client generated the right URL
mock_registry_store.create_model_version.assert_called_once_with(
"name", 'source', 'runid', [], workspace_url)
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_construct_run_url_errors(hostname, experiment_id, run_id,
workspace_id):
with pytest.raises(
MlflowException,
match="Hostname, experiment ID, and run ID are all required"):
construct_run_url(hostname, experiment_id, run_id, workspace_id)
def test_construct_run_url(hostname, experiment_id, run_id, workspace_id,
result):
assert construct_run_url(hostname, experiment_id, run_id,
workspace_id) == result
def test_construct_run_url_errors(hostname, experiment_id, run_id,
workspace_id):
with pytest.raises(MlflowException):
construct_run_url(hostname, experiment_id, run_id, workspace_id)
