datamint 2.3.3__py3-none-any.whl → 2.9.0__py3-none-any.whl

datamint/mlflow/lightning/callbacks/__init__.py

@@ -0,0 +1 @@
+from .modelcheckpoint import MLFlowModelCheckpoint

datamint/mlflow/lightning/callbacks/modelcheckpoint.py

@@ -0,0 +1,410 @@
+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
+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.models
+import mlflow.exceptions
+import mlflow.pytorch
+import logging
+import json
+import hashlib
+from lightning.pytorch.loggers import MLFlowLogger
+
+_LOGGER = logging.getLogger(__name__)
+
+
+def help_infer_signature(x):
+    import torch
+    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"] = 'test',
+                 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"]): The stage at which to register the model. It registers at the end of the specified stage.
+            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 not in ["train", "val", "test", "predict"]:
+            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.registered_model_info = None
+        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 []
+        self._last_registered_state_hash: str = "None"
+        self._has_been_trained: bool = False
+
+    def _compute_registration_state_hash(self) -> str:
+        """Compute a hash representing the current model state for registration comparison.
+
+        Returns:
+            A hash string of the current state, or None if state cannot be computed.
+        """
+        state_dict = {
+            'checkpoint_path': str(self._last_checkpoint_saved),
+            'global_step': self._last_global_step_saved,
+            'signature': str(self._inferred_signature) if self._inferred_signature else None,
+            'model_uri': self._last_model_uri,
+        }
+
+        state_str = json.dumps(state_dict, sort_keys=True)
+        return hashlib.md5(state_str.encode('utf-8')).hexdigest()
+
+    def _should_register_model(self) -> bool:
+        """Determine if the model should be registered.
+
+        Returns:
+            True if the model should be registered, False otherwise.
+        """
+
+        if self._last_model_uri is None:
+            _LOGGER.warning("No model URI available. Cannot register model.")
+            return False
+
+        # If never registered before, register
+        if self._last_registered_state_hash is None:
+            return True
+
+        # If model was retrained, register
+        if self._has_been_trained:
+            return True
+
+        # If state changed (signature, checkpoint, etc.), register
+        current_state_hash = self._compute_registration_state_hash()
+        if current_state_hash != self._last_registered_state_hash:
+            return True
+
+        _LOGGER.info("Model already registered with same configuration. Skipping registration.")
+        return False
+
+    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:
+        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:
+                    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__}')
+
+        modelinfo = mlflow.pytorch.log_model(
+            pytorch_model=model,
+            name=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."""
+        if not self._should_register_model():
+            return self.registered_model_info
+
+        # mlflow_client = _get_MLFlowLogger(trainer)._mlflow_client
+        self.registered_model_info = mlflow.register_model(
+            model_uri=self._last_model_uri,
+            name=self.register_model_name,
+        )
+
+        # Update the registered state hash after successful registration
+        self._last_registered_state_hash = self._compute_registration_state_hash()
+        self._has_been_trained = False  # Reset training flag after registration
+
+        _LOGGER.info(f"Model registered as '{self.register_model_name}' "
+                     f"version {self.registered_model_info.version}")
+
+        return self.registered_model_info
+
+    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
+
+        # update the signature
+        try:
+            mlflow.models.set_signature(
+                model_uri=self._last_model_uri,
+                signature=self._inferred_signature,
+            )
+        except mlflow.exceptions.MlflowException as e:
+            _LOGGER.warning(f"Failed to update model signature. Check if model actually exists. {e}")
+
+    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._has_been_trained = True
+        self.__wrap_forward(pl_module)
+        logger = _get_MLFlowLogger(trainer)
+        if logger._tracking_uri.startswith('file:'):
+            _LOGGER.error("MLFlowLogger tracking URI is a local file path. "
+                          "Model registration will likely fail if using MLflow Model Registry.")
+        if logger.experiment_id is not None:
+            mlflow.set_experiment(experiment_id=logger.experiment_id)
+        super().on_train_start(trainer, 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' and self.register_model_name:
+            self.register_model(trainer)
+
+    def _restore_model_uri(self, trainer: L.Trainer) -> None:
+        """Restore the last model URI from the trainer's checkpoint path.
+        """
+        logger = _get_MLFlowLogger(trainer)
+        self._last_model_uri = None
+        self.last_saved_model_info = None
+        if logger is None:
+            _LOGGER.warning("No MLFlowLogger found. Cannot restore model URI.")
+            return
+        if trainer.ckpt_path is None:
+            return
+        if logger.run_id is None:
+            _LOGGER.warning("MLFlowLogger has no run_id. Cannot restore model URI.")
+            return
+        if logger.run_id not in str(trainer.ckpt_path):
+            _LOGGER.warning(f"Run ID mismatch between checkpoint path and MLFlowLogger." +
+                            " Check `run_id` parameter in MLFlowLogger.")
+            return
+        retrieved_logged_models = mlflow.search_logged_models(
+            filter_string=f"name = '{Path(trainer.ckpt_path).stem[:256]}' AND source_run_id='{logger.run_id[:64]}'",
+            order_by=[{"field_name": "last_updated_timestamp", "ascending": False}],
+            output_format="list"
+        )
+        if not retrieved_logged_models:
+            _LOGGER.warning(f"No logged model found for checkpoint {trainer.ckpt_path}.")
+            return
+        # get the most recent one
+        self._last_model_uri = retrieved_logged_models[0].model_uri
+        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' and self.register_model_name:
+            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' and self.register_model_name:
+            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' and self.register_model_name:
+            self._update_signature(trainer)
+            self.register_model(trainer)

datamint/mlflow/models/__init__.py

@@ -0,0 +1,93 @@
+import logging
+import json
+import lightning as L
+from lightning.pytorch.loggers import MLFlowLogger
+import mlflow
+import os
+from tempfile import TemporaryDirectory
+
+_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,76 @@
+from mlflow.store.tracking.rest_store import RestStore
+from mlflow.exceptions import MlflowException
+from mlflow.utils.proto_json_utils import message_to_json
+from functools import partial
+import json
+from typing_extensions import override
+
+
+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 = None) -> str:
+        from mlflow.protos.service_pb2 import CreateExperiment
+        from datamint.mlflow.tracking.fluent import get_active_project_id
+
+        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
+
+    @override
+    def get_experiment_by_name(self, experiment_name, project_id: str | None = None):
+        from datamint.mlflow.tracking.fluent import get_active_project_id
+        from mlflow.protos.service_pb2 import GetExperimentByName
+        from mlflow.entities import Experiment
+        from mlflow.protos import databricks_pb2
+
+        if self.invalid:
+            return super().get_experiment_by_name(experiment_name)
+        if project_id is None:
+            project_id = get_active_project_id()
+        try:
+            req_body = message_to_json(GetExperimentByName(experiment_name=experiment_name))
+            if project_id:
+                body = json.loads(req_body)
+                body["project_id"] = project_id
+                req_body = json.dumps(body)
+
+            response_proto = self._call_endpoint(GetExperimentByName, req_body)
+            return Experiment.from_proto(response_proto.experiment)
+        except MlflowException as e:
+            if e.error_code == databricks_pb2.ErrorCode.Name(
+                databricks_pb2.RESOURCE_DOES_NOT_EXIST
+            ):
+                return None
+            else:
+                raise

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