def test_start_run_existing_run_deleted(empty_active_run_stack): mock_run = mock.Mock() mock_run.info.lifecycle_stage = LifecycleStage.DELETED run_id = uuid.uuid4().hex with mock.patch.object(MlflowClient, "get_run", return_value=mock_run): with pytest.raises(MlflowException): start_run(run_id)
def test_start_run_existing_run_from_environment_with_set_environment( empty_active_run_stack): mock_run = mock.Mock() mock_run.info.lifecycle_stage = LifecycleStage.ACTIVE run_id = uuid.uuid4().hex env_patch = mock.patch.dict("os.environ", {_RUN_ID_ENV_VAR: run_id}) with env_patch, mock.patch.object(MlflowClient, "get_run", return_value=mock_run): with pytest.raises(MlflowException): set_experiment("test-run") start_run()
def test_log_metric(): with start_run() as active_run, mock.patch("time.time") as time_mock: time_mock.side_effect = [123 for _ in range(100)] run_id = active_run.info.run_id kiwi.log_metric("name_1", 25) kiwi.log_metric("name_2", -3) kiwi.log_metric("name_1", 30, 5) kiwi.log_metric("name_1", 40, -2) kiwi.log_metric("nested/nested/name", 40) finished_run = tracking.MlflowClient().get_run(run_id) # Validate metrics assert len(finished_run.data.metrics) == 3 expected_pairs = {"name_1": 30, "name_2": -3, "nested/nested/name": 40} for key, value in finished_run.data.metrics.items(): assert expected_pairs[key] == value client = tracking.MlflowClient() metric_history_name1 = client.get_metric_history(run_id, "name_1") assert set([(m.value, m.timestamp, m.step) for m in metric_history_name1]) == set([ (25, 123 * 1000, 0), (30, 123 * 1000, 5), (40, 123 * 1000, -2), ]) metric_history_name2 = client.get_metric_history(run_id, "name_2") assert set([(m.value, m.timestamp, m.step) for m in metric_history_name2]) == set([ (-3, 123 * 1000, 0), ])
def test_log_batch_validates_entity_names_and_values(): bad_kwargs = { "metrics": [ [Metric(key="../bad/metric/name", value=0.3, timestamp=3, step=0)], [ Metric(key="ok-name", value="non-numerical-value", timestamp=3, step=0) ], [ Metric(key="ok-name", value=0.3, timestamp="non-numerical-timestamp", step=0) ], ], "params": [[Param(key="../bad/param/name", value="my-val")]], "tags": [[Param(key="../bad/tag/name", value="my-val")]], } with start_run() as active_run: for kwarg, bad_values in bad_kwargs.items(): for bad_kwarg_value in bad_values: final_kwargs = { "run_id": active_run.info.run_id, "metrics": [], "params": [], "tags": [], } final_kwargs[kwarg] = bad_kwarg_value with pytest.raises(MlflowException) as e: tracking.MlflowClient().log_batch(**final_kwargs) assert e.value.error_code == ErrorCode.Name( INVALID_PARAMETER_VALUE)
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 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_log_metrics_uses_millisecond_timestamp_resolution_fluent(): with start_run() as active_run, mock.patch("time.time") as time_mock: time_mock.side_effect = lambda: 123 kiwi.log_metrics({ "name_1": 25, "name_2": -3, }) kiwi.log_metrics({ "name_1": 30, }) kiwi.log_metrics({ "name_1": 40, }) run_id = active_run.info.run_id client = tracking.MlflowClient() metric_history_name1 = client.get_metric_history(run_id, "name_1") assert set([(m.value, m.timestamp) for m in metric_history_name1]) == set([ (25, 123 * 1000), (30, 123 * 1000), (40, 123 * 1000), ]) metric_history_name2 = client.get_metric_history(run_id, "name_2") assert set([(m.value, m.timestamp) for m in metric_history_name2]) == set([ (-3, 123 * 1000), ])
def test_log_metric_validation(): with start_run() as active_run: run_id = active_run.info.run_id with pytest.raises(MlflowException) as e: kiwi.log_metric("name_1", "apple") assert e.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE) finished_run = tracking.MlflowClient().get_run(run_id) assert len(finished_run.data.metrics) == 0
def test_start_run_context_manager(): with start_run() as first_run: first_uuid = first_run.info.run_id # Check that start_run() causes the run information to be persisted in the store persisted_run = tracking.MlflowClient().get_run(first_uuid) assert persisted_run is not None assert persisted_run.info == first_run.info finished_run = tracking.MlflowClient().get_run(first_uuid) assert finished_run.info.status == RunStatus.to_string(RunStatus.FINISHED) # Launch a separate run that fails, verify the run status is FAILED and the run UUID is # different with pytest.raises(Exception): with start_run() as second_run: second_run_id = second_run.info.run_id raise Exception("Failing run!") assert second_run_id != first_uuid finished_run2 = tracking.MlflowClient().get_run(second_run_id) assert finished_run2.info.status == RunStatus.to_string(RunStatus.FAILED)
def test_start_and_end_run(): # Use the start_run() and end_run() APIs without a `with` block, verify they work. with start_run() as active_run: kiwi.log_metric("name_1", 25) finished_run = tracking.MlflowClient().get_run(active_run.info.run_id) # Validate metrics assert len(finished_run.data.metrics) == 1 assert finished_run.data.metrics["name_1"] == 25
def test_set_experiment_with_deleted_experiment_name(): name = "dead_exp" kiwi.set_experiment(name) with start_run() as run: exp_id = run.info.experiment_id tracking.MlflowClient().delete_experiment(exp_id) with pytest.raises(MlflowException): kiwi.set_experiment(name)
def test_start_run_existing_run(empty_active_run_stack): mock_run = mock.Mock() mock_run.info.lifecycle_stage = LifecycleStage.ACTIVE run_id = uuid.uuid4().hex with mock.patch.object(MlflowClient, "get_run", return_value=mock_run): active_run = start_run(run_id) assert is_from_run(active_run, mock_run) MlflowClient.get_run.assert_called_once_with(run_id)
def test_set_experiment(): with pytest.raises(TypeError): kiwi.set_experiment() # pylint: disable=no-value-for-parameter with pytest.raises(Exception): kiwi.set_experiment(None) with pytest.raises(Exception): kiwi.set_experiment("") name = "random_exp" exp_id = kiwi.create_experiment(name) kiwi.set_experiment(name) with start_run() as run: assert run.info.experiment_id == exp_id another_name = "another_experiment" kiwi.set_experiment(another_name) exp_id2 = kiwi.tracking.MlflowClient().get_experiment_by_name(another_name) with start_run() as another_run: assert another_run.info.experiment_id == exp_id2.experiment_id
def test_log_params(): expected_params = {"name_1": "c", "name_2": "b", "nested/nested/name": 5} with start_run() as active_run: run_id = active_run.info.run_id kiwi.log_params(expected_params) finished_run = tracking.MlflowClient().get_run(run_id) # Validate params assert finished_run.data.params == { "name_1": "c", "name_2": "b", "nested/nested/name": "5" }
def test_log_artifact_with_dirs(tmpdir): # Test log artifact with a directory art_dir = tmpdir.mkdir("parent") file0 = art_dir.join("file0") file0.write("something") file1 = art_dir.join("file1") file1.write("something") sub_dir = art_dir.mkdir("child") with start_run(): artifact_uri = kiwi.get_artifact_uri() run_artifact_dir = local_file_uri_to_path(artifact_uri) kiwi.log_artifact(str(art_dir)) base = os.path.basename(str(art_dir)) assert os.listdir(run_artifact_dir) == [base] assert set(os.listdir(os.path.join(run_artifact_dir, base))) == \ {'child', 'file0', 'file1'} with open(os.path.join(run_artifact_dir, base, "file0")) as f: assert f.read() == "something" # Test log artifact with directory and specified parent folder art_dir = tmpdir.mkdir("dir") with start_run(): artifact_uri = kiwi.get_artifact_uri() run_artifact_dir = local_file_uri_to_path(artifact_uri) kiwi.log_artifact(str(art_dir), "some_parent") assert os.listdir(run_artifact_dir) == [ os.path.basename("some_parent") ] assert os.listdir(os.path.join(run_artifact_dir, "some_parent")) == \ [os.path.basename(str(art_dir))] sub_dir = art_dir.mkdir("another_dir") with start_run(): artifact_uri = kiwi.get_artifact_uri() run_artifact_dir = local_file_uri_to_path(artifact_uri) kiwi.log_artifact(str(art_dir), "parent/and_child") assert os.listdir(os.path.join(run_artifact_dir, "parent", "and_child")) == \ [os.path.basename(str(art_dir))] assert os.listdir(os.path.join(run_artifact_dir, "parent", "and_child", os.path.basename(str(art_dir)))) == \ [os.path.basename(str(sub_dir))]
def test_log_param(): with start_run() as active_run: run_id = active_run.info.run_id kiwi.log_param("name_1", "a") kiwi.log_param("name_2", "b") kiwi.log_param("nested/nested/name", 5) finished_run = tracking.MlflowClient().get_run(run_id) # Validate params assert finished_run.data.params == { "name_1": "a", "name_2": "b", "nested/nested/name": "5" }
def test_start_run_existing_run_from_environment(empty_active_run_stack): mock_run = mock.Mock() mock_run.info.lifecycle_stage = LifecycleStage.ACTIVE run_id = uuid.uuid4().hex env_patch = mock.patch.dict("os.environ", {_RUN_ID_ENV_VAR: run_id}) with env_patch, mock.patch.object(MlflowClient, "get_run", return_value=mock_run): active_run = start_run() assert is_from_run(active_run, mock_run) MlflowClient.get_run.assert_called_once_with(run_id)
def test_log_artifact(): artifact_src_dir = tempfile.mkdtemp() # Create artifacts _, path0 = tempfile.mkstemp(dir=artifact_src_dir) _, path1 = tempfile.mkstemp(dir=artifact_src_dir) for i, path in enumerate([path0, path1]): with open(path, "w") as handle: handle.write("%s" % str(i)) # Log an artifact, verify it exists in the directory returned by get_artifact_uri # after the run finishes artifact_parent_dirs = ["some_parent_dir", None] for parent_dir in artifact_parent_dirs: with start_run(): artifact_uri = kiwi.get_artifact_uri() run_artifact_dir = local_file_uri_to_path(artifact_uri) kiwi.log_artifact(path0, parent_dir) expected_dir = os.path.join(run_artifact_dir, parent_dir) \ if parent_dir is not None else run_artifact_dir assert os.listdir(expected_dir) == [os.path.basename(path0)] logged_artifact_path = os.path.join(expected_dir, path0) assert filecmp.cmp(logged_artifact_path, path0, shallow=False) # Log multiple artifacts, verify they exist in the directory returned by get_artifact_uri for parent_dir in artifact_parent_dirs: with start_run(): artifact_uri = kiwi.get_artifact_uri() run_artifact_dir = local_file_uri_to_path(artifact_uri) kiwi.log_artifacts(artifact_src_dir, parent_dir) # Check that the logged artifacts match expected_artifact_output_dir = os.path.join(run_artifact_dir, parent_dir) \ if parent_dir is not None else run_artifact_dir dir_comparison = filecmp.dircmp(artifact_src_dir, expected_artifact_output_dir) assert len(dir_comparison.left_only) == 0 assert len(dir_comparison.right_only) == 0 assert len(dir_comparison.diff_files) == 0 assert len(dir_comparison.funny_files) == 0
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_log_metrics_uses_common_timestamp_and_step_per_invocation(step_kwarg): expected_metrics = {"name_1": 30, "name_2": -3, "nested/nested/name": 40} with start_run() as active_run: run_id = active_run.info.run_id kiwi.log_metrics(expected_metrics, step=step_kwarg) finished_run = tracking.MlflowClient().get_run(run_id) # Validate metric key/values match what we expect, and that all metrics have the same timestamp assert len(finished_run.data.metrics) == len(expected_metrics) for key, value in finished_run.data.metrics.items(): assert expected_metrics[key] == value common_timestamp = finished_run.data._metric_objs[0].timestamp expected_step = step_kwarg if step_kwarg is not None else 0 for metric_obj in finished_run.data._metric_objs: assert metric_obj.timestamp == common_timestamp assert metric_obj.step == expected_step
def test_set_tags(): exact_expected_tags = { "name_1": "c", "name_2": "b", "nested/nested/name": 5 } approx_expected_tags = set( [MLFLOW_USER, MLFLOW_SOURCE_NAME, MLFLOW_SOURCE_TYPE]) with start_run() as active_run: run_id = active_run.info.run_id kiwi.set_tags(exact_expected_tags) finished_run = tracking.MlflowClient().get_run(run_id) # Validate tags assert len(finished_run.data.tags ) == len(exact_expected_tags) + len(approx_expected_tags) for tag_key, tag_val in finished_run.data.tags.items(): if tag_key in approx_expected_tags: pass else: assert str(exact_expected_tags[tag_key]) == tag_val
def test_log_metrics_uses_millisecond_timestamp_resolution_client(): with start_run() as active_run, mock.patch("time.time") as time_mock: time_mock.side_effect = lambda: 123 mlflow_client = tracking.MlflowClient() run_id = active_run.info.run_id mlflow_client.log_metric(run_id=run_id, key="name_1", value=25) mlflow_client.log_metric(run_id=run_id, key="name_2", value=-3) mlflow_client.log_metric(run_id=run_id, key="name_1", value=30) mlflow_client.log_metric(run_id=run_id, key="name_1", value=40) metric_history_name1 = mlflow_client.get_metric_history(run_id, "name_1") assert set([(m.value, m.timestamp) for m in metric_history_name1]) == set([ (25, 123 * 1000), (30, 123 * 1000), (40, 123 * 1000), ]) metric_history_name2 = mlflow_client.get_metric_history(run_id, "name_2") assert set([(m.value, m.timestamp) for m in metric_history_name2]) == set([ (-3, 123 * 1000), ])
def test_start_run_defaults(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.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 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 test_start_run_with_parent_non_nested(): with mock.patch("mlflow.tracking.fluent._active_run_stack", [mock.Mock()]): with pytest.raises(Exception): start_run()
def test_log_batch(): expected_metrics = {"metric-key0": 1.0, "metric-key1": 4.0} expected_params = {"param-key0": "param-val0", "param-key1": "param-val1"} exact_expected_tags = {"tag-key0": "tag-val0", "tag-key1": "tag-val1"} approx_expected_tags = set( [MLFLOW_USER, MLFLOW_SOURCE_NAME, MLFLOW_SOURCE_TYPE]) t = int(time.time()) sorted_expected_metrics = sorted(expected_metrics.items(), key=lambda kv: kv[0]) metrics = [ Metric(key=key, value=value, timestamp=t, step=i) for i, (key, value) in enumerate(sorted_expected_metrics) ] params = [ Param(key=key, value=value) for key, value in expected_params.items() ] tags = [ RunTag(key=key, value=value) for key, value in exact_expected_tags.items() ] with start_run() as active_run: run_id = active_run.info.run_id kiwi.tracking.MlflowClient().log_batch(run_id=run_id, metrics=metrics, params=params, tags=tags) client = tracking.MlflowClient() finished_run = client.get_run(run_id) # Validate metrics assert len(finished_run.data.metrics) == 2 for key, value in finished_run.data.metrics.items(): assert expected_metrics[key] == value metric_history0 = client.get_metric_history(run_id, "metric-key0") assert set([(m.value, m.timestamp, m.step) for m in metric_history0]) == set([ (1.0, t, 0), ]) metric_history1 = client.get_metric_history(run_id, "metric-key1") assert set([(m.value, m.timestamp, m.step) for m in metric_history1]) == set([ (4.0, t, 1), ]) # Validate tags (for automatically-set tags) assert len(finished_run.data.tags ) == len(exact_expected_tags) + len(approx_expected_tags) for tag_key, tag_value in finished_run.data.tags.items(): if tag_key in approx_expected_tags: pass else: assert exact_expected_tags[tag_key] == tag_value # Validate params assert finished_run.data.params == expected_params # test that log_batch works with fewer params new_tags = {"1": "2", "3": "4", "5": "6"} tags = [RunTag(key=key, value=value) for key, value in new_tags.items()] client.log_batch(run_id=run_id, tags=tags) finished_run_2 = client.get_run(run_id) # Validate tags (for automatically-set tags) assert len(finished_run_2.data.tags) == len(finished_run.data.tags) + 3 for tag_key, tag_value in finished_run_2.data.tags.items(): if tag_key in new_tags: assert new_tags[tag_key] == tag_value