datamint 2.3.4__py3-none-any.whl → 2.4.0__py3-none-any.whl

datamint/mlflow/lightning/callbacks/modelcheckpoint.py

@@ -0,0 +1,338 @@
+from lightning.pytorch.callbacks import ModelCheckpoint
+from pathlib import Path
+from weakref import proxy
+from mlflow.store.artifact.artifact_repository_registry import get_artifact_repository
+from typing import Literal, Any
+import inspect
+import torch
+from torch import nn
+import lightning.pytorch as L
+from datamint.mlflow.models import log_model_metadata, _get_MLFlowLogger
+from datamint.mlflow.env_utils import ensure_mlflow_configured
+import mlflow
+import logging
+from lightning.pytorch.loggers import MLFlowLogger
+
+_LOGGER = logging.getLogger(__name__)
+
+
+def help_infer_signature(x):
+    if isinstance(x, torch.Tensor):
+        return x.detach().cpu().numpy()
+    elif isinstance(x, dict):
+        return {k: v.detach().cpu().numpy() if isinstance(v, torch.Tensor) else v for k, v in x.items()}
+    elif isinstance(x, list):
+        return [v.detach().cpu().numpy() if isinstance(v, torch.Tensor) else v for v in x]
+    elif isinstance(x, tuple):
+        return tuple(v.detach().cpu().numpy() if isinstance(v, torch.Tensor) else v for v in x)
+
+    return x
+
+
+class MLFlowModelCheckpoint(ModelCheckpoint):
+    def __init__(self, *args,
+                 register_model_name: str | None = None,
+                 register_model_on: Literal["train", "val", "test", "predict"] | None = None,
+                 code_paths: list[str] | None = None,
+                 log_model_at_end_only: bool = True,
+                 additional_metadata: dict[str, Any] | None = None,
+                 extra_pip_requirements: list[str] | None = None,
+                 **kwargs):
+        """
+        MLFlowModelCheckpoint is a custom callback for PyTorch Lightning that integrates with MLFlow to log and register models.
+
+        Args:
+            register_model_name (str | None): The name to register the model under in MLFlow. If None, the model will not be registered.
+            register_model_on (Literal["train", "val", "test", "predict"] | None): The stage at which to register the model. If None, the model will not be registered.
+            code_paths (list[str] | None): List of paths to Python files that should be included in the MLFlow model.
+            log_model_at_end_only (bool): If True, only log the model to MLFlow at the end of the training instead of after every checkpoint save.
+            additional_metadata (dict[str, Any] | None): Additional metadata to log with the model as a JSON file.
+            extra_pip_requirements (list[str] | None): Additional pip requirements to include with the MLFlow model.
+            **kwargs: Keyword arguments for ModelCheckpoint.
+        """
+        # Ensure MLflow is configured when callback is initialized
+        ensure_mlflow_configured()
+
+        super().__init__(*args, **kwargs)
+        if self.save_top_k > 1:
+            raise NotImplementedError("save_top_k > 1 is not supported. "
+                                      "Please use save_top_k=1 to save only the best model.")
+        if self.save_last is not None and self.save_top_k != 0 and self.monitor is not None:
+            raise NotImplementedError("save_last is not supported with monitor and save_top_k!=0. "
+                                      "Please use two separate callbacks: one for saving the last model "
+                                      "and another for saving the best model based on the monitor metric.")
+
+        if register_model_name is not None and register_model_on is None:
+            raise ValueError("If you provide a register_model_name, you must also provide a register_model_on.")
+        if register_model_on is not None and register_model_name is None:
+            raise ValueError("If you provide a register_model_on, you must also provide a register_model_name.")
+        if register_model_on not in ["train", "val", "test", "predict", None]:
+            raise ValueError("register_model_on must be one of train, val, test or predict.")
+
+        self.register_model_name = register_model_name
+        self.register_model_on = register_model_on
+        self.log_model_at_end_only = log_model_at_end_only
+        self._last_model_uri = None
+        self.last_saved_model_info = None
+        self._inferred_signature = None
+        self._input_example = None
+        self.code_paths = code_paths
+        self.additional_metadata = additional_metadata or {}
+        self.extra_pip_requirements = extra_pip_requirements or []
+
+    def _infer_params(self, model: nn.Module) -> tuple[dict, ...]:
+        """Extract metadata from the model's forward method signature.
+
+        Returns:
+            A tuple of dicts, each containing parameter metadata ordered by position.
+        """
+        forward_method = getattr(model.__class__, 'forward', None)
+
+        if forward_method is None:
+            return ()
+
+        try:
+            sig = inspect.signature(forward_method)
+            params_list = []
+
+            for param_name, param in sig.parameters.items():
+                if param_name == 'self':
+                    continue
+
+                param_info = {
+                    'name': param_name,
+                    'kind': param.kind.name,
+                    'annotation': param.annotation if param.annotation != inspect.Parameter.empty else None,
+                    'default': param.default if param.default != inspect.Parameter.empty else None,
+                }
+                params_list.append(param_info)
+
+            # Add return annotation if available as the last element
+            return_annotation = sig.return_annotation
+            if return_annotation != inspect.Signature.empty:
+                return_info = {'_return_annotation': str(return_annotation)}
+                params_list.append(return_info)
+
+            return tuple(params_list)
+
+        except Exception as e:
+            _LOGGER.warning(f"Failed to infer forward method parameters: {e}")
+            return ()
+
+    def _save_checkpoint(self, trainer: L.Trainer, filepath: str) -> None:
+        _LOGGER.debug(f"Saving checkpoint to {filepath}...")
+        trainer.save_checkpoint(filepath, self.save_weights_only)
+
+        self._last_global_step_saved = trainer.global_step
+        self._last_checkpoint_saved = filepath
+
+        # notify loggers
+        if trainer.is_global_zero:
+            for logger in trainer.loggers:
+                logger.after_save_checkpoint(proxy(self))
+                if isinstance(logger, MLFlowLogger) and not self.log_model_at_end_only:
+                    _LOGGER.debug(f"_save_checkpoint: Logging model to MLFlow at {filepath}...")
+                    self.log_model_to_mlflow(trainer.model, run_id=logger.run_id)
+
+    def log_additional_metadata(self, logger: MLFlowLogger | L.Trainer,
+                                additional_metadata: dict) -> None:
+        """Log additional metadata as a JSON file to the model artifact.
+
+        Args:
+            logger: The MLFlowLogger or Lightning Trainer instance to use for logging.
+            additional_metadata: A dictionary containing additional metadata to log.
+        """
+        self.additional_metadata = additional_metadata
+        if not self.additional_metadata:
+            return
+
+        if self.last_saved_model_info is None:
+            _LOGGER.warning("No model has been saved yet. Cannot log additional metadata.")
+            return
+
+        try:
+            log_model_metadata(metadata=self.additional_metadata,
+                               logger=logger,
+                               model_path=self.last_saved_model_info.artifact_path)
+        except Exception as e:
+            _LOGGER.warning(f"Failed to log additional metadata: {e}")
+
+    def log_model_to_mlflow(self,
+                            model: nn.Module,
+                            run_id: str | MLFlowLogger
+                            ) -> None:
+        """Log the model to MLflow."""
+        if isinstance(run_id, MLFlowLogger):
+            logger = run_id
+            if logger.run_id is None:
+                raise ValueError("MLFlowLogger has no run_id. Cannot log model to MLFlow.")
+            run_id = logger.run_id
+
+        if self._last_checkpoint_saved is None or self._last_checkpoint_saved == '':
+            _LOGGER.warning("No checkpoint saved yet. Cannot log model to MLFlow.")
+            return
+
+        orig_device = next(model.parameters()).device
+        model = model.cpu()  # Ensure the model is on CPU for logging
+
+        requirements = list(self.extra_pip_requirements)
+        # check if lightning is in the requirements
+        if not any('lightning' in req.lower() for req in requirements):
+            requirements.append(f'lightning=={L.__version__}')
+
+        _LOGGER.debug(f"log_model_to_mlflow: Logging model to MLFlow at {self._last_checkpoint_saved}...")
+        modelinfo = mlflow.pytorch.log_model(
+            pytorch_model=model,
+            artifact_path=f'model/{Path(self._last_checkpoint_saved).stem}',
+            signature=self._inferred_signature,
+            run_id=run_id,
+            extra_pip_requirements=requirements,
+            code_paths=self.code_paths
+        )
+
+        model.to(device=orig_device)  # Move the model back to its original device
+        self._last_model_uri = modelinfo.model_uri
+        self.last_saved_model_info = modelinfo
+
+        # Log additional metadata after the model is saved
+        log_model_metadata(self.additional_metadata,
+                           model_path=modelinfo.artifact_path,
+                           run_id=run_id)
+
+    def _remove_checkpoint(self, trainer: L.Trainer, filepath: str) -> None:
+        super()._remove_checkpoint(trainer, filepath)
+        # remove the checkpoint from mlflow
+        if trainer.is_global_zero:
+            for logger in trainer.loggers:
+                if isinstance(logger, MLFlowLogger):
+                    artifact_uri = logger.experiment.get_run(logger.run_id).info.artifact_uri
+                    rep = get_artifact_repository(artifact_uri)
+                    rep.delete_artifacts(f'model/{Path(filepath).stem}')
+
+    def register_model(self, trainer=None):
+        """Register the model in MLFlow Model Registry."""
+        # mlflow_client = _get_MLFlowLogger(trainer)._mlflow_client
+        return mlflow.register_model(
+            model_uri=self._last_model_uri,
+            name=self.register_model_name,
+        )
+
+    def _update_signature(self, trainer):
+        if self._inferred_signature is None:
+            _LOGGER.warning("No signature found. Cannot update signature.")
+            return
+        if self._last_model_uri is None:
+            _LOGGER.warning("No model URI found. Cannot update signature.")
+            return
+
+        mllogger = _get_MLFlowLogger(trainer)
+        mlclient = mllogger._mlflow_client
+
+        # check if the model exists
+        for artifact_info in mlclient.list_artifacts(run_id=mllogger.run_id):
+            if artifact_info.path.startswith('model'):
+                break
+        else:
+            _LOGGER.warning(f"Model URI {self._last_model_uri} does not exist. Cannot update signature.")
+            return
+        _LOGGER.debug(f"Updating signature for model URI: {self._last_model_uri}...")
+        # update the signature
+        mlflow.models.set_signature(
+            model_uri=self._last_model_uri,
+            signature=self._inferred_signature,
+        )
+
+    def __wrap_forward(self, pl_module: nn.Module):
+        original_forward = pl_module.forward
+
+        def wrapped_forward(x, *args, **kwargs):
+            x0 = help_infer_signature(x)
+            infered_params = self._infer_params(pl_module)
+            if len(infered_params) > 1:
+                infered_params = {param['name']: param['default']
+                                  for param in infered_params[1:] if 'name' in param}
+            else:
+                infered_params = None
+
+            self._inferred_signature = mlflow.models.infer_signature(model_input=x0,
+                                                                     params=infered_params)
+
+
+            # run once and get back to the original forward
+            pl_module.forward = original_forward
+            method = getattr(pl_module, 'forward')
+            out = method(x, *args, **kwargs)
+
+            output_sig = mlflow.models.infer_signature(model_output=help_infer_signature(out))
+            self._inferred_signature.outputs = output_sig.outputs
+
+            return out
+
+        pl_module.forward = wrapped_forward
+
+    def on_train_start(self, trainer, pl_module):
+        self.__wrap_forward(pl_module)
+
+    def on_train_end(self, trainer: L.Trainer, pl_module: L.LightningModule) -> None:
+        super().on_train_end(trainer, pl_module)
+
+        if self.log_model_at_end_only and trainer.is_global_zero:
+            logger = _get_MLFlowLogger(trainer)
+            if logger is None:
+                _LOGGER.warning("No MLFlowLogger found. Cannot log model to MLFlow.")
+            else:
+                self.log_model_to_mlflow(trainer.model, run_id=logger.run_id)
+
+        self._update_signature(trainer)
+
+        if self.register_model_on == 'train':
+            self.register_model(trainer)
+
+    def _restore_model_uri(self, trainer: L.Trainer) -> None:
+        logger = _get_MLFlowLogger(trainer)
+        if logger is None:
+            _LOGGER.warning("No MLFlowLogger found. Cannot restore model URI.")
+            return
+        if trainer.ckpt_path is None:
+            return
+        extracted_run_id = Path(trainer.ckpt_path).parts[1]
+        if extracted_run_id != logger.run_id:
+            _LOGGER.warning(f"Run ID mismatch: {extracted_run_id} != {logger.run_id}." +
+                            " Check `run_id` parameter in MLFlowLogger.")
+        self._last_model_uri = f'runs:/{logger.run_id}/model/{Path(trainer.ckpt_path).stem}'
+        try:
+            self.last_saved_model_info = mlflow.models.get_model_info(self._last_model_uri)
+        except mlflow.exceptions.MlflowException as e:
+            _LOGGER.warning(f"Failed to get model info for URI {self._last_model_uri}: {e}")
+            self.last_saved_model_info = None
+
+    def on_test_start(self, trainer, pl_module):
+        self.__wrap_forward(pl_module)
+        self._restore_model_uri(trainer)
+        return super().on_test_start(trainer, pl_module)
+
+    def on_predict_start(self, trainer, pl_module):
+        self.__wrap_forward(pl_module)
+        self._restore_model_uri(trainer)
+        return super().on_predict_start(trainer, pl_module)
+
+    def on_test_end(self, trainer: L.Trainer, pl_module: L.LightningModule) -> None:
+        super().on_test_end(trainer, pl_module)
+
+        if self.register_model_on == 'test':
+            self._update_signature(trainer)
+            self.register_model(trainer)
+
+    def on_predict_end(self, trainer: L.Trainer, pl_module: L.LightningModule) -> None:
+        super().on_predict_end(trainer, pl_module)
+
+        if self.register_model_on == 'predict':
+            self._update_signature(trainer)
+            self.register_model(trainer)
+
+    def on_validation_end(self, trainer: L.Trainer, pl_module: L.LightningModule) -> None:
+        super().on_validation_end(trainer, pl_module)
+
+        if self.register_model_on == 'val':
+            self._update_signature(trainer)
+            self.register_model(trainer)

