Implement a function to clone a model (DAT-1045)

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()

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)

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
 +++++++++++++++++++++++++