def test_databricks_notebook_run_context_tags_nones(): patch_notebook_id = mock.patch( "mlflow.utils.databricks_utils.get_notebook_id", return_value=None) patch_notebook_path = mock.patch( "mlflow.utils.databricks_utils.get_notebook_path", return_value=None) patch_webapp_url = mock.patch( "mlflow.utils.databricks_utils.get_webapp_url", return_value=None) with patch_notebook_id, patch_notebook_path, patch_webapp_url: assert DatabricksNotebookRunContext().tags() == { MLFLOW_SOURCE_NAME: None, MLFLOW_SOURCE_TYPE: SourceType.to_string(SourceType.NOTEBOOK), }
def test_run_local_git_repo( local_git_repo, local_git_repo_uri, tracking_uri_mock, # pylint: disable=unused-argument use_start_run, version): if version is not None: uri = local_git_repo_uri + "#" + TEST_PROJECT_NAME else: uri = os.path.join("%s/" % local_git_repo, TEST_PROJECT_NAME) if version == "git-commit": version = _get_version_local_git_repo(local_git_repo) submitted_run = mlflow.projects.run( uri, entry_point="test_tracking", version=version, parameters={"use_start_run": use_start_run}, use_conda=False, experiment_id=0) # Blocking runs should be finished when they return validate_exit_status(submitted_run.get_status(), RunStatus.FINISHED) # Test that we can call wait() on a synchronous run & that the run has the correct # status after calling wait(). submitted_run.wait() validate_exit_status(submitted_run.get_status(), RunStatus.FINISHED) # Validate run contents in the FileStore run_uuid = submitted_run.run_id mlflow_service = mlflow.tracking.MlflowClient() run_infos = mlflow_service.list_run_infos( experiment_id=0, run_view_type=ViewType.ACTIVE_ONLY) assert "file:" in run_infos[0].source_name assert len(run_infos) == 1 store_run_uuid = run_infos[0].run_uuid assert run_uuid == store_run_uuid run = mlflow_service.get_run(run_uuid) assert run.info.status == RunStatus.FINISHED assert run.data.params == {"use_start_run": use_start_run} assert run.data.metrics == {"some_key": 3} tags = run.data.tags assert "file:" in tags[MLFLOW_SOURCE_NAME] assert tags[MLFLOW_SOURCE_TYPE] == SourceType.to_string(SourceType.PROJECT) assert tags[MLFLOW_PROJECT_ENTRY_POINT] == "test_tracking" if version == "master": assert tags[MLFLOW_GIT_BRANCH] == "master" assert tags[MLFLOW_GIT_REPO_URL] == local_git_repo_uri assert tags[LEGACY_MLFLOW_GIT_BRANCH_NAME] == "master" assert tags[LEGACY_MLFLOW_GIT_REPO_URL] == local_git_repo_uri
def test_start_run_defaults_databricks_notebook(empty_active_run_stack): mock_experiment_id = mock.Mock() experiment_id_patch = mock.patch( "mlflow.tracking.fluent._get_experiment_id", return_value=mock_experiment_id ) databricks_notebook_patch = mock.patch( "mlflow.utils.databricks_utils.is_in_databricks_notebook", return_value=True ) mock_user = mock.Mock() user_patch = mock.patch( "mlflow.tracking.context.default_context._get_user", return_value=mock_user ) mock_source_version = mock.Mock() source_version_patch = mock.patch( "mlflow.tracking.context.git_context._get_source_version", return_value=mock_source_version ) mock_notebook_id = mock.Mock() notebook_id_patch = mock.patch( "mlflow.utils.databricks_utils.get_notebook_id", return_value=mock_notebook_id ) mock_notebook_path = mock.Mock() notebook_path_patch = mock.patch( "mlflow.utils.databricks_utils.get_notebook_path", return_value=mock_notebook_path ) mock_webapp_url = mock.Mock() webapp_url_patch = mock.patch( "mlflow.utils.databricks_utils.get_webapp_url", return_value=mock_webapp_url ) expected_tags = { mlflow_tags.MLFLOW_USER: mock_user, mlflow_tags.MLFLOW_SOURCE_NAME: mock_notebook_path, mlflow_tags.MLFLOW_SOURCE_TYPE: SourceType.to_string(SourceType.NOTEBOOK), mlflow_tags.MLFLOW_GIT_COMMIT: mock_source_version, mlflow_tags.MLFLOW_DATABRICKS_NOTEBOOK_ID: mock_notebook_id, mlflow_tags.MLFLOW_DATABRICKS_NOTEBOOK_PATH: mock_notebook_path, mlflow_tags.MLFLOW_DATABRICKS_WEBAPP_URL: mock_webapp_url } create_run_patch = mock.patch.object(MlflowClient, "create_run") with experiment_id_patch, databricks_notebook_patch, user_patch, source_version_patch, \ notebook_id_patch, notebook_path_patch, webapp_url_patch, create_run_patch: active_run = start_run() MlflowClient.create_run.assert_called_once_with( experiment_id=mock_experiment_id, tags=expected_tags ) assert is_from_run(active_run, MlflowClient.create_run.return_value)
def test_databricks_job_run_context_tags_nones(): patch_job_id = mock.patch("mlflow.utils.databricks_utils.get_job_id", return_value=None) patch_job_run_id = mock.patch("mlflow.utils.databricks_utils.get_job_run_id", return_value=None) patch_job_type = mock.patch("mlflow.utils.databricks_utils.get_job_type", return_value=None) patch_webapp_url = mock.patch("mlflow.utils.databricks_utils.get_webapp_url", return_value=None) patch_workspace_info = mock.patch( "mlflow.utils.databricks_utils.get_workspace_info_from_dbutils", return_value=(None, None) ) with patch_job_id, patch_job_run_id, patch_job_type, patch_webapp_url, patch_workspace_info: assert DatabricksJobRunContext().tags() == { MLFLOW_SOURCE_NAME: None, MLFLOW_SOURCE_TYPE: SourceType.to_string(SourceType.JOB), }
def test_start_run_overrides_databricks_notebook(empty_active_run_stack): databricks_notebook_patch = mock.patch( "mlflow.utils.databricks_utils.is_in_databricks_notebook", return_value=True) mock_notebook_id = mock.Mock() notebook_id_patch = mock.patch( "mlflow.utils.databricks_utils.get_notebook_id", return_value=mock_notebook_id) mock_notebook_path = mock.Mock() notebook_path_patch = mock.patch( "mlflow.utils.databricks_utils.get_notebook_path", return_value=mock_notebook_path) mock_webapp_url = mock.Mock() webapp_url_patch = mock.patch( "mlflow.utils.databricks_utils.get_webapp_url", return_value=mock_webapp_url) mock_experiment_id = mock.Mock() mock_source_name = mock.Mock() source_type = SourceType.JOB mock_source_version = mock.Mock() mock_entry_point_name = mock.Mock() mock_run_name = mock.Mock() expected_tags = { mlflow_tags.MLFLOW_SOURCE_NAME: mock_source_name, mlflow_tags.MLFLOW_SOURCE_TYPE: SourceType.to_string(source_type), mlflow_tags.MLFLOW_GIT_COMMIT: mock_source_version, mlflow_tags.MLFLOW_PROJECT_ENTRY_POINT: mock_entry_point_name, mlflow_tags.MLFLOW_DATABRICKS_NOTEBOOK_ID: mock_notebook_id, mlflow_tags.MLFLOW_DATABRICKS_NOTEBOOK_PATH: mock_notebook_path, mlflow_tags.MLFLOW_DATABRICKS_WEBAPP_URL: mock_webapp_url } create_run_patch = mock.patch.object(MlflowClient, "create_run") with databricks_notebook_patch, notebook_id_patch, notebook_path_patch, webapp_url_patch, \ create_run_patch: active_run = start_run(experiment_id=mock_experiment_id, source_name=mock_source_name, source_version=mock_source_version, entry_point_name=mock_entry_point_name, source_type=source_type, run_name=mock_run_name) MlflowClient.create_run.assert_called_once_with( experiment_id=mock_experiment_id, run_name=mock_run_name, tags=expected_tags) assert is_from_run(active_run, MlflowClient.create_run.return_value)
def test_start_run_creates_new_run_with_user_specified_tags(): mock_experiment_id = mock.Mock() experiment_id_patch = mock.patch( "mlflow.tracking.fluent._get_experiment_id", return_value=mock_experiment_id) mock_user = mock.Mock() user_patch = mock.patch( "mlflow.tracking.context.default_context._get_user", return_value=mock_user) mock_source_name = mock.Mock() source_name_patch = mock.patch( "mlflow.tracking.context.default_context._get_source_name", return_value=mock_source_name) source_type_patch = mock.patch( "mlflow.tracking.context.default_context._get_source_type", return_value=SourceType.NOTEBOOK) mock_source_version = mock.Mock() source_version_patch = mock.patch( "mlflow.tracking.context.git_context._get_source_version", return_value=mock_source_version) user_specified_tags = { "ml_task": "regression", "num_layers": 7, mlflow_tags.MLFLOW_USER: "******", } expected_tags = { mlflow_tags.MLFLOW_SOURCE_NAME: mock_source_name, mlflow_tags.MLFLOW_SOURCE_TYPE: SourceType.to_string(SourceType.NOTEBOOK), mlflow_tags.MLFLOW_GIT_COMMIT: mock_source_version, mlflow_tags.MLFLOW_USER: "******", "ml_task": "regression", "num_layers": 7, } create_run_patch = mock.patch.object(MlflowClient, "create_run") with multi_context( experiment_id_patch, user_patch, source_name_patch, source_type_patch, source_version_patch, create_run_patch, ): active_run = start_run(tags=user_specified_tags) MlflowClient.create_run.assert_called_once_with( experiment_id=mock_experiment_id, tags=expected_tags) assert is_from_run(active_run, MlflowClient.create_run.return_value)
def test_databricks_notebook_run_context_tags(): patch_notebook_id = mock.patch("mlflow.utils.databricks_utils.get_notebook_id") patch_notebook_path = mock.patch("mlflow.utils.databricks_utils.get_notebook_path") patch_webapp_url = mock.patch("mlflow.utils.databricks_utils.get_webapp_url") with patch_notebook_id as notebook_id_mock, patch_notebook_path as notebook_path_mock, \ patch_webapp_url as webapp_url_mock: assert DatabricksNotebookRunContext().tags() == { MLFLOW_SOURCE_NAME: notebook_path_mock.return_value, MLFLOW_SOURCE_TYPE: SourceType.to_string(SourceType.NOTEBOOK), MLFLOW_DATABRICKS_NOTEBOOK_ID: notebook_id_mock.return_value, MLFLOW_DATABRICKS_NOTEBOOK_PATH: notebook_path_mock.return_value, MLFLOW_DATABRICKS_WEBAPP_URL: webapp_url_mock.return_value }
def _create_run(uri, experiment_id, work_dir, version, entry_point, parameters): """ Create a ``Run`` against the current MLflow tracking server, logging metadata (e.g. the URI, entry point, and parameters of the project) about the run. Return an ``ActiveRun`` that can be used to report additional data about the run (metrics/params) to the tracking server. """ if _is_local_uri(uri): source_name = tracking._tracking_service.utils._get_git_url_if_present(_expand_uri(uri)) else: source_name = _expand_uri(uri) source_version = _get_git_commit(work_dir) existing_run = fluent.active_run() if existing_run: parent_run_id = existing_run.info.run_id else: parent_run_id = None tags = { MLFLOW_USER: _get_user(), MLFLOW_SOURCE_NAME: source_name, MLFLOW_SOURCE_TYPE: SourceType.to_string(SourceType.PROJECT), MLFLOW_PROJECT_ENTRY_POINT: entry_point, } if source_version is not None: tags[MLFLOW_GIT_COMMIT] = source_version if parent_run_id is not None: tags[MLFLOW_PARENT_RUN_ID] = parent_run_id repo_url = _get_git_repo_url(work_dir) if repo_url is not None: tags[MLFLOW_GIT_REPO_URL] = repo_url tags[LEGACY_MLFLOW_GIT_REPO_URL] = repo_url # Add branch name tag if a branch is specified through -version if _is_valid_branch_name(work_dir, version): tags[MLFLOW_GIT_BRANCH] = version tags[LEGACY_MLFLOW_GIT_BRANCH_NAME] = version active_run = tracking.MlflowClient().create_run(experiment_id=experiment_id, tags=tags) project = _project_spec.load_project(work_dir) # Consolidate parameters for logging. # `storage_dir` is `None` since we want to log actual path not downloaded local path entry_point_obj = project.get_entry_point(entry_point) final_params, extra_params = entry_point_obj.compute_parameters(parameters, storage_dir=None) params_list = [ Param(key, value) for key, value in list(final_params.items()) + list(extra_params.items()) ] tracking.MlflowClient().log_batch(active_run.info.run_id, params=params_list) return active_run
def test_create_run(self): experiment_id = self._experiment_factory('test_create_run') expected = self._get_run_configs('booyya', experiment_id=experiment_id) tags = [RunTag('3', '4'), RunTag('1', '2')] actual = self.store.create_run(expected["experiment_id"], expected["user_id"], expected["name"], SourceType.from_string(expected["source_type"]), expected["source_name"], expected["entry_point_name"], expected["start_time"], expected["source_version"], tags, None) self.assertEqual(actual.info.experiment_id, expected["experiment_id"]) self.assertEqual(actual.info.user_id, expected["user_id"]) self.assertEqual(actual.info.name, 'booyya') self.assertEqual(actual.info.source_type, SourceType.from_string(expected["source_type"])) self.assertEqual(actual.info.source_name, expected["source_name"]) self.assertEqual(actual.info.source_version, expected["source_version"]) self.assertEqual(actual.info.entry_point_name, expected["entry_point_name"]) self.assertEqual(actual.info.start_time, expected["start_time"]) self.assertEqual(len(actual.data.tags), 3) name_tag = models.SqlTag(key='mlflow.runName', value='booyya').to_mlflow_entity() self.assertListEqual(actual.data.tags, tags + [name_tag])
def tags(self): notebook_id = databricks_utils.get_notebook_id() notebook_path = databricks_utils.get_notebook_path() webapp_url = databricks_utils.get_webapp_url() tags = { MLFLOW_SOURCE_NAME: notebook_path, MLFLOW_SOURCE_TYPE: SourceType.to_string(SourceType.NOTEBOOK) } if notebook_id is not None: tags[MLFLOW_DATABRICKS_NOTEBOOK_ID] = notebook_id if notebook_path is not None: tags[MLFLOW_DATABRICKS_NOTEBOOK_PATH] = notebook_path if webapp_url is not None: tags[MLFLOW_DATABRICKS_WEBAPP_URL] = webapp_url return tags
def create_run(self, experiment_id, user_id=None, run_name=None, start_time=None, parent_run_id=None, tags=None): """ Create a :py:class:`mlflow.entities.Run` object that can be associated with metrics, parameters, artifacts, etc. Unlike :py:func:`mlflow.projects.run`, creates objects but does not run code. Unlike :py:func:`mlflow.start_run`, does not change the "active run" used by :py:func:`mlflow.log_param`. :param user_id: If not provided, use the current user as a default. :param start_time: If not provided, use the current timestamp. :param parent_run_id Optional parent run ID - takes precedence over parent run ID included in the `tags` argument. :param tags: A dictionary of key-value pairs that are converted into :py:class:`mlflow.entities.RunTag` objects. :return: :py:class:`mlflow.entities.Run` that was created. """ tags = tags if tags else {} # Extract run attributes from tags # This logic is temporary; by the 1.0 release, this information will only be stored in tags # and will not be available as attributes of the run final_parent_run_id =\ tags.get(MLFLOW_PARENT_RUN_ID) if parent_run_id is None else parent_run_id source_name = tags.get(MLFLOW_SOURCE_NAME, "Python Application") source_version = tags.get(MLFLOW_GIT_COMMIT) entry_point_name = tags.get(MLFLOW_PROJECT_ENTRY_POINT) source_type_string = tags.get(MLFLOW_SOURCE_TYPE) if source_type_string is None: source_type = SourceType.LOCAL else: source_type = SourceType.from_string(source_type_string) return self.store.create_run( experiment_id=experiment_id, user_id=user_id if user_id is not None else _get_user_id(), run_name=run_name, start_time=start_time or int(time.time() * 1000), tags=[RunTag(key, value) for (key, value) in iteritems(tags)], # The below arguments remain set for backwards compatability: parent_run_id=final_parent_run_id, source_type=source_type, source_name=source_name, entry_point_name=entry_point_name, source_version=source_version )
def test_start_run_defaults(empty_active_run_stack): # pylint: disable=unused-argument mock_experiment_id = mock.Mock() experiment_id_patch = mock.patch( "mlflow.tracking.fluent._get_experiment_id", return_value=mock_experiment_id ) databricks_notebook_patch = mock.patch( "mlflow.tracking.fluent.is_in_databricks_notebook", return_value=False ) mock_user = mock.Mock() user_patch = mock.patch( "mlflow.tracking.context.default_context._get_user", return_value=mock_user ) mock_source_name = mock.Mock() source_name_patch = mock.patch( "mlflow.tracking.context.default_context._get_source_name", return_value=mock_source_name ) source_type_patch = mock.patch( "mlflow.tracking.context.default_context._get_source_type", return_value=SourceType.NOTEBOOK ) mock_source_version = mock.Mock() source_version_patch = mock.patch( "mlflow.tracking.context.git_context._get_source_version", return_value=mock_source_version ) expected_tags = { mlflow_tags.MLFLOW_USER: mock_user, mlflow_tags.MLFLOW_SOURCE_NAME: mock_source_name, mlflow_tags.MLFLOW_SOURCE_TYPE: SourceType.to_string(SourceType.NOTEBOOK), mlflow_tags.MLFLOW_GIT_COMMIT: mock_source_version, } create_run_patch = mock.patch.object(MlflowClient, "create_run") with multi_context( experiment_id_patch, databricks_notebook_patch, user_patch, source_name_patch, source_type_patch, source_version_patch, create_run_patch, ): active_run = start_run() MlflowClient.create_run.assert_called_once_with( experiment_id=mock_experiment_id, tags=expected_tags ) assert is_from_run(active_run, MlflowClient.create_run.return_value)
def _get_run_configs(self, name='test', experiment_id=None): return { 'experiment_id': experiment_id, 'name': name, 'user_id': 'Anderson', 'run_uuid': uuid.uuid4().hex, 'status': RunStatus.to_string(RunStatus.SCHEDULED), 'source_type': SourceType.to_string(SourceType.NOTEBOOK), 'source_name': 'Python application', 'entry_point_name': 'main.py', 'start_time': int(time.time()), 'end_time': int(time.time()), 'source_version': mlflow.__version__, 'lifecycle_stage': entities.LifecycleStage.ACTIVE, 'artifact_uri': '//' }
def test_databricks_job_run_context_tags(): patch_job_id = mock.patch("mlflow.utils.databricks_utils.get_job_id") patch_job_run_id = mock.patch("mlflow.utils.databricks_utils.get_job_run_id") patch_job_type = mock.patch("mlflow.utils.databricks_utils.get_job_type") patch_webapp_url = mock.patch("mlflow.utils.databricks_utils.get_webapp_url") with patch_job_id as job_id_mock, patch_job_run_id as job_run_id_mock, patch_job_type as job_type_mock, patch_webapp_url as webapp_url_mock: # noqa assert DatabricksJobRunContext().tags() == { MLFLOW_SOURCE_NAME: "jobs/{job_id}/run/{job_run_id}".format( job_id=job_id_mock.return_value, job_run_id=job_run_id_mock.return_value ), MLFLOW_SOURCE_TYPE: SourceType.to_string(SourceType.JOB), MLFLOW_DATABRICKS_JOB_ID: job_id_mock.return_value, MLFLOW_DATABRICKS_JOB_RUN_ID: job_run_id_mock.return_value, MLFLOW_DATABRICKS_JOB_TYPE: job_type_mock.return_value, MLFLOW_DATABRICKS_WEBAPP_URL: webapp_url_mock.return_value, }
def test_client_create_run_overrides(mock_store): experiment_id = mock.Mock() user_id = mock.Mock() run_name = mock.Mock() start_time = mock.Mock() tags = { MLFLOW_PARENT_RUN_ID: mock.Mock(), MLFLOW_SOURCE_TYPE: SourceType.to_string(SourceType.JOB), MLFLOW_SOURCE_NAME: mock.Mock(), MLFLOW_PROJECT_ENTRY_POINT: mock.Mock(), MLFLOW_GIT_COMMIT: mock.Mock(), "other-key": "other-value" } MlflowClient().create_run(experiment_id, user_id, run_name, start_time, None, tags) mock_store.create_run.assert_called_once_with( experiment_id=experiment_id, user_id=user_id, run_name=run_name, start_time=start_time, tags=[RunTag(key, value) for key, value in tags.items()], parent_run_id=tags[MLFLOW_PARENT_RUN_ID], source_type=SourceType.JOB, source_name=tags[MLFLOW_SOURCE_NAME], entry_point_name=tags[MLFLOW_PROJECT_ENTRY_POINT], source_version=tags[MLFLOW_GIT_COMMIT]) mock_store.reset_mock() parent_run_id = "mock-parent-run-id" MlflowClient().create_run(experiment_id, user_id, run_name, start_time, parent_run_id, tags) mock_store.create_run.assert_called_once_with( experiment_id=experiment_id, user_id=user_id, run_name=run_name, start_time=start_time, tags=[RunTag(key, value) for key, value in tags.items()], parent_run_id=parent_run_id, source_type=SourceType.JOB, source_name=tags[MLFLOW_SOURCE_NAME], entry_point_name=tags[MLFLOW_PROJECT_ENTRY_POINT], source_version=tags[MLFLOW_GIT_COMMIT])
def create_run(self, experiment_id, user_id, run_name, source_type, source_name, entry_point_name, start_time, source_version, tags, parent_run_id): with self.ManagedSessionMaker() as session: experiment = self.get_experiment(experiment_id) if experiment.lifecycle_stage != LifecycleStage.ACTIVE: raise MlflowException( 'Experiment id={} must be active'.format(experiment_id), INVALID_STATE) run_uuid = uuid.uuid4().hex artifact_location = build_path( experiment.artifact_location, run_uuid, SqlAlchemyStore.ARTIFACTS_FOLDER_NAME) run = SqlRun(name=run_name or "", artifact_uri=artifact_location, run_uuid=run_uuid, experiment_id=experiment_id, source_type=SourceType.to_string(source_type), source_name=source_name, entry_point_name=entry_point_name, user_id=user_id, status=RunStatus.to_string(RunStatus.RUNNING), start_time=start_time, end_time=None, source_version=source_version, lifecycle_stage=LifecycleStage.ACTIVE) tags_dict = {} for tag in tags: tags_dict[tag.key] = tag.value if parent_run_id: tags_dict[MLFLOW_PARENT_RUN_ID] = parent_run_id if run_name: tags_dict[MLFLOW_RUN_NAME] = run_name run.tags = [ SqlTag(key=key, value=value) for key, value in tags_dict.items() ] self._save_to_db(objs=run, session=session) return run.to_mlflow_entity()
def _create_entity(base, model): # create dict of kwargs properties for entity and return the initialized entity config = {} for k in base._properties(): # check if its mlflow entity and build it obj = getattr(model, k) if isinstance(model, SqlRun): if base is RunData: # Run data contains list for metrics, params and tags # so obj will be a list so we need to convert those items if k == 'metrics': # only get latest recorded metrics per key metrics = {} for o in obj: existing_metric = metrics.get(o.key) if (existing_metric is None)\ or ((o.step, o.timestamp, o.value) >= (existing_metric.step, existing_metric.timestamp, existing_metric.value)): metrics[o.key] = Metric(o.key, o.value, o.timestamp, o.step) obj = list(metrics.values()) elif k == 'params': obj = [Param(o.key, o.value) for o in obj] elif k == 'tags': obj = [RunTag(o.key, o.value) for o in obj] elif base is RunInfo: if k == 'source_type': obj = SourceType.from_string(obj) elif k == "status": obj = RunStatus.from_string(obj) elif k == "experiment_id": obj = str(obj) # Our data model defines experiment_ids as ints, but the in-memory representation was # changed to be a string in time for 1.0. if isinstance(model, SqlExperiment) and k == "experiment_id": obj = str(obj) config[k] = obj return base(**config)
def test_run( tmpdir, # pylint: disable=unused-argument patch_user, # pylint: disable=unused-argument tracking_uri_mock, # pylint: disable=unused-argument use_start_run): submitted_run = mlflow.projects.run( TEST_PROJECT_DIR, entry_point="test_tracking", parameters={"use_start_run": use_start_run}, use_conda=False, experiment_id=FileStore.DEFAULT_EXPERIMENT_ID) assert submitted_run.run_id is not None # Blocking runs should be finished when they return validate_exit_status(submitted_run.get_status(), RunStatus.FINISHED) # Test that we can call wait() on a synchronous run & that the run has the correct # status after calling wait(). submitted_run.wait() validate_exit_status(submitted_run.get_status(), RunStatus.FINISHED) # Validate run contents in the FileStore run_id = submitted_run.run_id mlflow_service = mlflow.tracking.MlflowClient() run_infos = mlflow_service.list_run_infos( experiment_id=FileStore.DEFAULT_EXPERIMENT_ID, run_view_type=ViewType.ACTIVE_ONLY) assert len(run_infos) == 1 store_run_id = run_infos[0].run_id assert run_id == store_run_id run = mlflow_service.get_run(run_id) assert run.info.status == RunStatus.to_string(RunStatus.FINISHED) assert run.data.params == {"use_start_run": use_start_run} assert run.data.metrics == {"some_key": 3} tags = run.data.tags assert tags[MLFLOW_USER] == MOCK_USER assert "file:" in tags[MLFLOW_SOURCE_NAME] assert tags[MLFLOW_SOURCE_TYPE] == SourceType.to_string(SourceType.PROJECT) assert tags[MLFLOW_PROJECT_ENTRY_POINT] == "test_tracking"
def test_start_run_with_parent(): parent_run = mock.Mock() mock_experiment_id = mock.Mock() mock_source_name = mock.Mock() active_run_stack_patch = mock.patch("mlflow.tracking.fluent._active_run_stack", [parent_run]) databricks_notebook_patch = mock.patch( "mlflow.tracking.fluent.is_in_databricks_notebook", return_value=False ) mock_user = mock.Mock() user_patch = mock.patch( "mlflow.tracking.context.default_context._get_user", return_value=mock_user ) source_name_patch = mock.patch( "mlflow.tracking.context.default_context._get_source_name", return_value=mock_source_name ) expected_tags = { mlflow_tags.MLFLOW_USER: mock_user, mlflow_tags.MLFLOW_SOURCE_NAME: mock_source_name, mlflow_tags.MLFLOW_SOURCE_TYPE: SourceType.to_string(SourceType.LOCAL), mlflow_tags.MLFLOW_PARENT_RUN_ID: parent_run.info.run_id, } create_run_patch = mock.patch.object(MlflowClient, "create_run") with multi_context( databricks_notebook_patch, active_run_stack_patch, create_run_patch, user_patch, source_name_patch, ): active_run = start_run(experiment_id=mock_experiment_id, nested=True) MlflowClient.create_run.assert_called_once_with( experiment_id=mock_experiment_id, tags=expected_tags ) assert is_from_run(active_run, MlflowClient.create_run.return_value)
def test_databricks_job_run_context_tags(): patch_job_id = mock.patch("mlflow.utils.databricks_utils.get_job_id") patch_job_run_id = mock.patch( "mlflow.utils.databricks_utils.get_job_run_id") patch_job_type = mock.patch("mlflow.utils.databricks_utils.get_job_type") patch_webapp_url = mock.patch( "mlflow.utils.databricks_utils.get_webapp_url") patch_workspace_info = mock.patch( "mlflow.utils.databricks_utils.get_workspace_info_from_dbutils", return_value=("https://databricks.com", "123456"), ) with multi_context(patch_job_id, patch_job_run_id, patch_job_type, patch_webapp_url, patch_workspace_info) as ( job_id_mock, job_run_id_mock, job_type_mock, webapp_url_mock, workspace_info_mock, ): assert DatabricksJobRunContext().tags() == { MLFLOW_SOURCE_NAME: "jobs/{job_id}/run/{job_run_id}".format( job_id=job_id_mock.return_value, job_run_id=job_run_id_mock.return_value), MLFLOW_SOURCE_TYPE: SourceType.to_string(SourceType.JOB), MLFLOW_DATABRICKS_JOB_ID: job_id_mock.return_value, MLFLOW_DATABRICKS_JOB_RUN_ID: job_run_id_mock.return_value, MLFLOW_DATABRICKS_JOB_TYPE: job_type_mock.return_value, MLFLOW_DATABRICKS_WEBAPP_URL: webapp_url_mock.return_value, MLFLOW_DATABRICKS_WORKSPACE_URL: workspace_info_mock.return_value[0], MLFLOW_DATABRICKS_WORKSPACE_ID: workspace_info_mock.return_value[1], }
def tags(self): job_id = databricks_utils.get_job_id() job_run_id = databricks_utils.get_job_run_id() job_type = databricks_utils.get_job_type() webapp_url = databricks_utils.get_webapp_url() tags = { MLFLOW_SOURCE_NAME: ("jobs/{job_id}/run/{job_run_id}".format(job_id=job_id, job_run_id=job_run_id) if job_id is not None and job_run_id is not None else None), MLFLOW_SOURCE_TYPE: SourceType.to_string(SourceType.JOB), } if job_id is not None: tags[MLFLOW_DATABRICKS_JOB_ID] = job_id if job_run_id is not None: tags[MLFLOW_DATABRICKS_JOB_RUN_ID] = job_run_id if job_type is not None: tags[MLFLOW_DATABRICKS_JOB_TYPE] = job_type if webapp_url is not None: tags[MLFLOW_DATABRICKS_WEBAPP_URL] = webapp_url return tags
def test_start_run_with_parent(): parent_run = mock.Mock() active_run_stack_patch = mock.patch( "mlflow.tracking.fluent._active_run_stack", [parent_run]) databricks_notebook_patch = mock.patch( "mlflow.tracking.fluent.is_in_databricks_notebook", return_value=False) mock_experiment_id = mock.Mock() mock_source_name = mock.Mock() source_type = SourceType.JOB mock_source_version = mock.Mock() mock_entry_point_name = mock.Mock() mock_run_name = mock.Mock() expected_tags = { mlflow_tags.MLFLOW_SOURCE_NAME: mock_source_name, mlflow_tags.MLFLOW_SOURCE_TYPE: SourceType.to_string(source_type), mlflow_tags.MLFLOW_GIT_COMMIT: mock_source_version, mlflow_tags.MLFLOW_PROJECT_ENTRY_POINT: mock_entry_point_name, mlflow_tags.MLFLOW_PARENT_RUN_ID: parent_run.info.run_id } create_run_patch = mock.patch.object(MlflowClient, "create_run") with databricks_notebook_patch, active_run_stack_patch, create_run_patch: active_run = start_run(experiment_id=mock_experiment_id, source_name=mock_source_name, source_version=mock_source_version, entry_point_name=mock_entry_point_name, source_type=source_type, run_name=mock_run_name, nested=True) MlflowClient.create_run.assert_called_once_with( experiment_id=mock_experiment_id, run_name=mock_run_name, tags=expected_tags) assert is_from_run(active_run, MlflowClient.create_run.return_value)
def create_run(self, experiment_id, user_id, run_name, source_type, source_name, entry_point_name, start_time, source_version, tags, parent_run_id): experiment = self.get_experiment(experiment_id) if experiment.lifecycle_stage != LifecycleStage.ACTIVE: raise MlflowException( 'Experiment id={} must be active'.format(experiment_id), INVALID_STATE) run_uuid = uuid.uuid4().hex artifact_location = build_path(experiment.artifact_location, run_uuid, SqlAlchemyStore.ARTIFACTS_FOLDER_NAME) run = SqlRun(name=run_name or "", artifact_uri=artifact_location, run_uuid=run_uuid, experiment_id=experiment_id, source_type=SourceType.to_string(source_type), source_name=source_name, entry_point_name=entry_point_name, user_id=user_id, status=RunStatus.to_string(RunStatus.RUNNING), start_time=start_time, end_time=None, source_version=source_version, lifecycle_stage=LifecycleStage.ACTIVE) for tag in tags: run.tags.append(SqlTag(key=tag.key, value=tag.value)) if parent_run_id: run.tags.append( SqlTag(key=MLFLOW_PARENT_RUN_ID, value=parent_run_id)) if run_name: run.tags.append(SqlTag(key=MLFLOW_RUN_NAME, value=run_name)) self._save_to_db([run]) return run.to_mlflow_entity()
def to_mlflow_entity(self): """ Convert DB model to corresponding MLflow entity. :return: :py:class:`mlflow.entities.Run`. """ run_info = RunInfo(run_uuid=self.run_uuid, run_id=self.run_uuid, experiment_id=str(self.experiment_id), name=self.name, source_type=SourceType.from_string( self.source_type), source_name=self.source_name, entry_point_name=self.entry_point_name, user_id=self.user_id, status=RunStatus.from_string(self.status), start_time=self.start_time, end_time=self.end_time, source_version=self.source_version, lifecycle_stage=self.lifecycle_stage, artifact_uri=self.artifact_uri) # only get latest recorded metrics per key all_metrics = [m.to_mlflow_entity() for m in self.metrics] metrics = {} for m in all_metrics: existing_metric = metrics.get(m.key) if (existing_metric is None)\ or ((m.step, m.timestamp, m.value) >= (existing_metric.step, existing_metric.timestamp, existing_metric.value)): metrics[m.key] = m run_data = RunData(metrics=list(metrics.values()), params=[p.to_mlflow_entity() for p in self.params], tags=[t.to_mlflow_entity() for t in self.tags]) return Run(run_info=run_info, run_data=run_data)
def _setup_remote(azureml_run): from azureml.mlflow._internal.constants import MLflowRunEnvVars logger.debug("Setting up a Remote MLflow run") tracking_uri = azureml_run.experiment.workspace.get_mlflow_tracking_uri() + "&is-remote=True" logger.debug("Setting MLflow tracking uri env var") os.environ[MLflowRunEnvVars.TRACKING_URI] = tracking_uri logger.debug("Setting MLflow run id env var with {}".format(azureml_run.id)) os.environ[MLflowRunEnvVars.ID] = azureml_run.id logger.debug("Setting Mlflow experiment with {}".format(azureml_run.experiment.name)) os.environ[MLflowRunEnvVars.EXPERIMENT_NAME] = azureml_run.experiment.name from mlflow.entities import SourceType if not os.environ.get("AZUREML_SECONDARY_INSTANCE"): mlflow_tags = {} mlflow_source_type_key = 'mlflow.source.type' if mlflow_source_type_key not in azureml_run.tags: logger.debug("Setting the mlflow tag {}".format(mlflow_source_type_key)) mlflow_tags[mlflow_source_type_key] = SourceType.to_string(SourceType.JOB) mlflow_source_name_key = 'mlflow.source.name' if mlflow_source_name_key not in azureml_run.tags: logger.debug("Setting the mlflow tag {}".format(mlflow_source_name_key)) mlflow_tags[mlflow_source_name_key] = azureml_run.get_details()['runDefinition']['script'] azureml_run.set_tags(mlflow_tags)
def create_run(self, experiment_id, user_id, start_time, tags): with self.ManagedSessionMaker() as session: experiment = self.get_experiment(experiment_id) self._check_experiment_is_active(experiment) run_id = uuid.uuid4().hex artifact_location = append_to_uri_path(experiment.artifact_location, run_id, SqlAlchemyStore.ARTIFACTS_FOLDER_NAME) run = SqlRun(name="", artifact_uri=artifact_location, run_uuid=run_id, experiment_id=experiment_id, source_type=SourceType.to_string(SourceType.UNKNOWN), source_name="", entry_point_name="", user_id=user_id, status=RunStatus.to_string(RunStatus.RUNNING), start_time=start_time, end_time=None, source_version="", lifecycle_stage=LifecycleStage.ACTIVE) tags_dict = {} for tag in tags: tags_dict[tag.key] = tag.value run.tags = [SqlTag(key=key, value=value) for key, value in tags_dict.items()] self._save_to_db(objs=run, session=session) return run.to_mlflow_entity()
def tags(self): notebook_id = databricks_utils.get_notebook_id() notebook_path = databricks_utils.get_notebook_path() webapp_url = databricks_utils.get_webapp_url() workspace_url = databricks_utils.get_workspace_url() workspace_url_fallback, workspace_id = databricks_utils.get_workspace_info_from_dbutils() tags = { MLFLOW_SOURCE_NAME: notebook_path, MLFLOW_SOURCE_TYPE: SourceType.to_string(SourceType.NOTEBOOK), } if notebook_id is not None: tags[MLFLOW_DATABRICKS_NOTEBOOK_ID] = notebook_id if notebook_path is not None: tags[MLFLOW_DATABRICKS_NOTEBOOK_PATH] = notebook_path if webapp_url is not None: tags[MLFLOW_DATABRICKS_WEBAPP_URL] = webapp_url if workspace_url is not None: tags[MLFLOW_DATABRICKS_WORKSPACE_URL] = workspace_url elif workspace_url_fallback is not None: tags[MLFLOW_DATABRICKS_WORKSPACE_URL] = workspace_url_fallback if workspace_id is not None: tags[MLFLOW_DATABRICKS_WORKSPACE_ID] = workspace_id return tags
RunTag, Metric, Param, RunData, RunInfo, SourceType, RunStatus, Run, ViewType, ExperimentTag, ) from mlflow.entities.lifecycle_stage import LifecycleStage from mlflow.store.db.base_sql_model import Base SourceTypes = [ SourceType.to_string(SourceType.NOTEBOOK), SourceType.to_string(SourceType.JOB), SourceType.to_string(SourceType.LOCAL), SourceType.to_string(SourceType.UNKNOWN), SourceType.to_string(SourceType.PROJECT), ] RunStatusTypes = [ RunStatus.to_string(RunStatus.SCHEDULED), RunStatus.to_string(RunStatus.FAILED), RunStatus.to_string(RunStatus.FINISHED), RunStatus.to_string(RunStatus.RUNNING), RunStatus.to_string(RunStatus.KILLED), ]
class SqlRun(Base): """ DB model for :py:class:`mlflow.entities.Run`. These are recorded in ``runs`` table. """ __tablename__ = "runs" run_uuid = Column(String(32), nullable=False) """ Run UUID: `String` (limit 32 characters). *Primary Key* for ``runs`` table. """ name = Column(String(250)) """ Run name: `String` (limit 250 characters). """ source_type = Column(String(20), default=SourceType.to_string(SourceType.LOCAL)) """ Source Type: `String` (limit 20 characters). Can be one of ``NOTEBOOK``, ``JOB``, ``PROJECT``, ``LOCAL`` (default), or ``UNKNOWN``. """ source_name = Column(String(500)) """ Name of source recording the run: `String` (limit 500 characters). """ entry_point_name = Column(String(50)) """ Entry-point name that launched the run run: `String` (limit 50 characters). """ user_id = Column(String(256), nullable=True, default=None) """ User ID: `String` (limit 256 characters). Defaults to ``null``. """ status = Column(String(20), default=RunStatus.to_string(RunStatus.SCHEDULED)) """ Run Status: `String` (limit 20 characters). Can be one of ``RUNNING``, ``SCHEDULED`` (default), ``FINISHED``, ``FAILED``. """ start_time = Column(BigInteger, default=int(time.time())) """ Run start time: `BigInteger`. Defaults to current system time. """ end_time = Column(BigInteger, nullable=True, default=None) """ Run end time: `BigInteger`. """ source_version = Column(String(50)) """ Source version: `String` (limit 50 characters). """ lifecycle_stage = Column(String(20), default=LifecycleStage.ACTIVE) """ Lifecycle Stage of run: `String` (limit 32 characters). Can be either ``active`` (default) or ``deleted``. """ artifact_uri = Column(String(200), default=None) """ Default artifact location for this run: `String` (limit 200 characters). """ experiment_id = Column(Integer, ForeignKey("experiments.experiment_id")) """ Experiment ID to which this run belongs to: *Foreign Key* into ``experiment`` table. """ experiment = relationship("SqlExperiment", backref=backref("runs", cascade="all")) """ SQLAlchemy relationship (many:one) with :py:class:`mlflow.store.dbmodels.models.SqlExperiment`. """ __table_args__ = ( CheckConstraint(source_type.in_(SourceTypes), name="source_type"), CheckConstraint(status.in_(RunStatusTypes), name="status"), CheckConstraint( lifecycle_stage.in_( LifecycleStage.view_type_to_stages(ViewType.ALL)), name="runs_lifecycle_stage", ), PrimaryKeyConstraint("run_uuid", name="run_pk"), ) @staticmethod def get_attribute_name(mlflow_attribute_name): """ Resolves an MLflow attribute name to a `SqlRun` attribute name. """ # Currently, MLflow Search attributes defined in `SearchUtils.VALID_SEARCH_ATTRIBUTE_KEYS` # share the same names as their corresponding `SqlRun` attributes. Therefore, this function # returns the same attribute name return mlflow_attribute_name def to_mlflow_entity(self): """ Convert DB model to corresponding MLflow entity. :return: :py:class:`mlflow.entities.Run`. """ run_info = RunInfo( run_uuid=self.run_uuid, run_id=self.run_uuid, experiment_id=str(self.experiment_id), user_id=self.user_id, status=self.status, start_time=self.start_time, end_time=self.end_time, lifecycle_stage=self.lifecycle_stage, artifact_uri=self.artifact_uri, ) run_data = RunData( metrics=[m.to_mlflow_entity() for m in self.latest_metrics], params=[p.to_mlflow_entity() for p in self.params], tags=[t.to_mlflow_entity() for t in self.tags], ) return Run(run_info=run_info, run_data=run_data)
class SqlRun(Base): """ DB model for :py:class:`mlflow.entities.Run`. These are recorded in ``runs`` table. """ __tablename__ = 'runs' run_uuid = Column(String(32), nullable=False) """ Run UUID: `String` (limit 32 characters). *Primary Key* for ``runs`` table. """ name = Column(String(250)) """ Run name: `String` (limit 250 characters). """ source_type = Column(String(20), default=SourceType.to_string(SourceType.LOCAL)) """ Source Type: `String` (limit 20 characters). Can be one of ``NOTEBOOK``, ``JOB``, ``PROJECT``, ``LOCAL`` (default), or ``UNKNOWN``. """ source_name = Column(String(500)) """ Name of source recording the run: `String` (limit 500 characters). """ entry_point_name = Column(String(50)) """ Entry-point name that launched the run run: `String` (limit 50 characters). """ user_id = Column(String(256), nullable=True, default=None) """ User ID: `String` (limit 256 characters). Defaults to ``null``. """ status = Column(String(20), default=RunStatus.to_string(RunStatus.SCHEDULED)) """ Run Status: `String` (limit 20 characters). Can be one of ``RUNNING``, ``SCHEDULED`` (default), ``FINISHED``, ``FAILED``. """ start_time = Column(BigInteger, default=int(time.time())) """ Run start time: `BigInteger`. Defaults to current system time. """ end_time = Column(BigInteger, nullable=True, default=None) """ Run end time: `BigInteger`. """ source_version = Column(String(50)) """ Source version: `String` (limit 50 characters). """ lifecycle_stage = Column(String(20), default=LifecycleStage.ACTIVE) """ Lifecycle Stage of run: `String` (limit 32 characters). Can be either ``active`` (default) or ``deleted``. """ artifact_uri = Column(String(200), default=None) """ Default artifact location for this run: `String` (limit 200 characters). """ experiment_id = Column(Integer, ForeignKey('experiments.experiment_id')) """ Experiment ID to which this run belongs to: *Foreign Key* into ``experiment`` table. """ experiment = relationship('SqlExperiment', backref=backref('runs', cascade='all')) """ SQLAlchemy relationship (many:one) with :py:class:`mlflow.store.dbmodels.models.SqlExperiment`. """ __table_args__ = (CheckConstraint(source_type.in_(SourceTypes), name='source_type'), CheckConstraint(status.in_(RunStatusTypes), name='status'), CheckConstraint(lifecycle_stage.in_( LifecycleStage.view_type_to_stages(ViewType.ALL)), name='runs_lifecycle_stage'), PrimaryKeyConstraint('run_uuid', name='run_pk')) def to_mlflow_entity(self): """ Convert DB model to corresponding MLflow entity. :return: :py:class:`mlflow.entities.Run`. """ run_info = RunInfo(run_uuid=self.run_uuid, run_id=self.run_uuid, experiment_id=str(self.experiment_id), user_id=self.user_id, status=self.status, start_time=self.start_time, end_time=self.end_time, lifecycle_stage=self.lifecycle_stage, artifact_uri=self.artifact_uri) # only get latest recorded metrics per key all_metrics = [m.to_mlflow_entity() for m in self.metrics] metrics = {} for m in all_metrics: existing_metric = metrics.get(m.key) if (existing_metric is None)\ or ((m.step, m.timestamp, m.value) >= (existing_metric.step, existing_metric.timestamp, existing_metric.value)): metrics[m.key] = m run_data = RunData(metrics=list(metrics.values()), params=[p.to_mlflow_entity() for p in self.params], tags=[t.to_mlflow_entity() for t in self.tags]) return Run(run_info=run_info, run_data=run_data)