datamint/mlflow/models/__init__.py

@@ -0,0 +1,94 @@
+import logging
+import json
+import lightning as L
+from lightning.pytorch.loggers import MLFlowLogger
+import mlflow
+import os
+from tempfile import TemporaryDirectory
+from torch import nn
+
+_LOGGER = logging.getLogger(__name__)
+
+
+def download_model_metadata(model_uri: str) -> dict:
+    from mlflow.tracking.artifact_utils import get_artifact_repository
+
+    art_repo = get_artifact_repository(artifact_uri=model_uri)
+    try:
+        out_artifact_path = art_repo.download_artifacts(artifact_path='metadata.json')
+    except OSError as e:
+        _LOGGER.warning(f"Error downloading model metadata: {e}")
+        return {}
+
+    with open(out_artifact_path, 'r') as f:
+        metadata = json.load(f)
+    return metadata
+
+
+def _get_MLFlowLogger(trainer: L.Trainer) -> MLFlowLogger:
+    for logger in trainer.loggers:
+        if isinstance(logger, MLFlowLogger):
+            return logger
+    raise ValueError("No MLFlowLogger found in the trainer loggers.")
+
+
+def log_model_metadata(metadata: dict,
+                       mlflow_model: mlflow.models.Model | None = None,
+                       logger: MLFlowLogger | L.Trainer | None = None,
+                       model_path: str | None = None,
+                       run_id: str | None = None,
+                       ) -> None:
+    """
+    Log additional metadata to the MLflow model.
+    It should be provided the one of the following combination of parameters:
+    1. `mlflow_model`
+    2. `logger` and `model_path`
+    3. `run_id` and `model_path`
+
+    Args:
+        self: The instance of the class where this method is called.
+        metadata (dict): The metadata to log.
+        mlflow_model (mlflow.models.Model, optional): The MLflow model object. Defaults to None.
+        logger (MLFlowLogger or L.Trainer, optional): The MLFlow logger or Lightning Trainer instance. Defaults to None.
+        model_path (str, optional): The path where the model is stored in MLflow. Defaults to None.
+        run_id (str, optional): The run ID of the MLflow run. Defaults to None.
+    """
+
+    # Validate inputs
+    if mlflow_model is None and (logger is None or model_path is None) and (run_id is None or model_path is None):
+        raise ValueError(
+            "You must provide either `mlflow_model`, or both `logger` and `model_path`, "
+            "or both `run_id` and `model_path`."
+        )
+    # not both
+    if mlflow_model is not None and logger is not None:
+        raise ValueError("Only one of mlflow_model or logger can be provided.")
+
+    if logger is not None and isinstance(logger, L.Trainer):
+        logger = _get_MLFlowLogger(logger)
+        if logger is None:
+            raise ValueError("MLFlowLogger not found in the Trainer's loggers.")
+        run_id = logger.run_id
+        artifact_path = model_path
+        mlfclient = logger.experiment
+    elif mlflow_model is not None:
+        run_id = mlflow_model.run_id
+        artifact_path = mlflow_model.artifact_path
+        mlfclient = mlflow.client.MlflowClient()
+    elif run_id is not None and model_path is not None:
+        mlfclient = mlflow.client.MlflowClient()
+        artifact_path = model_path
+    else:
+        raise ValueError("Invalid logger or mlflow_model provided.")
+
+    with TemporaryDirectory() as tmpdir:
+        metadata_path = os.path.join(tmpdir, "metadata.json")
+        with open(metadata_path, "w") as f:
+            json.dump(metadata, f, indent=2)
+
+        mlfclient.log_artifact(
+            run_id=run_id,
+            local_path=metadata_path,
+            artifact_path=artifact_path,
+        )
+        _LOGGER.debug(f"Additional metadata logged to {artifact_path}/metadata.json")

