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

datamint/mlflow/env_utils.py

@@ -0,0 +1,131 @@
+"""
+Utility functions for automatically configuring MLflow environment variables
+based on Datamint configuration.
+"""
+
+import os
+import logging
+from urllib.parse import urlparse
+from datamint import configs
+import sys
+
+
+_LOGGER = logging.getLogger(__name__)
+
+
+def get_datamint_api_url() -> str | None:
+    """Get the Datamint API URL from configuration or environment variables."""
+    api_url = configs.get_value(configs.APIURL_KEY, include_envvars=True)  # configs checks env vars first
+    return api_url
+
+
+def get_datamint_api_key() -> str | None:
+    """Get the Datamint API key from configuration or environment variables."""
+    # First check environment variable
+    api_key = os.getenv('DATAMINT_API_KEY')
+    if api_key:
+        return api_key
+
+    # Then check configuration
+    api_key = configs.get_value(configs.APIKEY_KEY)
+    if api_key:
+        return api_key
+
+    return None
+
+
+def _get_mlflowdatamint_uri() -> str | None:
+    api_url = get_datamint_api_url()
+    if not api_url:
+        return None
+    _LOGGER.debug(f"Retrieved Datamint API URL: {api_url}")
+
+    # Remove trailing slash if present
+    api_url = api_url.rstrip('/')
+    # api_url samples:
+    # https://api.datamint.io
+    # http://localhost:3001
+
+    parsed_url = urlparse(api_url)
+    base_url = f"{parsed_url.scheme}://{parsed_url.hostname}"
+    _LOGGER.debug(f"Derived base URL for MLflow Datamint: {base_url}")
+    # FIXME: It should work with https or datamint-api server should forward https requests.
+    base_url = base_url.replace('https://', 'http://')
+    if len(base_url.replace('http:', '')) == 0:
+        return None
+
+    mlflow_uri = f"{base_url}:5000"
+    return mlflow_uri
+
+
+def setup_mlflow_environment(overwrite: bool = False,
+                             set_mlflow: bool = True) -> bool:
+    """
+    Set up MLflow environment variables based on Datamint configuration.
+
+    Args:
+        overwrite (bool): If True, overwrite existing MLflow environment variables.
+        set_mlflow (bool): If True, set the MLflow tracking URI using mlflow.set_tracking_uri().
+
+    Returns:
+        bool: True if success, False otherwise.
+    """
+    api_key = get_datamint_api_key()
+    mlflow_uri = _get_mlflowdatamint_uri()
+    _LOGGER.debug(f"Setting up MLflow environment variables from Datamint configuration: URI='{mlflow_uri}', API_KEY={'***' if api_key is not None else None}")
+    if not mlflow_uri or not api_key:
+        _LOGGER.warning("Datamint configuration incomplete, cannot auto-configure MLflow")
+        return False
+
+    if overwrite or not os.getenv('MLFLOW_TRACKING_TOKEN'):
+        os.environ['MLFLOW_TRACKING_TOKEN'] = api_key
+    if overwrite or not os.getenv('MLFLOW_TRACKING_URI'):
+        os.environ['MLFLOW_TRACKING_URI'] = mlflow_uri
+
+    _LOGGER.debug(f'Final MLflow environment variables: MLFLOW_TRACKING_URI={os.getenv("MLFLOW_TRACKING_URI")}, MLFLOW_TRACKING_TOKEN={"***" if os.getenv("MLFLOW_TRACKING_TOKEN") is not None else None}')
+
+    if set_mlflow:
+        import mlflow
+        _LOGGER.debug(f"Setting MLflow tracking URI to: {mlflow_uri}")
+        mlflow.set_tracking_uri(mlflow_uri)
+
+    if 'lightning.pytorch.loggers' in sys.modules:
+        # import lightning.pytorch.loggers
+        # importlib.reload(lightning.pytorch.loggers)
+        from lightning.pytorch.loggers import MLFlowLogger 
+
+        # 1. Convert the immutable defaults tuple to a mutable list
+        current_defaults = list(MLFlowLogger.__init__.__defaults__)
+
+        # 2. Update the default value for 'tracking_uri'
+        # Based on the signature, 'tracking_uri' is the 3rd argument with a default (index 2)
+        # Signature: (experiment_name, run_name, tracking_uri, ...)
+        current_defaults[2] = mlflow_uri
+
+        # 3. Apply the modified defaults back to the class
+        MLFlowLogger.__init__.__defaults__ = tuple(current_defaults)
+
+
+    return True
+
+
+def ensure_mlflow_configured() -> None:
+    """
+    Ensure MLflow environment is properly configured.
+    Raises ValueError if configuration is incomplete.
+    """
+    if not setup_mlflow_environment():
+        if not os.getenv('MLFLOW_TRACKING_URI'):
+            raise ValueError(
+                "MLflow environment not configured. Please either:\n"
+                "1. Run 'datamint-config --default-url <url>',  or\n"
+                "2. Set DATAMINT_API_URL environment variable, or\n"
+                "3. Manually set MLFLOW_TRACKING_URI environment variable"
+            )
+        if not os.getenv('MLFLOW_TRACKING_TOKEN'):
+            raise ValueError(
+                "MLflow environment not configured. Please either:\n"
+                "1. Run 'datamint-config', or\n"
+                "2. Set DATAMINT_API_KEY environment variable, or\n"
+                "3. Manually set MLFLOW_TRACKING_TOKEN environment variable"
+            )

