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