datamint/mlflow/tracking/datamint_store.py

@@ -0,0 +1,46 @@
+from mlflow.store.tracking.rest_store import RestStore
+from functools import partial
+from .fluent import get_active_project_id
+import json
+
+
+class DatamintStore(RestStore):
+    """
+    DatamintStore is a subclass of RestStore that provides a tracking store
+    implementation for Datamint.
+    """
+
+    def __init__(self, store_uri: str, artifact_uri=None, force_valid=True):
+        # Ensure MLflow environment is configured when store is initialized
+        from datamint.mlflow.env_utils import setup_mlflow_environment
+        from mlflow.utils.credentials import get_default_host_creds
+        setup_mlflow_environment()
+        
+        if store_uri.startswith('datamint://') or 'datamint.io' in store_uri or force_valid:
+            self.invalid = False
+        else:
+            self.invalid = True
+
+        store_uri = store_uri.split('datamint://', maxsplit=1)[-1]
+        get_host_creds = partial(get_default_host_creds, store_uri)
+        super().__init__(get_host_creds=get_host_creds)
+
+    def create_experiment(self, name, artifact_location=None, tags=None, project_id: str = None) -> str:
+        from mlflow.protos.service_pb2 import CreateExperiment
+        from mlflow.utils.proto_json_utils import message_to_json
+
+        if self.invalid:
+            return super().create_experiment(name, artifact_location, tags)
+        if project_id is None:
+            project_id = get_active_project_id()
+        tag_protos = [tag.to_proto() for tag in tags] if tags else []
+        req_body = message_to_json(
+            CreateExperiment(name=name, artifact_location=artifact_location, tags=tag_protos)
+        )
+
+        req_body = json.loads(req_body)
+        req_body["project_id"] = project_id  # FIXME: this should be in the proto
+        req_body = json.dumps(req_body)
+
+        response_proto = self._call_endpoint(CreateExperiment, req_body)
+        return response_proto.experiment_id

