MLflow Integration

The datamint.mlflow module provides seamless MLflow integration for experiment tracking, model logging, and model deployment.

MLflow Components

Overview

Key features:

  • Automatic Experiment Tracking: Experiments are automatically associated with Datamint projects

  • Model Registration: Trained models are registered in MLflow and can be deployed to Datamint

  • Dataset Versioning: Datasets are logged as MLflow artifacts with metadata

  • Custom Checkpointing: MLflowModelCheckpoint integrates with Lightning callbacks

  • Flavor Support: Custom Datamint flavors for segmentation and classification models

Automatic Configuration

The MLflow module auto-configures itself on first import:

# This automatically sets up MLflow environment
import datamint.mlflow

# Or explicitly:
from datamint.mlflow import ensure_mlflow_configured
ensure_mlflow_configured()

Environment Setup

Utility functions for automatically configuring MLflow environment variables based on Datamint configuration.

datamint.mlflow.env_utils.ensure_mlflow_configured()

Ensure MLflow environment is properly configured. Raises ValueError if configuration is incomplete.

Return type:

None

datamint.mlflow.env_utils.setup_mlflow_environment(overwrite=False, set_mlflow=True)

Set up MLflow environment variables based on Datamint configuration.

Parameters:

  • 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:

True if success, False otherwise.

Return type:

bool

class datamint.mlflow.env_vars.EnvVars(*values)
DATAMINT_PROJECT_ID = 'DATAMINT_PROJECT_ID'
DATAMINT_PROJECT_NAME = 'DATAMINT_PROJECT_NAME'

MLflow Dataset

MLflow Dataset adapter for Datamint project splits.

class datamint.mlflow.data.datamint_dataset.DatamintMLflowDataset(project_id, project_name, split, resources, extra_params=None)

Bases: Dataset

MLflow Dataset wrapping a Datamint project split for lineage tracking.

Parameters:

  • project_id (str)

  • project_name (str)

  • split (str | None)

  • resources (Sequence[str] | Sequence[Resource])

  • extra_params (dict[str, Any] | None)

property profile: Any | None

Optional summary statistics for the dataset, such as the number of rows in a table, the mean / median / std of each table column, etc.

property schema

Optional dataset schema, such as an instance of mlflow.types.Schema representing the features and targets of the dataset.

to_dict()

Create config dictionary for the dataset.

Subclasses should override this method to provide additional fields in the config dict, e.g., schema, profile, etc.

Returns a string dictionary containing the following fields: name, digest, source, source type.

Return type:

dict[str, str]

The DatamintMLflowDataset wraps a DatamintBaseDataset for MLflow artifact logging.

Model Flavors

class datamint.mlflow.flavors.model.BaseDatamintModel(settings=None)

Bases: PythonModel, ABC

Core prediction gateway that any MLflow PythonModel can build on.

Owns:

Use DatamintModel when you need to load linked models at serve time.

Subclasses only need to implement predict_default() (and optionally other predict_* hooks registered with @prediction_mode).

Parameters:

settings (ModelSettings | dict[str, Any] | None)

annotation_specs: ClassVar[list[AnnotationSpec] | None] = None

Annotation specs this model produces. Subclasses can override at the class body level. When None, specs may be provided via log_model/save_model instead.

get_supported_modes()

Return the list of prediction modes supported by this model.

Return type:

list[str]

property inference_device: str

The device that will be used for inference.

Returns _inference_device if already set, then falls back to the MLFLOW_DEFAULT_PREDICTION_DEVICE environment variable, then 'cpu'.

load_context(context)

Detect the inference device and move any attached torch modules to it.

Override in subclasses to perform additional loading (e.g. linked models) — but always call super().load_context(context) first so that inference_device is set before any model loading.

Parameters:

context (PythonModelContext)

Return type:

None

predict(model_input, params=None)

Main prediction entry point.

Routes to the appropriate handlesr based on params['mode']. Do not override — implement predict_default() (or other predict_* hooks) instead.

Pass params={'log_predictions': True, 'model_name': ...} to also upload the resulting annotations to the Datamint server, tagged source='model_deploy'. This is meant for local inference against a model loaded via datamint.mlflow.flavors.load_model() (testing a registered model outside of Trainer.fit()), as opposed to predictions made automatically during training, which are tagged source='model_pipeline' instead.

If params['prompts'] is a dict, it is converted to a DeployedModelPrompts instance before dispatch, so predict_* handlers receive a typed object rather than a raw dict.

Parameters:

  • model_input (list[BaseResource])

  • params (dict[str, Any] | None)

Return type:

list[list[Annotation]]

predict_all_frames(model_input, **kwargs)

Process all video frames independently through predict_image().

Parameters:

  • model_input (list[BaseResource])

  • kwargs (Any)

Return type:

