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)