datamint/mlflow/tracking/default_experiment.py

@@ -0,0 +1,27 @@
+import sys
+import os
+from mlflow.tracking.default_experiment.abstract_context import DefaultExperimentProvider
+
+
+class DatamintExperimentProvider(DefaultExperimentProvider):
+    _experiment_id = None
+
+    def in_context(self):
+        return True
+
+    def get_experiment_id(self):
+        from mlflow.tracking.client import MlflowClient
+        
+        if DatamintExperimentProvider._experiment_id is not None:
+            return self._experiment_id
+        # Get the filename of the main source file
+        source_code_filename = os.path.basename(sys.argv[0])
+        mlflowclient = MlflowClient()
+        exp = mlflowclient.get_experiment_by_name(source_code_filename)
+        if exp is None:
+            experiment_id = mlflowclient.create_experiment(source_code_filename)
+        else:
+            experiment_id = exp.experiment_id
+        DatamintExperimentProvider._experiment_id = experiment_id
+
+        return experiment_id

datamint/mlflow/tracking/fluent.py

@@ -0,0 +1,78 @@
+from typing import Optional
+import threading
+import logging
+from datamint import Api
+from datamint.exceptions import DatamintException
+import os
+from datamint.mlflow.env_vars import EnvVars
+from datamint.mlflow.env_utils import ensure_mlflow_configured
+
+_PROJECT_LOCK = threading.Lock()
+_LOGGER = logging.getLogger(__name__)
+
+_ACTIVE_PROJECT_ID: Optional[str] = None
+
+
+def get_active_project_id() -> str | None:
+    """
+    Get the active project ID from the environment variable or the global variable.
+    """
+    global _ACTIVE_PROJECT_ID
+
+    if _ACTIVE_PROJECT_ID is not None:
+        return _ACTIVE_PROJECT_ID
+    # Check if the environment variable is set
+    project_id = os.getenv(EnvVars.DATAMINT_PROJECT_ID.value)
+    if project_id is not None:
+        _ACTIVE_PROJECT_ID = project_id
+        return project_id
+    project_name = os.getenv(EnvVars.DATAMINT_PROJECT_NAME.value)
+    if project_name is not None:
+        project = _find_project_by_name(project_name)
+        if project is not None:
+            _ACTIVE_PROJECT_ID = project['id']
+            return _ACTIVE_PROJECT_ID
+
+    return None
+
+
+def _find_project_by_name(project_name: str):
+    dt_client = Api(check_connection=False)
+    project = dt_client.projects.get_by_name(project_name)
+    if project is None:
+        raise DatamintException(f"Project with name '{project_name}' does not exist.")
+    return project
+
+
+def set_project(project_name: Optional[str] = None, project_id: Optional[str] = None):
+    from mlflow.exceptions import MlflowException
+    global _ACTIVE_PROJECT_ID
+
+    # Ensure MLflow is properly configured before proceeding
+    ensure_mlflow_configured()
+
+    if project_name is None and project_id is None:
+        raise MlflowException("You must specify either a project name or a project id")
+
+    if project_name is not None and project_id is not None:
+        raise MlflowException("You cannot specify both a project name and a project id")
+
+    with _PROJECT_LOCK:
+        dt_client = Api(check_connection=False)
+        if project_id is None:
+            project = dt_client.projects.get_by_name(project_name)
+            if project is None:
+                raise DatamintException(f"Project with name '{project_name}' does not exist.")
+            project_id = project.id
+        else:
+            project = dt_client.projects.get_by_id(project_id)
+            if project is None:
+                raise DatamintException(f"Project with id '{project_id}' does not exist.")
+
+        _ACTIVE_PROJECT_ID = project_id
+
+    # Set 'DATAMINT_PROJECT_ID' environment variable
+    # so that subprocess can inherit it.
+    os.environ[EnvVars.DATAMINT_PROJECT_ID.value] = project_id
+
+    return project