diff --git a/mlflow_mongodb/__init__.py b/mlflow_mongodb/__init__.py index e07197c..50d3688 100644 --- a/mlflow_mongodb/__init__.py +++ b/mlflow_mongodb/__init__.py @@ -1,6 +1,6 @@ """MongoDB model registry store plugin for MLflow.""" -from mlflow_mongodb.model_registry_store import MongoDBModelRegistryStore +from mlflow_mongodb.model_registry.store import MongoDBModelRegistryStore __all__ = ["MongoDBModelRegistryStore"] __version__ = "0.1.0.dev0" diff --git a/mlflow_mongodb/infrastructure/__init__.py b/mlflow_mongodb/infrastructure/__init__.py new file mode 100644 index 0000000..358f4b5 --- /dev/null +++ b/mlflow_mongodb/infrastructure/__init__.py @@ -0,0 +1 @@ +"""Shared infrastructure for the MongoDB MLflow stores.""" diff --git a/mlflow_mongodb/repositories/_helpers.py b/mlflow_mongodb/infrastructure/_array_updates.py similarity index 100% rename from mlflow_mongodb/repositories/_helpers.py rename to mlflow_mongodb/infrastructure/_array_updates.py diff --git a/mlflow_mongodb/settings.py b/mlflow_mongodb/infrastructure/settings.py similarity index 100% rename from mlflow_mongodb/settings.py rename to mlflow_mongodb/infrastructure/settings.py diff --git a/mlflow_mongodb/model_registry/__init__.py b/mlflow_mongodb/model_registry/__init__.py new file mode 100644 index 0000000..d6f546e --- /dev/null +++ b/mlflow_mongodb/model_registry/__init__.py @@ -0,0 +1 @@ +"""MongoDB-backed MLflow model-registry store.""" diff --git a/mlflow_mongodb/model_registry/errors.py b/mlflow_mongodb/model_registry/errors.py new file mode 100644 index 0000000..d04e9a8 --- /dev/null +++ b/mlflow_mongodb/model_registry/errors.py @@ -0,0 +1,17 @@ +"""Domain exceptions raised by the model-registry store and repositories.""" + + +class ModelVersionAlreadyExistsError(Exception): + """Raised when a model version number is already stored for a registered model.""" + + +class ModelVersionNotFoundError(Exception): + """Raised when a model version is not stored for a registered model.""" + + +class RegisteredModelAlreadyExistsError(Exception): + """Raised when a registered model name is already stored.""" + + +class RegisteredModelNotFoundError(Exception): + """Raised when a registered model name is not stored.""" diff --git a/mlflow_mongodb/model_registry/repositories/__init__.py b/mlflow_mongodb/model_registry/repositories/__init__.py new file mode 100644 index 0000000..3e9e2b9 --- /dev/null +++ b/mlflow_mongodb/model_registry/repositories/__init__.py @@ -0,0 +1,25 @@ +"""MongoDB repositories owned by the model-registry store.""" + +from mlflow_mongodb.model_registry.repositories.model_versions import ( + ModelVersionFilter, + ModelVersionOrder, + ModelVersionPage, + ModelVersionRepository, +) +from mlflow_mongodb.model_registry.repositories.registered_models import ( + RegisteredModelFilter, + RegisteredModelOrder, + RegisteredModelPage, + RegisteredModelRepository, +) + +__all__ = [ + "ModelVersionFilter", + "ModelVersionOrder", + "ModelVersionPage", + "ModelVersionRepository", + "RegisteredModelFilter", + "RegisteredModelOrder", + "RegisteredModelPage", + "RegisteredModelRepository", +] diff --git a/mlflow_mongodb/repositories/model_versions.py b/mlflow_mongodb/model_registry/repositories/model_versions.py similarity index 98% rename from mlflow_mongodb/repositories/model_versions.py rename to mlflow_mongodb/model_registry/repositories/model_versions.py index 165a6ff..3f27143 100644 --- a/mlflow_mongodb/repositories/model_versions.py +++ b/mlflow_mongodb/model_registry/repositories/model_versions.py @@ -16,16 +16,20 @@ from pymongo.database import Database from pymongo.errors import DuplicateKeyError -from mlflow_mongodb.repositories._helpers import ( +from mlflow_mongodb.infrastructure._array_updates import ( build_remove_array_element_update, build_replace_array_element_pipeline, ) -from mlflow_mongodb.repositories.types import ( +from mlflow_mongodb.infrastructure.settings import MongoDBSettings +from mlflow_mongodb.model_registry.errors import ( + ModelVersionAlreadyExistsError, + ModelVersionNotFoundError, +) +from mlflow_mongodb.model_registry.types import ( ModelVersionRecord, ModelVersionSearchResult, RegisteredModelRecord, ) -from mlflow_mongodb.settings import MongoDBSettings logger = logging.getLogger(__name__) @@ -61,14 +65,6 @@ class ModelVersionPage: has_more: bool -class ModelVersionAlreadyExistsError(Exception): - """Raised when a model version number is already stored for a registered model.""" - - -class ModelVersionNotFoundError(Exception): - """Raised when a model version is not stored for a registered model.""" - - class ModelVersionRepository: """Read and write model-version documents in MongoDB.""" diff --git a/mlflow_mongodb/repositories/registered_models.py b/mlflow_mongodb/model_registry/repositories/registered_models.py similarity index 98% rename from mlflow_mongodb/repositories/registered_models.py rename to mlflow_mongodb/model_registry/repositories/registered_models.py index f3e1068..736efbb 100644 --- a/mlflow_mongodb/repositories/registered_models.py +++ b/mlflow_mongodb/model_registry/repositories/registered_models.py @@ -13,16 +13,20 @@ from pymongo.database import Database from pymongo.errors import DuplicateKeyError -from mlflow_mongodb.repositories._helpers import ( +from mlflow_mongodb.infrastructure._array_updates import ( build_remove_array_element_update, build_replace_array_element_pipeline, ) -from mlflow_mongodb.repositories.types import ( +from mlflow_mongodb.infrastructure.settings import MongoDBSettings +from mlflow_mongodb.model_registry.errors import ( + RegisteredModelAlreadyExistsError, + RegisteredModelNotFoundError, +) +from mlflow_mongodb.model_registry.types import ( ModelVersionRecord, RegisteredModelDetails, RegisteredModelRecord, ) -from mlflow_mongodb.settings import MongoDBSettings logger = logging.getLogger(__name__) @@ -53,14 +57,6 @@ class RegisteredModelPage: has_more: bool -class RegisteredModelAlreadyExistsError(Exception): - """Raised when a registered model name is already stored.""" - - -class RegisteredModelNotFoundError(Exception): - """Raised when a registered model name is not stored.""" - - class RegisteredModelRepository: """Store registered-model documents in MongoDB.""" diff --git a/mlflow_mongodb/model_registry_store.py b/mlflow_mongodb/model_registry/store.py similarity index 99% rename from mlflow_mongodb/model_registry_store.py rename to mlflow_mongodb/model_registry/store.py index fdeaa69..a2f043f 100644 --- a/mlflow_mongodb/model_registry_store.py +++ b/mlflow_mongodb/model_registry/store.py @@ -54,22 +54,26 @@ from pymongo import MongoClient from pymongo.errors import ConfigurationError -from mlflow_mongodb.repositories import ( +from mlflow_mongodb.infrastructure.settings import MongoDBSettings +from mlflow_mongodb.model_registry.errors import ( ModelVersionAlreadyExistsError, - ModelVersionFilter, ModelVersionNotFoundError, + RegisteredModelAlreadyExistsError, + RegisteredModelNotFoundError, +) +from mlflow_mongodb.model_registry.repositories import ( + ModelVersionFilter, ModelVersionOrder, - ModelVersionRecord, ModelVersionRepository, - RegisteredModelAlreadyExistsError, - RegisteredModelDetails, RegisteredModelFilter, - RegisteredModelNotFoundError, RegisteredModelOrder, - RegisteredModelRecord, RegisteredModelRepository, ) -from mlflow_mongodb.settings import MongoDBSettings +from mlflow_mongodb.model_registry.types import ( + ModelVersionRecord, + RegisteredModelDetails, + RegisteredModelRecord, +) logger = logging.getLogger(__name__) diff --git a/mlflow_mongodb/repositories/types.py b/mlflow_mongodb/model_registry/types.py similarity index 98% rename from mlflow_mongodb/repositories/types.py rename to mlflow_mongodb/model_registry/types.py index dce0df1..28e1207 100644 --- a/mlflow_mongodb/repositories/types.py +++ b/mlflow_mongodb/model_registry/types.py @@ -1,4 +1,4 @@ -"""Shared persistence DTOs for the MongoDB repositories.""" +"""Persistence DTOs owned by the model-registry store.""" from collections.abc import Mapping from dataclasses import dataclass diff --git a/mlflow_mongodb/repositories/__init__.py b/mlflow_mongodb/repositories/__init__.py deleted file mode 100644 index 568e23f..0000000 --- a/mlflow_mongodb/repositories/__init__.py +++ /dev/null @@ -1,49 +0,0 @@ -"""MongoDB repositories used by the model registry store.""" - -from mlflow_mongodb.repositories.model_versions import ( - ModelVersionAlreadyExistsError, - ModelVersionFilter, - ModelVersionNotFoundError, - ModelVersionOrder, - ModelVersionPage, - ModelVersionRepository, -) -from mlflow_mongodb.repositories.registered_models import ( - RegisteredModelAlreadyExistsError, - RegisteredModelFilter, - RegisteredModelNotFoundError, - RegisteredModelOrder, - RegisteredModelPage, - RegisteredModelRepository, -) -from mlflow_mongodb.repositories.types import ( - ModelVersionRecord, - ModelVersionSearchResult, - ModelVersionTagRecord, - RegisteredModelAliasRecord, - RegisteredModelDetails, - RegisteredModelRecord, - RegisteredModelTagRecord, -) - -__all__ = [ - "ModelVersionAlreadyExistsError", - "ModelVersionFilter", - "ModelVersionNotFoundError", - "ModelVersionOrder", - "ModelVersionPage", - "ModelVersionRecord", - "ModelVersionRepository", - "ModelVersionSearchResult", - "ModelVersionTagRecord", - "RegisteredModelAliasRecord", - "RegisteredModelAlreadyExistsError", - "RegisteredModelDetails", - "RegisteredModelFilter", - "RegisteredModelNotFoundError", - "RegisteredModelOrder", - "RegisteredModelPage", - "RegisteredModelRecord", - "RegisteredModelRepository", - "RegisteredModelTagRecord", -] diff --git a/pyproject.toml b/pyproject.toml index 83e0ef7..3026a59 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,7 +52,7 @@ ignore = [ known-first-party = ["mlflow_mongodb"] [tool.ruff.lint.per-file-ignores] -"mlflow_mongodb/model_registry_store.py" = ["import-private-name"] +"mlflow_mongodb/model_registry/store.py" = ["import-private-name"] [tool.ruff.lint.pydocstyle] convention = "google" diff --git a/setup.py b/setup.py index eb9c4f0..f4f6713 100644 --- a/setup.py +++ b/setup.py @@ -32,8 +32,8 @@ }, entry_points={ "mlflow.model_registry_store": [ - "mongodb=mlflow_mongodb.model_registry_store:MongoDBModelRegistryStore", - "mongodb+srv=mlflow_mongodb.model_registry_store:MongoDBModelRegistryStore", + "mongodb=mlflow_mongodb.model_registry.store:MongoDBModelRegistryStore", + "mongodb+srv=mlflow_mongodb.model_registry.store:MongoDBModelRegistryStore", ], }, classifiers=[ diff --git a/tests/functional/test_indexes.py b/tests/functional/test_indexes.py index 47306fc..343c8ab 100644 --- a/tests/functional/test_indexes.py +++ b/tests/functional/test_indexes.py @@ -1,7 +1,7 @@ """Functional assertions for production MongoDB index definitions.""" from mlflow_mongodb import MongoDBModelRegistryStore -from mlflow_mongodb.repositories import ( +from mlflow_mongodb.model_registry.repositories import ( ModelVersionRepository, RegisteredModelRepository, ) diff --git a/tests/functional/test_model_version_contract.py b/tests/functional/test_model_version_contract.py index 44c065e..71f3838 100644 --- a/tests/functional/test_model_version_contract.py +++ b/tests/functional/test_model_version_contract.py @@ -20,7 +20,7 @@ def test_model_version_create_update_and_stage_validation_contract( name = "model-version-metadata-contract" clock = {"now": 1_700_001_000_000} monkeypatch.setattr( - "mlflow_mongodb.model_registry_store.get_current_time_millis", + "mlflow_mongodb.model_registry.store.get_current_time_millis", lambda: clock["now"], ) store.create_registered_model(name) @@ -129,7 +129,7 @@ def test_transition_with_archiving_rejects_inactive_stages_and_archives_peers( transition_timestamp = 1_700_002_000_000 monkeypatch.setattr( - "mlflow_mongodb.model_registry_store.get_current_time_millis", + "mlflow_mongodb.model_registry.store.get_current_time_millis", lambda: transition_timestamp, ) store.transition_model_version_stage(name, version_three.version, "STAGING", True) diff --git a/tests/functional/test_registered_model_contract.py b/tests/functional/test_registered_model_contract.py index ff633a2..8c1eef6 100644 --- a/tests/functional/test_registered_model_contract.py +++ b/tests/functional/test_registered_model_contract.py @@ -23,7 +23,7 @@ def test_registered_model_create_get_and_update_contract( ): clock = {"now": 1_700_000_000_000} monkeypatch.setattr( - "mlflow_mongodb.model_registry_store.get_current_time_millis", + "mlflow_mongodb.model_registry.store.get_current_time_millis", lambda: clock["now"], ) @@ -93,7 +93,7 @@ def test_rename_registered_model_preserves_versions_and_rejects_conflicts( rename_timestamp = 1_700_000_100_000 monkeypatch.setattr( - "mlflow_mongodb.model_registry_store.get_current_time_millis", + "mlflow_mongodb.model_registry.store.get_current_time_millis", lambda: rename_timestamp, ) diff --git a/tests/functional/test_search_contract.py b/tests/functional/test_search_contract.py index 305cea3..4fea68b 100644 --- a/tests/functional/test_search_contract.py +++ b/tests/functional/test_search_contract.py @@ -113,7 +113,7 @@ def test_search_model_versions_orders_and_paginates_stably( ): clock = {"now": 1_700_003_000_000} monkeypatch.setattr( - "mlflow_mongodb.model_registry_store.get_current_time_millis", + "mlflow_mongodb.model_registry.store.get_current_time_millis", lambda: clock["now"], ) names = [ @@ -294,7 +294,7 @@ def test_search_registered_models_orders_and_paginates_stably( ): clock = {"now": 1_700_004_000_000} monkeypatch.setattr( - "mlflow_mongodb.model_registry_store.get_current_time_millis", + "mlflow_mongodb.model_registry.store.get_current_time_millis", lambda: clock["now"], ) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 48b72a1..cd4ab10 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -10,14 +10,16 @@ from mlflow.tracking.client import MlflowClient from mlflow_mongodb import MongoDBModelRegistryStore -from mlflow_mongodb.repositories import ( - ModelVersionRecord, +from mlflow_mongodb.model_registry.repositories import ( ModelVersionRepository, + RegisteredModelRepository, +) +from mlflow_mongodb.model_registry.types import ( + ModelVersionRecord, ModelVersionTagRecord, RegisteredModelAliasRecord, RegisteredModelDetails, RegisteredModelRecord, - RegisteredModelRepository, RegisteredModelTagRecord, ) @@ -48,7 +50,7 @@ def tracking_client(): def store(monkeypatch, registered_model_repository, model_version_repository, tracking_client): """Return a store wired to mocks so unit tests never connect to MongoDB.""" monkeypatch.setattr( - "mlflow_mongodb.model_registry_store.get_current_time_millis", + "mlflow_mongodb.model_registry.store.get_current_time_millis", lambda: FIXED_TIMESTAMP, ) diff --git a/tests/unit/test_prompt_store.py b/tests/unit/test_prompt_store.py index db7fccc..98ecf85 100644 --- a/tests/unit/test_prompt_store.py +++ b/tests/unit/test_prompt_store.py @@ -4,7 +4,7 @@ from mlflow.exceptions import MlflowException from mlflow.prompt.constants import IS_PROMPT_TAG_KEY, PROMPT_TEXT_TAG_KEY -from mlflow_mongodb.repositories import ( +from mlflow_mongodb.model_registry.errors import ( RegisteredModelAlreadyExistsError, RegisteredModelNotFoundError, ) diff --git a/tests/unit/test_store.py b/tests/unit/test_store.py index 4bcc6db..cf03479 100644 --- a/tests/unit/test_store.py +++ b/tests/unit/test_store.py @@ -15,7 +15,7 @@ ) from mlflow_mongodb import MongoDBModelRegistryStore -from mlflow_mongodb.repositories import ( +from mlflow_mongodb.model_registry.errors import ( ModelVersionAlreadyExistsError, ModelVersionNotFoundError, RegisteredModelAlreadyExistsError, @@ -33,7 +33,7 @@ def test_tracking_client_is_lazy_and_cached(monkeypatch): tracking_client = Mock() client_factory = Mock(return_value=tracking_client) monkeypatch.setattr( - "mlflow_mongodb.model_registry_store.MlflowClient", + "mlflow_mongodb.model_registry.store.MlflowClient", client_factory, ) diff --git a/tests/unit/test_store_configuration.py b/tests/unit/test_store_configuration.py index cff016f..fa0aa20 100644 --- a/tests/unit/test_store_configuration.py +++ b/tests/unit/test_store_configuration.py @@ -6,7 +6,7 @@ from mlflow.exceptions import MlflowException from mlflow_mongodb import MongoDBModelRegistryStore -from mlflow_mongodb.settings import MongoDBSettings +from mlflow_mongodb.infrastructure.settings import MongoDBSettings def test_mongo_client_requires_registry_uri(): diff --git a/tests/unit/test_store_search.py b/tests/unit/test_store_search.py index 70ea9a9..964328d 100644 --- a/tests/unit/test_store_search.py +++ b/tests/unit/test_store_search.py @@ -7,7 +7,7 @@ from mlflow.utils.search_utils import SearchModelUtils, SearchModelVersionUtils, SearchUtils from mlflow_mongodb import MongoDBModelRegistryStore -from mlflow_mongodb.repositories import ( +from mlflow_mongodb.model_registry.repositories import ( ModelVersionFilter, ModelVersionOrder, RegisteredModelFilter,