wehub-resource-sync
142 lines
4.8 KiB
Python
142 lines
4.8 KiB
Python
from unittest import mock
|
|
|
|
import pytest
|
|
|
|
import mlflow
|
|
import mlflow.tracking.fluent as fluent_module
|
|
from mlflow.entities import Experiment
|
|
from mlflow.entities.experiment_tag import ExperimentTag
|
|
from mlflow.entities.trace_location import UnityCatalog
|
|
from mlflow.exceptions import MlflowException
|
|
|
|
|
|
def _experiment(experiment_id="789"):
|
|
return Experiment(
|
|
experiment_id=experiment_id,
|
|
name="test-experiment",
|
|
artifact_location="file:/tmp",
|
|
lifecycle_stage="active",
|
|
tags=[ExperimentTag("key", "val")],
|
|
)
|
|
|
|
|
|
def test_without_trace_location_unchanged():
|
|
with mock.patch("mlflow.tracking.fluent.MlflowClient") as mock_client_cls:
|
|
client = mock_client_cls.return_value
|
|
client.create_experiment.return_value = "789"
|
|
|
|
assert mlflow.create_experiment("my-exp", tags={"k": "v"}) == "789"
|
|
client.create_experiment.assert_called_once_with("my-exp", None, {"k": "v"})
|
|
client.get_experiment.assert_not_called()
|
|
|
|
|
|
def test_defaults_empty_prefix_to_experiment_id():
|
|
with (
|
|
mock.patch("mlflow.tracking.fluent.MlflowClient") as mock_client_cls,
|
|
mock.patch(
|
|
"mlflow.tracking.fluent._resolve_experiment_to_trace_location",
|
|
return_value=None,
|
|
) as mock_resolve,
|
|
):
|
|
client = mock_client_cls.return_value
|
|
client.create_experiment.return_value = "789"
|
|
client.get_experiment.return_value = _experiment()
|
|
|
|
original = UnityCatalog("catalog", "schema")
|
|
mlflow.create_experiment("my-exp", trace_location=original)
|
|
|
|
passed_location = mock_resolve.call_args.kwargs["trace_location"]
|
|
assert passed_location.table_prefix == "789"
|
|
assert original.table_prefix is None
|
|
|
|
|
|
def test_explicit_prefix_preserved():
|
|
with (
|
|
mock.patch("mlflow.tracking.fluent.MlflowClient") as mock_client_cls,
|
|
mock.patch(
|
|
"mlflow.tracking.fluent._resolve_experiment_to_trace_location",
|
|
return_value=None,
|
|
) as mock_resolve,
|
|
):
|
|
client = mock_client_cls.return_value
|
|
client.create_experiment.return_value = "789"
|
|
client.get_experiment.return_value = _experiment()
|
|
|
|
mlflow.create_experiment(
|
|
"my-exp",
|
|
trace_location=UnityCatalog("catalog", "schema", "custom_prefix"),
|
|
)
|
|
|
|
assert mock_resolve.call_args.kwargs["trace_location"].table_prefix == "custom_prefix"
|
|
|
|
|
|
def test_resolve_receives_experiment_object():
|
|
exp = _experiment()
|
|
|
|
with (
|
|
mock.patch("mlflow.tracking.fluent.MlflowClient") as mock_client_cls,
|
|
mock.patch(
|
|
"mlflow.tracking.fluent._resolve_experiment_to_trace_location",
|
|
return_value=None,
|
|
) as mock_resolve,
|
|
):
|
|
client = mock_client_cls.return_value
|
|
client.create_experiment.return_value = "789"
|
|
client.get_experiment.return_value = exp
|
|
|
|
result = mlflow.create_experiment(
|
|
"my-exp",
|
|
trace_location=UnityCatalog("catalog", "schema", "pfx"),
|
|
)
|
|
|
|
assert result == "789"
|
|
assert mock_resolve.call_args.kwargs["experiment"] is exp
|
|
|
|
|
|
def test_does_not_sync_provider_or_set_active():
|
|
original_active = fluent_module._active_experiment_id
|
|
|
|
with (
|
|
mock.patch("mlflow.tracking.fluent.MlflowClient") as mock_client_cls,
|
|
mock.patch(
|
|
"mlflow.tracking.fluent._resolve_experiment_to_trace_location",
|
|
return_value=None,
|
|
),
|
|
mock.patch(
|
|
"mlflow.tracking.fluent._sync_trace_destination_and_provider",
|
|
) as mock_sync,
|
|
):
|
|
client = mock_client_cls.return_value
|
|
client.create_experiment.return_value = "789"
|
|
client.get_experiment.return_value = _experiment()
|
|
|
|
mlflow.create_experiment(
|
|
"my-exp",
|
|
trace_location=UnityCatalog("catalog", "schema", "pfx"),
|
|
)
|
|
|
|
mock_sync.assert_not_called()
|
|
assert fluent_module._active_experiment_id == original_active
|
|
|
|
|
|
def test_link_failure_includes_retry_guidance():
|
|
with (
|
|
mock.patch("mlflow.tracking.fluent.MlflowClient") as mock_client_cls,
|
|
mock.patch(
|
|
"mlflow.tracking.fluent._resolve_experiment_to_trace_location",
|
|
side_effect=MlflowException("backend error"),
|
|
),
|
|
):
|
|
client = mock_client_cls.return_value
|
|
client.create_experiment.return_value = "456"
|
|
client.get_experiment.return_value = _experiment(experiment_id="456")
|
|
|
|
with pytest.raises(MlflowException, match="delete the experiment and retry") as exc_info:
|
|
mlflow.create_experiment(
|
|
"new-exp",
|
|
trace_location=UnityCatalog("cat", "sch", "pfx"),
|
|
)
|
|
|
|
assert "was created" in exc_info.value.message
|
|
assert "backend error" in exc_info.value.message
|