list[list[Annotation]]

predict_default(model_input, **kwargs)

Default prediction on entire resources.

Override this in your subclass.

Parameters:

  • model_input (list[BaseResource])

  • kwargs (Any)

Return type:

list[list[Annotation]]

predict_frame(model_input, frame_index, **kwargs)

Process a single video frame through predict_image().

Parameters:

  • model_input (list[BaseResource])

  • frame_index (int)

  • kwargs (Any)

Return type:

Sequence[Sequence[ImageSegmentation | ImageClassification]]

predict_image(model_input, **kwargs)
Parameters:

  • model_input (list[BaseResource])

  • kwargs (Any)

Return type:

Sequence[Sequence[ImageSegmentation | ImageClassification]]

predict_slice(model_input, slice_index, axis, **kwargs)

Process a single volume slice through predict_image().

Parameters:

  • model_input (list[BaseResource])

  • slice_index (int)

  • axis (Literal['axial', 'sagittal', 'coronal'])

  • kwargs (Any)

Return type:

Sequence[Sequence[ImageSegmentation | ImageClassification]]

predict_volume(model_input, **kwargs)

Process each volume by iterating axial slices through predict_slice().

Parameters:

  • model_input (list[BaseResource])

  • kwargs (Any)

Return type:

list[list[Annotation]]

property settings: ModelSettings
task_type: ClassVar[TaskType | None] = None

Semantic task category for this model class. Subclasses should override at the class body level (e.g. task_type = TaskType.IMAGE_SEGMENTATION).

to_cpu_for_save()

Move all reachable torch modules to CPU before pickling.

This is a safety net for datamint.mlflow.flavors.datamint_flavor.save_model() — it ensures that any torch.nn.Module held as an attribute (e.g. inside LinkedModelLoader) is relocated to CPU so the resulting pickle is device-agnostic and loads on CPU-only containers.

Safe to call multiple times; idempotent.

Return type:

None

class datamint.mlflow.flavors.model.DatamintModel(settings=None, mlflow_torch_models_uri=None, mlflow_models_uri=None, torch_model=None)

Bases: BaseDatamintModel

Abstract adapter for wrapping ML models to produce Datamint annotations.

Extends BaseDatamintModel with support for loading external (“linked”) MLflow models at serve time via LinkedModelLoader. Subclasses only need to override predict_default (and optionally other predict_* hooks).

Quick Start:

class MyModel(DatamintModel):
    def __init__(self):
        super().__init__(
            mlflow_models_uri={'model': 'models:/MyModel/latest'},
            settings=ModelSettings(need_gpu=True),
        )

    def predict_default(self, model_input, **kwargs):
        device = self.inference_device
        model = self.get_mlflow_models()['model'].get_raw_model().to(device)
        return predictions

You can also pass pre-instantiated torch.nn.Module objects directly:

class MyModel(DatamintModel):
    def __init__(self):
        net = MyTorchNet()
        super().__init__(torch_model=net)

    def predict_default(self, model_input, **kwargs):
        net = self.get_mlflow_torch_models()['net']
        return net(preprocess(model_input))
Parameters:

  • settings (ModelSettings | dict[str, Any] | None)

  • mlflow_torch_models_uri (dict[str, str] | None)

  • mlflow_models_uri (dict[str, str] | None)

  • torch_model (Module | None)

LINKED_MODELS_DIR: ClassVar[str] = 'linked_models'
get_mlflow_models()

Access loaded MLflow pyfunc models.

Return type:

dict[str, PyFuncModel]

get_mlflow_torch_models()

Access loaded MLflow PyTorch models.

Return type:

dict[str, Any]

get_pytorch_model()
Return type:

Module | None

load_context(context)

Detect device and load all linked MLflow models.

Parameters:

context (PythonModelContext)

Return type:

None

property mlflow_models_uri: dict[str, str]
property mlflow_torch_models_uri: dict[str, str]

Task type enumeration for Datamint MLflow models.

class datamint.mlflow.flavors.task_type.TaskType(*values)

Medical-AI task categories for Datamint models.

The str mixin ensures values are JSON-serialisable in MLflow metadata without explicit .value calls.

ANOMALY_DETECTION = 'anomaly_detection'
IMAGE_CLASSIFICATION = 'image_classification'
IMAGE_SEGMENTATION = 'image_segmentation'
INSTANCE_SEGMENTATION = 'instance_segmentation'
LANDMARK_DETECTION = 'landmark_detection'
MULTILABEL_IMAGE_CLASSIFICATION = 'multilabel_image_classification'
OBJECT_DETECTION = 'object_detection'
REPORT_GENERATION = 'report_generation'
VIDEO_FRAME_CLASSIFICATION = 'video_frame_classification'
VIDEO_SEGMENTATION = 'video_segmentation'
VOLUME_CLASSIFICATION = 'volume_classification'
VOLUME_SEGMENTATION = 'volume_segmentation'