datamint/mlflow/env_vars.py

@@ -0,0 +1,5 @@
+from enum import Enum
+
+class EnvVars(Enum):
+    DATAMINT_PROJECT_ID = "DATAMINT_PROJECT_ID"
+    DATAMINT_PROJECT_NAME = "DATAMINT_PROJECT_NAME"

datamint/mlflow/flavors/__init__.py

@@ -0,0 +1,17 @@
+"""
+Datamint MLflow custom flavor for wrapping PyTorch models with preprocessing.
+"""
+
+from .datamint_flavor import (
+    save_model,
+    log_model,
+    load_model,
+    _load_pyfunc,
+)
+
+__all__ = [
+    "save_model",
+    "log_model", 
+    "load_model",
+    "_load_pyfunc",
+]

datamint/mlflow/flavors/datamint_flavor.py

@@ -0,0 +1,150 @@
+import mlflow
+from mlflow.models import Model, ModelInputExample, ModelSignature
+import datamint
+import datamint.mlflow.flavors
+from mlflow import pyfunc
+from .model import DatamintModel
+import logging
+from collections.abc import Sequence
+from dataclasses import asdict
+from packaging.requirements import Requirement
+
+FLAVOR_NAME = 'datamint'
+
+_LOGGER = logging.getLogger(__name__)
+
+
+def save_model(datamint_model: DatamintModel,
+               path,
+               supported_modes: Sequence[str] | None = None,
+               data_path=None,
+               code_paths=None,
+               infer_code_paths=False,
+               conda_env=None,
+               mlflow_model: Model | None = None,
+               artifacts=None,
+               signature: ModelSignature | None = None,
+               input_example: ModelInputExample | None = None,
+               pip_requirements=None,
+               extra_pip_requirements=None,
+               metadata=None,
+               model_config=None,
+               streamable=None,
+               **kwargs):
+    import medimgkit
+
+    if mlflow_model is None:
+        mlflow_model = Model()
+
+    mlflow_model.add_flavor(
+        FLAVOR_NAME,
+        datamint_version=datamint.__version__,
+        supported_modes=supported_modes or datamint_model.get_supported_modes(),
+        model_settings=asdict(datamint_model.settings),
+        linked_models=datamint_model._get_linked_models_uri()
+    )
+
+    model_config = model_config or {}
+    model_config.setdefault('device', 'cuda' if datamint_model.settings.need_gpu else 'cpu')
+    
+    def _get_req_name(req):
+        try:
+            return Requirement(req).name.lower()
+        except Exception:
+            return req.split("==")[0].strip().lower()
+
+    datamint_requirements = ['datamint=={}'.format(datamint.__version__), 'medimgkit=={}'.format(medimgkit.__version__)]
+    
+    user_requirements = []
+    # Check if requirements are lists (not strings which are also Sequences)
+    if pip_requirements and isinstance(pip_requirements, Sequence) and not isinstance(pip_requirements, str):
+        user_requirements.extend(pip_requirements)
+    if extra_pip_requirements and isinstance(extra_pip_requirements, Sequence) and not isinstance(extra_pip_requirements, str):
+        user_requirements.extend(extra_pip_requirements)
+        
+    user_req_names = {_get_req_name(req) for req in user_requirements}
+    
+    missing_requirements = [req for req in datamint_requirements if _get_req_name(req) not in user_req_names]
+
+    if missing_requirements:
+        if extra_pip_requirements is None:
+            extra_pip_requirements = missing_requirements
+        elif isinstance(extra_pip_requirements, Sequence) and not isinstance(extra_pip_requirements, str):
+            extra_pip_requirements = list(extra_pip_requirements) + missing_requirements
+        elif pip_requirements and isinstance(pip_requirements, Sequence) and not isinstance(pip_requirements, str):
+             pip_requirements = list(pip_requirements) + missing_requirements
+
+
+    datamint_model._clear_linked_models_cache()
+
+    return mlflow.pyfunc.save_model(
+        path=path,
+        python_model=datamint_model,
+        data_path=data_path,
+        conda_env=conda_env,
+        mlflow_model=mlflow_model,
+        # loader_module=None,
+        artifacts=artifacts,
+        code_paths=code_paths,
+        infer_code_paths=infer_code_paths,
+        signature=signature,
+        input_example=input_example,
+        pip_requirements=pip_requirements,
+        extra_pip_requirements=extra_pip_requirements,
+        metadata=metadata,
+        model_config=model_config,
+        streamable=streamable,
+        **kwargs
+    )
+
+
+def log_model(
+    datamint_model: DatamintModel,
+    supported_modes: Sequence[str] | None = None,
+    name: str = "datamint_model",
+    data_path=None,
+    code_paths=None,
+    infer_code_paths=False,
+    artifacts=None,
+    registered_model_name: str | None = None,
+    signature: ModelSignature | None = None,
+    input_example: ModelInputExample | None = None,
+    pip_requirements=None,
+    extra_pip_requirements=None,
+    metadata=None,
+    model_config=None,
+    **kwargs
+):
+    return Model.log(
+        datamint_model=datamint_model,
+        supported_modes=supported_modes,
+        name=name,
+        flavor=datamint.mlflow.flavors.datamint_flavor,
+        # loader_module=loader_module,
+        data_path=data_path,
+        code_paths=code_paths,
+        artifacts=artifacts,
+        registered_model_name=registered_model_name,
+        signature=signature,
+        input_example=input_example,
+        pip_requirements=pip_requirements,
+        extra_pip_requirements=extra_pip_requirements,
+        metadata=metadata,
+        model_config=model_config,
+        infer_code_paths=infer_code_paths,
+        **kwargs
+    )
+
+
+def load_model(model_uri: str, device: str | None = None) -> DatamintModel:
+    if device is not None:
+        model_config = {'device': device}
+    else:
+        model_config = None
+    return mlflow.pyfunc.load_model(model_uri=model_uri,
+                                    model_config=model_config
+                                    ).unwrap_python_model()
+
+
+def _load_pyfunc(path: str, model_config=None) -> pyfunc.PyFuncModel:
+    return mlflow.pyfunc.load_model(model_uri=path, model_config=model_config)