wehub-resource-sync
342 lines
12 KiB
Python
342 lines
12 KiB
Python
import io
|
|
import pickle
|
|
from unittest import mock
|
|
|
|
import pytest
|
|
|
|
from mlflow.environment_variables import MLFLOW_TRACKING_URI
|
|
from mlflow.store._unity_catalog.registry.rest_store import UcModelRegistryStore
|
|
from mlflow.store.db.db_types import DATABASE_ENGINES
|
|
from mlflow.store.model_registry.rest_store import RestStore
|
|
from mlflow.store.model_registry.sqlalchemy_store import SqlAlchemyStore
|
|
from mlflow.tracking._model_registry.utils import (
|
|
_get_store,
|
|
_resolve_registry_uri,
|
|
get_registry_uri,
|
|
set_registry_uri,
|
|
)
|
|
from mlflow.tracking.registry import UnsupportedModelRegistryStoreURIException
|
|
|
|
from tests.helpers.db_mocks import mock_get_managed_session_maker
|
|
|
|
# Disable mocking tracking URI here, as we want to test setting the tracking URI via
|
|
# environment variable. See
|
|
# http://doc.pytest.org/en/latest/skipping.html#skip-all-test-functions-of-a-class-or-module
|
|
# and https://github.com/mlflow/mlflow/blob/master/CONTRIBUTING.md#writing-python-tests
|
|
# for more information.
|
|
pytestmark = pytest.mark.notrackingurimock
|
|
|
|
|
|
@pytest.fixture
|
|
def reset_registry_uri():
|
|
yield
|
|
set_registry_uri(None)
|
|
|
|
|
|
def test_set_get_registry_uri():
|
|
with mock.patch(
|
|
"mlflow.tracking._model_registry.utils._resolve_tracking_uri"
|
|
) as get_tracking_uri_mock:
|
|
get_tracking_uri_mock.return_value = "databricks://tracking_sldkfj"
|
|
uri = "databricks://registry/path"
|
|
set_registry_uri(uri)
|
|
assert get_registry_uri() == uri
|
|
set_registry_uri(None)
|
|
|
|
|
|
def test_set_get_empty_registry_uri():
|
|
with mock.patch(
|
|
"mlflow.tracking._model_registry.utils._resolve_tracking_uri"
|
|
) as get_tracking_uri_mock:
|
|
get_tracking_uri_mock.return_value = None
|
|
set_registry_uri("")
|
|
assert get_registry_uri() is None
|
|
set_registry_uri(None)
|
|
|
|
|
|
def test_default_get_registry_uri_no_tracking_uri():
|
|
with mock.patch(
|
|
"mlflow.tracking._model_registry.utils._resolve_tracking_uri"
|
|
) as get_tracking_uri_mock:
|
|
get_tracking_uri_mock.return_value = None
|
|
set_registry_uri(None)
|
|
assert get_registry_uri() is None
|
|
|
|
|
|
def test_default_get_registry_uri_with_databricks_tracking_uri_defaults_to_uc():
|
|
tracking_uri = "databricks://tracking_werohoz"
|
|
with mock.patch(
|
|
"mlflow.tracking._model_registry.utils._resolve_tracking_uri"
|
|
) as resolve_tracking_uri_mock:
|
|
resolve_tracking_uri_mock.return_value = tracking_uri
|
|
set_registry_uri(None)
|
|
# Should default to Unity Catalog when tracking URI starts with 'databricks'
|
|
# and include the profile when present
|
|
assert get_registry_uri() == "databricks-uc://tracking_werohoz"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"tracking_uri",
|
|
[
|
|
"http://localhost:5000",
|
|
"https://remote-server.com",
|
|
"sqlite:///path/to/db.sqlite",
|
|
"postgresql://user:pass@localhost/db",
|
|
"file:///local/path",
|
|
],
|
|
)
|
|
def test_default_registry_uri_non_databricks_tracking_uri(tracking_uri):
|
|
with mock.patch(
|
|
"mlflow.tracking._tracking_service.utils.get_tracking_uri"
|
|
) as get_tracking_uri_mock:
|
|
get_tracking_uri_mock.return_value = tracking_uri
|
|
set_registry_uri(None)
|
|
# Non-databricks URIs should be used directly as registry URI
|
|
assert get_registry_uri() == tracking_uri
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("tracking_uri", "expected_registry_uri"),
|
|
[
|
|
("databricks", "databricks-uc"),
|
|
("databricks://profile", "databricks-uc://profile"),
|
|
("databricks://profile_name", "databricks-uc://profile_name"),
|
|
("databricks://workspace_url", "databricks-uc://workspace_url"),
|
|
(
|
|
"databricks://some.databricks.workspace.com",
|
|
"databricks-uc://some.databricks.workspace.com",
|
|
),
|
|
],
|
|
)
|
|
def test_databricks_tracking_uri_variations_default_to_uc(tracking_uri, expected_registry_uri):
|
|
with mock.patch(
|
|
"mlflow.tracking._tracking_service.utils.get_tracking_uri"
|
|
) as get_tracking_uri_mock:
|
|
get_tracking_uri_mock.return_value = tracking_uri
|
|
set_registry_uri(None)
|
|
# All databricks tracking URIs should default to Unity Catalog
|
|
registry_uri = get_registry_uri()
|
|
assert registry_uri == expected_registry_uri
|
|
|
|
|
|
def test_explicit_registry_uri_overrides_databricks_default():
|
|
tracking_uri = "databricks://workspace"
|
|
explicit_registry_uri = "databricks://different_workspace"
|
|
|
|
with mock.patch(
|
|
"mlflow.tracking._tracking_service.utils.get_tracking_uri"
|
|
) as get_tracking_uri_mock:
|
|
get_tracking_uri_mock.return_value = tracking_uri
|
|
set_registry_uri(explicit_registry_uri)
|
|
# Explicit registry URI should override the databricks-uc default
|
|
assert get_registry_uri() == explicit_registry_uri
|
|
set_registry_uri(None) # Reset for other tests
|
|
|
|
|
|
def test_registry_uri_from_environment_overrides_databricks_default():
|
|
from mlflow.environment_variables import MLFLOW_REGISTRY_URI
|
|
|
|
tracking_uri = "databricks://workspace"
|
|
env_registry_uri = "http://env-registry-server:5000"
|
|
|
|
with (
|
|
mock.patch(
|
|
"mlflow.tracking._tracking_service.utils.get_tracking_uri"
|
|
) as get_tracking_uri_mock,
|
|
mock.patch.object(MLFLOW_REGISTRY_URI, "get", return_value=env_registry_uri),
|
|
):
|
|
get_tracking_uri_mock.return_value = tracking_uri
|
|
set_registry_uri(None)
|
|
# Environment variable should override the databricks-uc default
|
|
assert get_registry_uri() == env_registry_uri
|
|
|
|
|
|
def test_registry_uri_from_spark_session_overrides_databricks_default():
|
|
tracking_uri = "databricks://workspace"
|
|
spark_registry_uri = "databricks-uc://spark_profile"
|
|
|
|
with (
|
|
mock.patch(
|
|
"mlflow.tracking._tracking_service.utils.get_tracking_uri"
|
|
) as get_tracking_uri_mock,
|
|
mock.patch(
|
|
"mlflow.tracking._model_registry.utils._get_registry_uri_from_spark_session"
|
|
) as get_spark_uri_mock,
|
|
):
|
|
get_tracking_uri_mock.return_value = tracking_uri
|
|
get_spark_uri_mock.return_value = spark_registry_uri
|
|
set_registry_uri(None)
|
|
# Spark session URI should override the databricks-uc default
|
|
assert get_registry_uri() == spark_registry_uri
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("tracking_uri", "expected_result"),
|
|
[
|
|
("mydatabricks://custom", None), # Should not match partial
|
|
("databricks", "databricks-uc"), # Should match exact
|
|
("", None), # Empty string should return None
|
|
],
|
|
)
|
|
def test_edge_cases_for_databricks_uri_detection(tracking_uri, expected_result):
|
|
with mock.patch(
|
|
"mlflow.tracking._tracking_service.utils.get_tracking_uri"
|
|
) as get_tracking_uri_mock:
|
|
get_tracking_uri_mock.return_value = tracking_uri
|
|
set_registry_uri(None)
|
|
result = get_registry_uri()
|
|
if expected_result is None:
|
|
assert result == tracking_uri # Should fallback to tracking URI
|
|
else:
|
|
assert result == expected_result
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("tracking_uri", "registry_uri_param", "expected_result"),
|
|
[
|
|
# (tracking_uri, registry_uri_param, expected_result)
|
|
("databricks://workspace", None, "databricks-uc://workspace"),
|
|
("databricks", None, "databricks-uc"),
|
|
("http://localhost:5000", None, "http://localhost:5000"),
|
|
("databricks://workspace", "explicit://registry", "explicit://registry"),
|
|
(None, None, None),
|
|
("", None, ""),
|
|
],
|
|
)
|
|
def test_resolve_registry_uri_consistency_with_get_registry_uri(
|
|
tracking_uri, registry_uri_param, expected_result
|
|
):
|
|
with mock.patch(
|
|
"mlflow.tracking._model_registry.utils._resolve_tracking_uri"
|
|
) as mock_resolve_tracking:
|
|
mock_resolve_tracking.return_value = tracking_uri
|
|
set_registry_uri(None) # Clear context
|
|
|
|
result = _resolve_registry_uri(registry_uri_param, tracking_uri)
|
|
assert result == expected_result, (
|
|
f"Failed for tracking_uri={tracking_uri}, registry_uri={registry_uri_param}"
|
|
)
|
|
|
|
|
|
def test_resolve_registry_uri_with_environment_variable():
|
|
from mlflow.environment_variables import MLFLOW_REGISTRY_URI
|
|
|
|
env_registry_uri = "http://env-registry:5000"
|
|
tracking_uri = "databricks://workspace"
|
|
|
|
with (
|
|
mock.patch(
|
|
"mlflow.tracking._model_registry.utils._resolve_tracking_uri"
|
|
) as mock_resolve_tracking,
|
|
mock.patch.object(MLFLOW_REGISTRY_URI, "get", return_value=env_registry_uri),
|
|
):
|
|
mock_resolve_tracking.return_value = tracking_uri
|
|
set_registry_uri(None) # Clear explicit setting
|
|
|
|
# Environment variable should override databricks default
|
|
result = _resolve_registry_uri(None, tracking_uri)
|
|
assert result == env_registry_uri
|
|
|
|
|
|
def test_resolve_registry_uri_with_spark_session():
|
|
spark_registry_uri = "databricks-uc://spark_profile"
|
|
tracking_uri = "databricks://workspace"
|
|
|
|
with (
|
|
mock.patch(
|
|
"mlflow.tracking._model_registry.utils._resolve_tracking_uri"
|
|
) as mock_resolve_tracking,
|
|
mock.patch(
|
|
"mlflow.tracking._model_registry.utils._get_registry_uri_from_spark_session"
|
|
) as mock_spark_uri,
|
|
):
|
|
mock_resolve_tracking.return_value = tracking_uri
|
|
mock_spark_uri.return_value = spark_registry_uri
|
|
set_registry_uri(None) # Clear explicit setting
|
|
|
|
# Spark session URI should override databricks default
|
|
result = _resolve_registry_uri(None, tracking_uri)
|
|
assert result == spark_registry_uri
|
|
|
|
|
|
def test_get_store_rest_store_from_arg(monkeypatch):
|
|
monkeypatch.setenv(MLFLOW_TRACKING_URI.name, "https://my-tracking-server:5050")
|
|
store = _get_store("http://some/path")
|
|
assert isinstance(store, RestStore)
|
|
assert store.get_host_creds().host == "http://some/path"
|
|
|
|
|
|
def test_fallback_to_tracking_store(monkeypatch):
|
|
monkeypatch.setenv(MLFLOW_TRACKING_URI.name, "https://my-tracking-server:5050")
|
|
store = _get_store()
|
|
assert isinstance(store, RestStore)
|
|
assert store.get_host_creds().host == "https://my-tracking-server:5050"
|
|
assert store.get_host_creds().token is None
|
|
|
|
|
|
@pytest.mark.parametrize("db_type", DATABASE_ENGINES)
|
|
def test_get_store_sqlalchemy_store(db_type, monkeypatch):
|
|
uri = f"{db_type}://hostname/database"
|
|
monkeypatch.setenv(MLFLOW_TRACKING_URI.name, uri)
|
|
monkeypatch.delenv("MLFLOW_SQLALCHEMYSTORE_POOLCLASS", raising=False)
|
|
with (
|
|
mock.patch("sqlalchemy.create_engine") as mock_create_engine,
|
|
mock.patch("sqlalchemy.event.listens_for"),
|
|
mock.patch("mlflow.store.db.utils._initialize_tables"),
|
|
mock.patch("mlflow.store.model_registry.sqlalchemy_store._all_tables_exist"),
|
|
mock.patch(
|
|
"mlflow.store.model_registry.sqlalchemy_store._get_managed_session_maker",
|
|
new=mock_get_managed_session_maker,
|
|
),
|
|
mock.patch(
|
|
"mlflow.store.model_registry.sqlalchemy_store.SqlAlchemyStore."
|
|
"_verify_registry_tables_exist"
|
|
),
|
|
):
|
|
store = _get_store()
|
|
assert isinstance(store, SqlAlchemyStore)
|
|
assert store.db_uri == uri
|
|
|
|
mock_create_engine.assert_called_once_with(uri, pool_pre_ping=True)
|
|
|
|
|
|
@pytest.mark.parametrize("bad_uri", ["badsql://imfake", "yoursql://hi"])
|
|
def test_get_store_bad_uris(bad_uri, monkeypatch):
|
|
monkeypatch.setenv(MLFLOW_TRACKING_URI.name, bad_uri)
|
|
with pytest.raises(
|
|
UnsupportedModelRegistryStoreURIException,
|
|
match="Model registry functionality is unavailable",
|
|
):
|
|
_get_store()
|
|
|
|
|
|
def test_get_store_caches_on_store_uri(tmp_path):
|
|
store_uri_1 = f"sqlite:///{tmp_path.joinpath('store1.db')}"
|
|
store_uri_2 = f"sqlite:///{tmp_path.joinpath('store2.db')}"
|
|
|
|
store1 = _get_store(store_uri_1)
|
|
store2 = _get_store(store_uri_1)
|
|
assert store1 is store2
|
|
|
|
store3 = _get_store(store_uri_2)
|
|
store4 = _get_store(store_uri_2)
|
|
assert store3 is store4
|
|
|
|
assert store1 is not store3
|
|
|
|
|
|
@pytest.mark.parametrize("store_uri", ["databricks-uc", "databricks-uc://profile"])
|
|
def test_get_store_uc_registry_uri(store_uri):
|
|
assert isinstance(_get_store(store_uri), UcModelRegistryStore)
|
|
|
|
|
|
def test_store_object_can_be_serialized_by_pickle():
|
|
"""
|
|
This test ensures a store object generated by `_get_store` can be serialized by pickle
|
|
to prevent issues such as https://github.com/mlflow/mlflow/issues/2954
|
|
"""
|
|
pickle.dump(_get_store("https://example.com"), io.BytesIO())
|
|
pickle.dump(_get_store("databricks"), io.BytesIO())
|
|
# pickle.dump(_get_store(f"sqlite:///{tmpdir.strpath}/mlflow.db"), io.BytesIO())
|
|
# This throws `AttributeError: Can't pickle local object 'create_engine.<locals>.connect'`
|