From 410acd3f6760e9a69144d9f320a289e23462c76b Mon Sep 17 00:00:00 2001 From: luandalmazo Date: Wed, 5 Aug 2026 14:23:32 -0300 Subject: [PATCH] add clone_model function --- datamint/api/endpoints/models_api.py | 115 ++++++++++++++++++++++++++- datamint/mlflow/tracking/fluent.py | 10 +++ docs/source/client_api_content.rst | 43 +++++++++- 3 files changed, 166 insertions(+), 2 deletions(-) diff --git a/datamint/api/endpoints/models_api.py b/datamint/api/endpoints/models_api.py index 1ee638fe..92c08d9b 100644 --- a/datamint/api/endpoints/models_api.py +++ b/datamint/api/endpoints/models_api.py @@ -1,11 +1,18 @@ """API handler for the model registry, backed by MLflow.""" +from collections.abc import Sequence +from typing import TYPE_CHECKING + import httpx import mlflow.exceptions import mlflow.tracking +from datamint.exceptions import ItemNotFoundError from ..entity_base_api import ApiConfig, BaseApi from .deploy_model_api import DeployModelApi -from .model_types import Model +from .model_types import Model, ModelVersion + +if TYPE_CHECKING: + from datamint.entities.project import Project class ModelsApi(BaseApi): @@ -102,3 +109,109 @@ def delete_registered_model(self, name: str) -> None: name: Name of the registered model to delete. """ self._mlflow_client.delete_registered_model(name) + + def _resolve_version(self, + model: 'str | Model | ModelVersion', + version: str | int | None, + alias: str | None) -> ModelVersion: + if isinstance(model, ModelVersion): + if version is not None or alias is not None: + raise TypeError("'version'/'alias' must not be passed when 'model' is already a ModelVersion.") + return model + + if (version is None) == (alias is None): + raise TypeError("clone_model() requires exactly one of 'version' or 'alias'.") + + if isinstance(model, Model): + registered_model = model + else: + registered_model = self.get_by_name(model) + if registered_model is None: + raise ItemNotFoundError('Model', {'name': model}) + + if alias is not None: + resolved = registered_model.get_latest_version(alias=alias) + if resolved is None: + raise ItemNotFoundError('ModelVersion', {'name': registered_model.name, 'alias': alias}) + return resolved + + for v in registered_model.get_versions(): + if str(v.version) == str(version): + return v + raise ItemNotFoundError('ModelVersion', {'name': registered_model.name, 'version': version}) + + def clone_model(self, + model: 'str | Model | ModelVersion', + target_project: 'str | Project', + *, + version: str | int | None = None, + alias: str | None = None, + target_model_name: str | None = None, + code_paths: 'Sequence[str] | None' = None) -> Model: + """Clone a model version from the active project into another project. + + The source version is resolved against whichever project is currently + active (see :func:`datamint.mlflow.set_project`) -- same as every other + ``ModelsApi`` lookup. Only models logged with the ``datamint`` MLflow + flavor are supported, since that's what carries the task type, supported + modes, and annotation specs this method copies over. + + Args: + model: Registered model name, :class:`~.model_types.Model`, or a specific + :class:`~.model_types.ModelVersion` to clone. When a name or ``Model`` + is passed, exactly one of ``version``/``alias`` must also be given. + target_project: Project (name, ID, or :class:`~datamint.entities.project.Project`) + to register the cloned model under. + version: Version number to clone. Mutually exclusive with ``alias``. + alias: Alias to clone (e.g. ``"champion"``). Mutually exclusive with ``version``. + target_model_name: Name to register the clone under in ``target_project``. + Defaults to the source model's name. + code_paths: Local paths to custom code files/dirs the model's class depends + on (same meaning as ``datamint_flavor.log_model``'s ``code_paths``). + + Returns: + The newly registered :class:`~.model_types.Model` in ``target_project``. + + Raises: + ValueError: If the source version wasn't logged with the ``datamint`` flavor. + """ + + from datamint.mlflow.flavors import datamint_flavor + from datamint.mlflow.flavors.datamint_flavor import FLAVOR_NAME + from datamint.mlflow.tracking.fluent import _reset_active_project, get_active_project_id, set_project + + source_version = self._resolve_version(model, version, alias) + + model_info = mlflow.models.get_model_info(source_version.source) + if FLAVOR_NAME not in model_info.flavors: + raise ValueError( + f"Model {source_version.name!r} version {source_version.version!r} was not logged " + f"with the 'datamint' MLflow flavor; clone_model() only supports datamint-flavor models." + ) + + task_type = source_version.get_task_type() + supported_modes = source_version.get_supported_modes() + annotation_specs = source_version.get_annotation_specs() + loaded_model = datamint_flavor.load_model(source_version.source) + + target_name = target_model_name or source_version.name + + previous_project_id = get_active_project_id() + try: + set_project(target_project) + mlflow.set_experiment(target_name) + with mlflow.start_run(run_name=f"clone_{source_version.name}_v{source_version.version}"): + datamint_flavor.log_model( + loaded_model, + task_type=task_type, + supported_modes=supported_modes, + annotation_specs=annotation_specs, + model_name=target_name, + code_paths=code_paths, + ) + return self.get_by_name(target_name) + finally: + if previous_project_id is not None: + set_project(previous_project_id) + else: + _reset_active_project() diff --git a/datamint/mlflow/tracking/fluent.py b/datamint/mlflow/tracking/fluent.py index 6bf6115f..643a4898 100644 --- a/datamint/mlflow/tracking/fluent.py +++ b/datamint/mlflow/tracking/fluent.py @@ -92,3 +92,13 @@ def set_project(project: 'Project | str'): os.environ[EnvVars.DATAMINT_PROJECT_ID.value] = project_id return project + + +def _reset_active_project(): + """Clear the active project, restoring the pre-``set_project()`` state. """ + global _ACTIVE_PROJECT_ID + + with _PROJECT_LOCK: + _ACTIVE_PROJECT_ID = None + + os.environ.pop(EnvVars.DATAMINT_PROJECT_ID.value, None) diff --git a/docs/source/client_api_content.rst b/docs/source/client_api_content.rst index 9de1d32e..83e2afec 100644 --- a/docs/source/client_api_content.rst +++ b/docs/source/client_api_content.rst @@ -515,7 +515,7 @@ a Datamint :mod:`~datamint.lightning.trainers`), rather than raising. to the latest version when you don't need a specific one. Model registry (MLflow) operations -+++++++++++++++++++++++++++++++++ +++++++++++++++++++++++++++++++++++ ``api.models`` wraps the underlying MLflow model registry client directly, so these calls map one-to-one onto MLflow's own registry API: @@ -542,6 +542,47 @@ these calls map one-to-one onto MLflow's own registry API: # Delete a registered model and all of its remaining versions api.models.delete_registered_model("my-model") +Clone a model to another project +++++++++++++++++++++++++++++++++ + +Models live in the MLflow registry of whichever project is active (see +:func:`datamint.mlflow.set_project`), so using a model trained in one project +against another normally means manually reloading and re-logging it. +``api.models.clone_model()`` does that for you: + +.. code-block:: python + + from datamint.mlflow import set_project + + set_project("Project A") # clone_model() resolves the source here + + cloned = api.models.clone_model( + "my-model", + target_project="Project B", + version=3, # or alias="champion" + target_model_name="my-model-v2", # optional, defaults to the source name + ) + +Only models logged with the ``datamint`` MLflow flavor are supported, since +that's what carries the task type, supported modes, and annotation specs the +clone copies over. If the source model's class depends on custom code (not an +installed package), pass ``code_paths`` the same way you would to +``log_model()`` -- otherwise the clone registers successfully but fails to +load later, since its class gets pickled by reference to code the new +artifact never bundled: + +.. code-block:: python + + api.models.clone_model( + "my-model", + target_project="Project B", + version=3, + code_paths=["my_adapter.py"], + ) + +The project that was active before the call is always restored afterward, +even if cloning fails partway through. + Deploy a registered model +++++++++++++++++++++++++