wehub-resource-sync
597 lines
19 KiB
Python
597 lines
19 KiB
Python
import json
|
|
import logging
|
|
import os
|
|
import textwrap
|
|
from unittest import mock
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
from click.testing import CliRunner
|
|
|
|
import mlflow
|
|
from mlflow import experiments
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.runs import create_run, link_traces, list_run
|
|
from mlflow.tracing.constant import TraceExperimentTagKey
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def suppress_logging():
|
|
"""Suppress logging for all tests to ensure clean CLI output."""
|
|
# Suppress all logging
|
|
logging.disable(logging.CRITICAL)
|
|
|
|
yield
|
|
|
|
# Re-enable logging
|
|
logging.disable(logging.NOTSET)
|
|
|
|
|
|
def test_list_run():
|
|
with mlflow.start_run(run_name="apple"):
|
|
pass
|
|
result = CliRunner().invoke(list_run, ["--experiment-id", "0"])
|
|
assert "apple" in result.output
|
|
|
|
|
|
def test_list_run_experiment_id_required():
|
|
result = CliRunner().invoke(list_run, [])
|
|
assert "Missing option '--experiment-id'" in result.output
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
"MLFLOW_SKINNY" in os.environ,
|
|
reason="Skinny Client does not support predict due to the pandas dependency",
|
|
)
|
|
def test_csv_generation(tmp_path):
|
|
import numpy as np
|
|
import pandas as pd
|
|
|
|
with mock.patch(
|
|
"mlflow.experiments.fluent.search_runs",
|
|
return_value=pd.DataFrame(
|
|
{
|
|
"run_id": np.array(["all_set", "with_none", "with_nan"]),
|
|
"experiment_id": np.array([1, 1, 1]),
|
|
"param_optimizer": np.array(["Adam", None, "Adam"]),
|
|
"avg_loss": np.array([42.0, None, np.nan], dtype=np.float32),
|
|
},
|
|
columns=["run_id", "experiment_id", "param_optimizer", "avg_loss"],
|
|
),
|
|
):
|
|
expected_csv = textwrap.dedent(
|
|
"""\
|
|
run_id,experiment_id,param_optimizer,avg_loss
|
|
all_set,1,Adam,42.0
|
|
with_none,1,,
|
|
with_nan,1,Adam,
|
|
"""
|
|
)
|
|
result_filename = os.path.join(tmp_path, "result.csv")
|
|
CliRunner().invoke(
|
|
experiments.generate_csv_with_runs,
|
|
["--experiment-id", "1", "--filename", result_filename],
|
|
)
|
|
with open(result_filename) as fd:
|
|
assert expected_csv == fd.read()
|
|
|
|
|
|
def test_create_run_with_experiment_id():
|
|
mlflow.create_experiment("test_create_run_exp")
|
|
exp = mlflow.get_experiment_by_name("test_create_run_exp")
|
|
|
|
result = CliRunner().invoke(create_run, ["--experiment-id", exp.experiment_id])
|
|
assert result.exit_code == 0
|
|
|
|
output = json.loads(result.output)
|
|
assert "run_id" in output
|
|
assert output["experiment_id"] == exp.experiment_id
|
|
assert output["status"] == "FINISHED"
|
|
|
|
# Verify the run was created
|
|
run = mlflow.get_run(output["run_id"])
|
|
assert run.info.experiment_id == exp.experiment_id
|
|
assert run.info.status == "FINISHED"
|
|
|
|
|
|
def test_create_run_with_experiment_name():
|
|
exp_name = "test_create_run_by_name"
|
|
|
|
result = CliRunner().invoke(create_run, ["--experiment-name", exp_name])
|
|
assert result.exit_code == 0
|
|
|
|
output = json.loads(result.output)
|
|
assert "run_id" in output
|
|
assert output["status"] == "FINISHED"
|
|
|
|
# Verify experiment was created
|
|
exp = mlflow.get_experiment_by_name(exp_name)
|
|
assert exp is not None
|
|
assert output["experiment_id"] == exp.experiment_id
|
|
|
|
|
|
def test_create_run_with_custom_name_and_description():
|
|
mlflow.create_experiment("test_run_with_details")
|
|
exp = mlflow.get_experiment_by_name("test_run_with_details")
|
|
|
|
run_name = "my-custom-run"
|
|
description = "This is a test run"
|
|
|
|
result = CliRunner().invoke(
|
|
create_run,
|
|
[
|
|
"--experiment-id",
|
|
exp.experiment_id,
|
|
"--run-name",
|
|
run_name,
|
|
"--description",
|
|
description,
|
|
],
|
|
)
|
|
assert result.exit_code == 0
|
|
|
|
output = json.loads(result.output)
|
|
assert output["run_name"] == run_name
|
|
|
|
# Verify run details
|
|
run = mlflow.get_run(output["run_id"])
|
|
assert run.data.tags.get("mlflow.note.content") == description
|
|
assert run.info.run_name == run_name
|
|
|
|
|
|
def test_create_run_with_tags():
|
|
mlflow.create_experiment("test_run_with_tags")
|
|
exp = mlflow.get_experiment_by_name("test_run_with_tags")
|
|
|
|
result = CliRunner().invoke(
|
|
create_run,
|
|
[
|
|
"--experiment-id",
|
|
exp.experiment_id,
|
|
"--tags",
|
|
"env=test",
|
|
"--tags",
|
|
"model=linear",
|
|
"--tags",
|
|
"version=1.0",
|
|
],
|
|
)
|
|
assert result.exit_code == 0
|
|
|
|
output = json.loads(result.output)
|
|
run = mlflow.get_run(output["run_id"])
|
|
|
|
assert run.data.tags["env"] == "test"
|
|
assert run.data.tags["model"] == "linear"
|
|
assert run.data.tags["version"] == "1.0"
|
|
|
|
|
|
@pytest.mark.parametrize("status", ["FAILED", "KILLED"])
|
|
def test_create_run_with_different_status(status):
|
|
mlflow.create_experiment("test_run_statuses")
|
|
exp = mlflow.get_experiment_by_name("test_run_statuses")
|
|
|
|
result = CliRunner().invoke(
|
|
create_run, ["--experiment-id", exp.experiment_id, "--status", status]
|
|
)
|
|
assert result.exit_code == 0
|
|
output = json.loads(result.output)
|
|
assert output["status"] == status
|
|
|
|
run = mlflow.get_run(output["run_id"])
|
|
assert run.info.status == status
|
|
|
|
|
|
def test_create_run_missing_experiment():
|
|
result = CliRunner().invoke(create_run, [])
|
|
assert result.exit_code != 0
|
|
assert "Must specify exactly one of --experiment-id or --experiment-name" in result.output
|
|
|
|
|
|
def test_create_run_both_experiment_params():
|
|
result = CliRunner().invoke(create_run, ["--experiment-id", "0", "--experiment-name", "test"])
|
|
assert result.exit_code != 0
|
|
assert "Must specify exactly one of --experiment-id or --experiment-name" in result.output
|
|
|
|
|
|
def test_create_run_invalid_tag_format():
|
|
mlflow.create_experiment("test_invalid_tag")
|
|
exp = mlflow.get_experiment_by_name("test_invalid_tag")
|
|
|
|
result = CliRunner().invoke(
|
|
create_run, ["--experiment-id", exp.experiment_id, "--tags", "invalid-tag"]
|
|
)
|
|
assert result.exit_code != 0
|
|
assert "Invalid tag format" in result.output
|
|
|
|
|
|
def test_create_run_duplicate_tag_key():
|
|
mlflow.create_experiment("test_duplicate_tag")
|
|
exp = mlflow.get_experiment_by_name("test_duplicate_tag")
|
|
|
|
result = CliRunner().invoke(
|
|
create_run,
|
|
["--experiment-id", exp.experiment_id, "--tags", "env=test", "--tags", "env=prod"],
|
|
)
|
|
assert result.exit_code != 0
|
|
assert "Duplicate tag key" in result.output
|
|
|
|
|
|
def test_link_traces_single_trace():
|
|
with patch("mlflow.runs.MlflowClient.link_traces_to_run") as mock_link_traces:
|
|
result = CliRunner().invoke(
|
|
link_traces,
|
|
["--run-id", "test_run_123", "--trace-id", "trace_abc"],
|
|
)
|
|
|
|
assert result.exit_code == 0
|
|
assert "Successfully linked 1 trace(s) to run 'test_run_123'" in result.output
|
|
mock_link_traces.assert_called_once_with(["trace_abc"], "test_run_123")
|
|
|
|
|
|
def test_link_traces_multiple_traces():
|
|
with patch("mlflow.runs.MlflowClient.link_traces_to_run") as mock_link_traces:
|
|
result = CliRunner().invoke(
|
|
link_traces,
|
|
[
|
|
"--run-id",
|
|
"test_run_456",
|
|
"--trace-id",
|
|
"trace_1",
|
|
"--trace-id",
|
|
"trace_2",
|
|
"--trace-id",
|
|
"trace_3",
|
|
],
|
|
)
|
|
|
|
assert result.exit_code == 0
|
|
assert "Successfully linked 3 trace(s) to run 'test_run_456'" in result.output
|
|
mock_link_traces.assert_called_once_with(["trace_1", "trace_2", "trace_3"], "test_run_456")
|
|
|
|
|
|
def test_link_traces_with_short_option():
|
|
with patch("mlflow.runs.MlflowClient.link_traces_to_run") as mock_link_traces:
|
|
result = CliRunner().invoke(
|
|
link_traces,
|
|
["--run-id", "run_789", "-t", "trace_x", "-t", "trace_y"],
|
|
)
|
|
|
|
assert result.exit_code == 0
|
|
assert "Successfully linked 2 trace(s) to run 'run_789'" in result.output
|
|
mock_link_traces.assert_called_once_with(["trace_x", "trace_y"], "run_789")
|
|
|
|
|
|
def test_link_traces_file_store_error():
|
|
with patch(
|
|
"mlflow.runs.MlflowClient.link_traces_to_run",
|
|
side_effect=MlflowException(
|
|
"Linking traces to runs is not supported in FileStore. "
|
|
"Please use a database-backed store (e.g., SQLAlchemy store) for this feature."
|
|
),
|
|
):
|
|
result = CliRunner().invoke(
|
|
link_traces,
|
|
["--run-id", "test_run", "--trace-id", "trace_1"],
|
|
)
|
|
|
|
assert result.exit_code != 0
|
|
assert "Failed to link traces" in result.output
|
|
assert "not supported in FileStore" in result.output
|
|
|
|
|
|
def test_link_traces_too_many_traces_error():
|
|
with patch(
|
|
"mlflow.runs.MlflowClient.link_traces_to_run",
|
|
side_effect=MlflowException(
|
|
"Cannot link more than 100 traces to a run in a single request. Provided 101 traces."
|
|
),
|
|
):
|
|
result = CliRunner().invoke(
|
|
link_traces,
|
|
["--run-id", "test_run", "--trace-id", "trace_1"],
|
|
)
|
|
|
|
assert result.exit_code != 0
|
|
assert "Failed to link traces" in result.output
|
|
assert "100" in result.output
|
|
|
|
|
|
def test_link_traces_missing_run_id():
|
|
result = CliRunner().invoke(link_traces, ["--trace-id", "trace_1"])
|
|
|
|
assert result.exit_code != 0
|
|
assert "Missing option '--run-id'" in result.output
|
|
|
|
|
|
def test_link_traces_missing_trace_id():
|
|
result = CliRunner().invoke(link_traces, ["--run-id", "test_run"])
|
|
|
|
assert result.exit_code != 0
|
|
assert "Missing option '--trace-id'" in result.output
|
|
|
|
|
|
def test_link_traces_generic_error():
|
|
with patch(
|
|
"mlflow.runs.MlflowClient.link_traces_to_run",
|
|
side_effect=MlflowException("Some other error"),
|
|
):
|
|
result = CliRunner().invoke(
|
|
link_traces,
|
|
["--run-id", "test_run", "--trace-id", "trace_1"],
|
|
)
|
|
|
|
assert result.exit_code != 0
|
|
assert "Failed to link traces: Some other error" in result.output
|
|
|
|
|
|
def test_get_experiment_default():
|
|
result = CliRunner().invoke(experiments.get_experiment, ["--experiment-id", "0"])
|
|
assert result.exit_code == 0
|
|
|
|
# Default output is table format
|
|
assert "Experiment ID" in result.output
|
|
assert "Name" in result.output
|
|
assert "Artifact Location" in result.output
|
|
assert "Lifecycle Stage" in result.output
|
|
assert ":" in result.output
|
|
|
|
|
|
def test_get_experiment_json():
|
|
exp_id = mlflow.create_experiment("test_get_exp_json", tags={"env": "test"})
|
|
exp = mlflow.get_experiment(exp_id)
|
|
|
|
result = CliRunner().invoke(
|
|
experiments.get_experiment, ["--experiment-id", exp_id, "--output", "json"]
|
|
)
|
|
assert result.exit_code == 0
|
|
|
|
output = json.loads(result.output)
|
|
expected = {
|
|
"experiment_id": exp_id,
|
|
"name": "test_get_exp_json",
|
|
"artifact_location": exp.artifact_location,
|
|
"lifecycle_stage": "active",
|
|
"tags": {"env": "test"},
|
|
"creation_time": exp.creation_time,
|
|
"last_update_time": exp.last_update_time,
|
|
"effective_trace_archival_retention": exp.effective_trace_archival_retention,
|
|
"trace_location": exp.trace_location,
|
|
"workspace": exp.workspace,
|
|
}
|
|
assert output == expected
|
|
|
|
|
|
def test_get_experiment_table():
|
|
exp_id = mlflow.create_experiment("test_get_exp_table", tags={"env": "test", "team": "ml"})
|
|
|
|
result = CliRunner().invoke(
|
|
experiments.get_experiment, ["--experiment-id", exp_id, "--output", "table"]
|
|
)
|
|
assert result.exit_code == 0
|
|
|
|
# Verify table format
|
|
assert "Experiment ID" in result.output
|
|
assert exp_id in result.output
|
|
assert "Name" in result.output
|
|
assert "test_get_exp_table" in result.output
|
|
assert "Lifecycle Stage" in result.output
|
|
assert "active" in result.output
|
|
assert "Tags" in result.output
|
|
assert "env=test" in result.output
|
|
assert "team=ml" in result.output
|
|
|
|
|
|
def test_get_experiment_table_no_tags():
|
|
exp_id = mlflow.create_experiment("test_get_exp_no_tags")
|
|
|
|
result = CliRunner().invoke(experiments.get_experiment, ["-x", exp_id, "--output", "table"])
|
|
assert result.exit_code == 0
|
|
|
|
assert "Experiment ID" in result.output
|
|
assert exp_id in result.output
|
|
assert "Tags" in result.output
|
|
|
|
|
|
def test_get_experiment_missing_id():
|
|
result = CliRunner().invoke(experiments.get_experiment, [])
|
|
assert result.exit_code != 0
|
|
assert "Must specify exactly one of --experiment-id or --experiment-name" in result.output
|
|
|
|
|
|
def test_get_experiment_invalid_id():
|
|
result = CliRunner().invoke(experiments.get_experiment, ["-x", "999999"])
|
|
assert result.exit_code != 0
|
|
|
|
|
|
def test_get_experiment_deleted():
|
|
exp_id = mlflow.create_experiment("test_deleted")
|
|
mlflow.delete_experiment(exp_id)
|
|
|
|
result = CliRunner().invoke(experiments.get_experiment, ["-x", exp_id, "--output", "json"])
|
|
assert result.exit_code == 0
|
|
|
|
output = json.loads(result.output)
|
|
assert output["lifecycle_stage"] == "deleted"
|
|
assert output["experiment_id"] == exp_id
|
|
|
|
|
|
def test_get_experiment_by_name_table():
|
|
exp_name = "test_get_by_name"
|
|
exp_id = mlflow.create_experiment(exp_name, tags={"env": "test"})
|
|
|
|
result = CliRunner().invoke(
|
|
experiments.get_experiment, ["--experiment-name", exp_name, "--output", "table"]
|
|
)
|
|
assert result.exit_code == 0
|
|
assert "Experiment ID" in result.output
|
|
assert exp_id in result.output
|
|
assert "Name" in result.output
|
|
assert exp_name in result.output
|
|
assert "Tags" in result.output
|
|
assert "env=test" in result.output
|
|
|
|
|
|
def test_get_experiment_by_name_json():
|
|
exp_name = "test_get_by_name_json"
|
|
exp_id = mlflow.create_experiment(exp_name, tags={"team": "ml"})
|
|
exp = mlflow.get_experiment(exp_id)
|
|
|
|
result = CliRunner().invoke(
|
|
experiments.get_experiment, ["--experiment-name", exp_name, "--output", "json"]
|
|
)
|
|
assert result.exit_code == 0
|
|
|
|
output = json.loads(result.output)
|
|
expected = {
|
|
"experiment_id": exp_id,
|
|
"name": exp_name,
|
|
"artifact_location": exp.artifact_location,
|
|
"lifecycle_stage": "active",
|
|
"tags": {"team": "ml"},
|
|
"creation_time": exp.creation_time,
|
|
"last_update_time": exp.last_update_time,
|
|
"effective_trace_archival_retention": exp.effective_trace_archival_retention,
|
|
"trace_location": exp.trace_location,
|
|
"workspace": "default",
|
|
}
|
|
assert output == expected
|
|
|
|
|
|
def test_get_experiment_by_name_short_option():
|
|
exp_name = "test_short_option"
|
|
exp_id = mlflow.create_experiment(exp_name)
|
|
|
|
result = CliRunner().invoke(experiments.get_experiment, ["-n", exp_name])
|
|
assert result.exit_code == 0
|
|
assert exp_id in result.output
|
|
assert exp_name in result.output
|
|
|
|
|
|
def test_get_experiment_by_name_not_found():
|
|
result = CliRunner().invoke(
|
|
experiments.get_experiment, ["--experiment-name", "nonexistent_experiment"]
|
|
)
|
|
assert result.exit_code != 0
|
|
|
|
|
|
def test_get_experiment_both_options_provided():
|
|
result = CliRunner().invoke(
|
|
experiments.get_experiment, ["--experiment-id", "0", "--experiment-name", "Default"]
|
|
)
|
|
assert result.exit_code != 0
|
|
assert "Must specify exactly one of --experiment-id or --experiment-name" in result.output
|
|
|
|
|
|
def test_get_experiment_by_name_deleted():
|
|
exp_name = "test_deleted_by_name"
|
|
exp_id = mlflow.create_experiment(exp_name)
|
|
mlflow.delete_experiment(exp_id)
|
|
|
|
result = CliRunner().invoke(
|
|
experiments.get_experiment, ["--experiment-name", exp_name, "--output", "json"]
|
|
)
|
|
assert result.exit_code == 0
|
|
|
|
output = json.loads(result.output)
|
|
assert output["lifecycle_stage"] == "deleted"
|
|
assert output["name"] == exp_name
|
|
|
|
|
|
def test_create_experiment_with_trace_archival_retention():
|
|
store = mock.Mock()
|
|
store.create_experiment.return_value = "123"
|
|
|
|
with patch("mlflow.experiments._get_store", return_value=store):
|
|
result = CliRunner().invoke(
|
|
experiments.create,
|
|
[
|
|
"--experiment-name",
|
|
"trace-policy-exp",
|
|
"--trace-archival-retention",
|
|
"30d",
|
|
],
|
|
)
|
|
|
|
assert result.exit_code == 0
|
|
_, kwargs = store.create_experiment.call_args
|
|
assert kwargs["tags"] is not None
|
|
assert len(kwargs["tags"]) == 1
|
|
assert kwargs["tags"][0].key == TraceExperimentTagKey.ARCHIVAL_RETENTION
|
|
assert kwargs["tags"][0].value == json.dumps({"type": "duration", "value": "30d"})
|
|
|
|
|
|
def test_update_experiment_sets_trace_archival_controls():
|
|
store = mock.Mock()
|
|
store.get_experiment.return_value = mock.Mock(tags={})
|
|
|
|
with patch("mlflow.experiments._get_store", return_value=store):
|
|
result = CliRunner().invoke(
|
|
experiments.update_experiment,
|
|
[
|
|
"--experiment-id",
|
|
"123",
|
|
"--trace-archival-retention",
|
|
"30d",
|
|
"--trace-archive-now-older-than",
|
|
"1d",
|
|
],
|
|
)
|
|
|
|
assert result.exit_code == 0
|
|
assert len(store.set_experiment_tag.call_args_list) == 2
|
|
first_tag = store.set_experiment_tag.call_args_list[0].args[1]
|
|
second_tag = store.set_experiment_tag.call_args_list[1].args[1]
|
|
assert first_tag.key == TraceExperimentTagKey.ARCHIVAL_RETENTION
|
|
assert first_tag.value == json.dumps({"type": "duration", "value": "30d"})
|
|
assert second_tag.key == TraceExperimentTagKey.ARCHIVE_NOW
|
|
assert second_tag.value == json.dumps({"older_than": "1d"})
|
|
|
|
|
|
def test_update_experiment_clears_trace_archival_controls():
|
|
store = mock.Mock()
|
|
store.get_experiment.return_value = mock.Mock(
|
|
tags={
|
|
TraceExperimentTagKey.ARCHIVAL_RETENTION: json.dumps({
|
|
"type": "duration",
|
|
"value": "30d",
|
|
}),
|
|
TraceExperimentTagKey.ARCHIVE_NOW: json.dumps({}),
|
|
}
|
|
)
|
|
|
|
with patch("mlflow.experiments._get_store", return_value=store):
|
|
result = CliRunner().invoke(
|
|
experiments.update_experiment,
|
|
[
|
|
"--experiment-id",
|
|
"123",
|
|
"--clear-trace-archival-retention",
|
|
"--clear-trace-archive-now",
|
|
],
|
|
)
|
|
|
|
assert result.exit_code == 0
|
|
assert store.delete_experiment_tag.call_args_list == [
|
|
mock.call("123", TraceExperimentTagKey.ARCHIVAL_RETENTION),
|
|
mock.call("123", TraceExperimentTagKey.ARCHIVE_NOW),
|
|
]
|
|
|
|
|
|
def test_update_experiment_rejects_conflicting_archive_now_flags():
|
|
result = CliRunner().invoke(
|
|
experiments.update_experiment,
|
|
[
|
|
"--experiment-id",
|
|
"123",
|
|
"--trace-archive-now",
|
|
"--trace-archive-now-older-than",
|
|
"1d",
|
|
],
|
|
)
|
|
|
|
assert result.exit_code != 0
|
|
assert "Cannot specify both --trace-archive-now and --trace-archive-now-older-than" in (
|
|
result.output
|
|
)
|