Registry-driven prediction dispatcher and @prediction_mode decorator.

Replaces the hardcoded mode_param_keys dict, stub-method pattern, and fragile _is_mode_implemented introspection previously in DatamintModel.

class datamint.mlflow.flavors.prediction_router.DeployedModelPrompts(**data)

Prompts for a promptable model, as received by a predict_* handler.

Parameters:

data (Any)

boxes: list[list[float]] | None
classmethod from_dict(data)

Build from the server’s preprocessed prompt dict.

Parameters:

data (dict)

Return type:

DeployedModelPrompts

model_config: ClassVar[ConfigDict] = {'extra': 'ignore'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

point_coordinates: list[list[float]] | None
point_labels: list[int] | None
text: str | None
class datamint.mlflow.flavors.prediction_router.ModeSpec(mode, param_keys=(), fallback_to_default=True)

Metadata for a registered prediction mode handler.

Parameters:

  • mode (PredictionMode)

  • param_keys (tuple[str, ...])

  • fallback_to_default (bool)

fallback_to_default: bool
mode: PredictionMode
param_keys: tuple[str, ...]
class datamint.mlflow.flavors.prediction_router.PredictionRouter(model_instance, base_class=None)

Registry-driven prediction dispatcher.

Discovers mode handlers in two ways (in order of priority):

  1. Methods decorated with @prediction_mode.

  2. Convention-named methods predict_<mode_value> not defined on the abstract base (backward compatibility with old-style overrides).

Parameters:

  • model_instance (Any)

  • base_class (type | None)

static discover(model, base_class)

Build the mode -> handler registry from the model instance.

Parameters:

  • model (Any)

  • base_class (type | None)

Return type:

dict[PredictionMode, tuple[Callable, ModeSpec]]

dispatch(model_input, params)
Parameters:

  • model_input (list)

  • params (dict[str, Any])

Return type:

list

supported_modes()
Return type:

list[str]

update_registry(model_instance, base_class=None, overwrite=False)

Update the registry with handlers from a new model instance (e.g. linked model).

Parameters:

  • model_instance (Any)

  • base_class (type | None)

  • overwrite (bool)

Return type:

None

datamint.mlflow.flavors.prediction_router.bridge_mode(fn)

Marks a method as a bridge implementation.

Parameters:

fn (Callable)

Return type:

Callable

datamint.mlflow.flavors.prediction_router.prediction_mode(mode, *, param_keys=(), fallback_to_default=True)

Decorator that registers a method as a prediction mode handler.

Usage:

class MyModel(DatamintModel):
    @prediction_mode(PredictionMode.SLICE, param_keys=("slice_index", "axis"))
    def predict_slice(self, model_input, *, slice_index, axis="axial", **kw):
        ...
Parameters:

  • mode (PredictionMode)

  • param_keys (tuple[str, ...])

  • fallback_to_default (bool)

Return type:

Callable

Prediction mode enumeration for DataMint models.

class datamint.mlflow.flavors.prediction_modes.PredictionMode(*values)

Enumeration of supported prediction modes.

Each mode corresponds to a specific method signature in DatamintModel.

ALL_FRAMES = 'all_frames'
DEFAULT = 'default'
FEW_SHOT = 'few_shot'
FRAME = 'frame'
FRAME_RANGE = 'frame_range'
IMAGE = 'image'
INTERACTIVE = 'interactive'
PRIMARY_SLICE = 'primary_slice'
SLICE = 'slice'
SLICE_RANGE = 'slice_range'
TEMPORAL_SEQUENCE = 'temporal_sequence'
VOLUME = 'volume'
datamint.mlflow.flavors.datamint_flavor.load_model(model_uri, device=None)
Parameters:

  • model_uri (str)

  • device (str | None)

Return type:

DatamintModel

datamint.mlflow.flavors.datamint_flavor.log_model(datamint_model, task_type=None, supported_modes=None, annotation_specs=None, name='datamint_model', data_path=None, code_paths=None, infer_code_paths=False, artifacts=None, model_name=None, signature=None, input_example=None, pip_requirements=None, extra_pip_requirements=None, metadata=None, model_config=None, **kwargs)
Parameters:

  • datamint_model (BaseDatamintModel)

  • task_type (TaskType | str | None)

  • supported_modes (Sequence[str] | None)

  • annotation_specs (Sequence[AnnotationSpec] | None)

  • name (str)

  • model_name (str | None)

  • signature (ModelSignature | None)

  • input_example (DataFrame | ndarray | dict | list | csr_matrix | csc_matrix | str | bytes | tuple | None)

datamint.mlflow.flavors.datamint_flavor.save_model(datamint_model, path, task_type=None, supported_modes=None, annotation_specs=None, data_path=None, code_paths=None, infer_code_paths=False, conda_env=None, mlflow_model=None, artifacts=None, signature=None, input_example=None, pip_requirements=None, extra_pip_requirements=None, metadata=None, model_config=None, streamable=None, **kwargs)
Parameters:

  • datamint_model (BaseDatamintModel)

  • task_type (TaskType | str | None)

  • supported_modes (Sequence[str] | None)

  • annotation_specs (Sequence[AnnotationSpec] | None)

  • mlflow_model (Model | None)

  • signature (ModelSignature | None)

  • input_example (DataFrame | ndarray | dict | list | csr_matrix | csc_matrix | str | bytes | tuple | None)

Checkpointing

The MLflowModelCheckpoint extends Lightning’s ModelCheckpoint to automatically log checkpoints as MLflow artifacts.

MLflow Tracking

datamint.mlflow.tracking.fluent.set_project(project)

Set the active project for the current session.

The project’s name also becomes the default MLflow experiment name for any run started without an explicit experiment (see DatamintExperimentProvider).

Parameters:

project (Project | str) – The Project instance or project name/ID to set as active.

class datamint.mlflow.tracking.default_experiment.DatamintExperimentProvider
get_experiment_id()

Provide the MLflow Experiment ID for the current MLflow client context.

Assumes that in_context() is True.

Returns:

The ID of the MLflow Experiment associated with the current context.

in_context()

Determine if the MLflow client is running in a context where this provider can identify an associated MLflow Experiment ID.

Returns:

True if the MLflow client is running in a context where the provider can identify an associated MLflow Experiment ID. False otherwise.

class datamint.mlflow.tracking.datamint_store.DatamintStore(store_uri, artifact_uri=None, force_valid=True)

DatamintStore is a subclass of RestStore that provides a tracking store implementation for Datamint.

Parameters:

store_uri (str)

create_experiment(name, artifact_location=None, tags=None, project_id=None)

Create a new experiment. If an experiment with the given name already exists, throws exception.

Parameters:

  • name – Desired name for an experiment.

  • artifact_location – Location to store run artifacts.

  • tags – A list of mlflow.entities.ExperimentTag instances to set for the experiment.

  • project_id (str | None)

Return type:

str

Returns:

experiment_id for the newly created experiment if successful, else None

get_experiment_by_name(experiment_name, project_id=None)

Fetch the experiment by name from the backend store.

Parameters:

  • experiment_name – Name of experiment

  • project_id (str | None)

Returns:

A single mlflow.entities.Experiment object if it exists.

log_batch(run_id, metrics=(), params=(), tags=())

Log multiple metrics, params, and tags for the specified run

Parameters:

  • run_id – String id for the run

  • metrics – List of mlflow.entities.Metric instances to log

  • params – List of mlflow.entities.Param instances to log

  • tags – List of mlflow.entities.RunTag instances to log

Returns:

None.

log_metric(run_id, metric)

Log a metric for the specified run

Parameters:

  • run_id – String id for the run

  • metric – Metric instance to log

log_param(run_id, param)

Log a param for the specified run

Parameters:

  • run_id – String id for the run

  • param – Param instance to log

sync_offline_logs(run_id)

Send everything buffered locally for run_id to the remote Datamint MLflow backend, then clear the buffer. Returns the number of entries synced.

If the connection is still down, the underlying network exception propagates and the buffer is left untouched, so a later retry can pick up where this left off.

Parameters:

run_id (str)

Return type:

int

Artifact Repository

class datamint.mlflow.artifact.datamint_artifacts_repo.DatamintArtifactsRepository(artifact_uri, tracking_uri=None, registry_uri=None)
Parameters:

  • artifact_uri (str)

  • tracking_uri (str | None)

  • registry_uri (str | None)

classmethod resolve_uri(artifact_uri, tracking_uri)

Usage Example

import lightning as L
from datamint.dataset import ImageDataset
from datamint.lightning import DatamintDataModule, UNetPPTrainer

# Create dataset and datamodule
dataset = ImageDataset(project="Liver Segmentation")
datamodule = DatamintDataModule(
    dataset,
    batch_size=16,
    split={'train': 0.8, 'val': 0.1, 'test': 0.1},
)

# Use trainer with automatic MLflow integration
trainer = UNetPPTrainer(
    project="Liver Segmentation",
    image_size=256,
    batch_size=16,
    max_epochs=50,
    accelerator="gpu",
    register_model=True,  # Auto-register in MLflow
)

results = trainer.fit()

# Model is automatically logged and registered
print(results["test_results"])

import mlflow
mlflow.get_experiment_by_name("Liver Segmentation")  # Verify experiment exists
mlflow.search_runs(experiment_names=["Liver Segmentation"])  # View runs