Utils

---

source

set_mlflow_tracking_uri

def set_mlflow_tracking_uri(
    uri:NoneType=None
):

Point MLflow at a tracking store for fastMONAI runs (opt-in).

Call once, before mlflow.set_experiment/mlflow.start_run or before creating fastMONAI’s MLflow callbacks, to choose where runs are logged.

Args: uri: A tracking URI (e.g. "sqlite:///runs.db" or an HTTP server). If None, uses fastMONAI’s default sqlite store at the package/repo root.

Returns: The tracking URI that was set.

---

source

store_variables

def store_variables(
    config_fn:str | pathlib.Path, size:list, apply_reorder:bool, target_spacing:int | list
):

Save inference variables as JSON (numpy values coerced to native types).

Written as JSON for safe, pickle-free sharing; load_variables reads JSON and still falls back to legacy pickle files.

---

source

load_variables

def load_variables(
    config_fn:str | pathlib.Path
):

Load stored inference variables.

Tries JSON first. Falls back to legacy pickle ONLY for files named .pkl (so a non-JSON payload disguised as .json is refused rather than unpickled, since pickle can execute arbitrary code). Existing .pkl artifacts still load.

Args: config_fn: File path to load.

Returns: The deserialized [size, apply_reorder, target_spacing] list.

Patch-based inference settings

---

source

decode_model_output

def decode_model_output(
    logits:Tensor, output_spec:Mapping
):

Apply the declared safe output contract to raw model logits.

---

source

standard_inference_config

def standard_inference_config(
    apply_reorder:bool, target_spacing, item_tfms:list, output_spec:Mapping
)->dict:

Build a complete standard/whole-volume inference contract from transform specs.

---

source

patch_inference_config

def patch_inference_config(
    config, output_spec:Mapping
)->dict:

Build the complete versioned inference contract for a patch workflow.

---

source

make_output_spec

def make_output_spec(
    kind:str, threshold:float=None, classes:NoneType=None, output_shape:NoneType=None, class_axis:int=1
)->dict:

Build a versioned, allow-listed inference output/decoding contract.

Supported kind values are binary_segmentation, multiclass_segmentation, classification, and regression.

---

source

patch_config_to_dict

def patch_config_to_dict(
    config, inference_only:bool=False
)->dict:

Return PatchConfig values as a JSON-native dictionary.

config may be a PatchConfig-like object or a mapping. Missing fields in historical mappings receive the current PatchConfig defaults. Set inference_only=True to omit training-only sampling, queue, and data-state settings.

---

source

print_colab_gpu_info

def print_colab_gpu_info():

Check if we have a GPU attached to the runtime.

torch.compile helpers

---

source

load_model_artifact

def load_model_artifact(
    path:str | pathlib.Path, device:str='cpu'
):

Load a Safetensors model artifact.

---

source

load_safetensors_model

def load_safetensors_model(
    path:str | pathlib.Path, device:str='cpu'
):

Strict-load a registered Safetensors model and return it in eval mode.

---

source

save_safetensors_model

def save_safetensors_model(
    model, path:str | pathlib.Path, model_spec:Mapping, inference_config:Mapping, artifact_role:str,
    mlflow_run:str=None, extra_metadata:Mapping=None
)->Path:

Save a rebuildable nn.Module with the model-level Safetensors API.

---

source

read_safetensors_metadata

def read_safetensors_metadata(
    path:str | pathlib.Path
)->dict:

Read and validate parsed fastMONAI Safetensors metadata.

---

source

build_model_from_spec

def build_model_from_spec(
    spec:Mapping
):

Rebuild a model only through registered, trusted constructors and wrappers.

---

source

make_model_spec

def make_model_spec(
    arch_id:str, arch_kwargs:Mapping, wrapper_spec:NoneType=None
)->dict:

Create a validated, versioned model specification.

Wrappers are listed from inner to outer and use {'wrapper_id': str, 'wrapper_kwargs': dict}.

---

source

register_model_wrapper

def register_model_wrapper(
    wrapper_id:str, constructor:Callable, replace:bool=False
):

Register a trusted constructor(model, **kwargs) wrapper.

---

source

register_model_constructor

def register_model_constructor(
    arch_id:str, constructor:Callable, replace:bool=False
):

Register a trusted model constructor for Safetensors reconstruction.

---

source

strip_compile_prefix

def strip_compile_prefix(
    state_dict, prefix:str='_orig_mod.'
):

Remove the torch.compile key prefix from a state dict, in place.

learn.save() on a compiled model writes keys prefixed with _orig_mod., which an uncompiled module cannot load.

Args: state_dict: A state dict, e.g. from torch.load. fastai checkpoints written with with_opt=True are {‘model’: …, ‘opt’: …}; pass the ‘model’ value. prefix: Key prefix to remove.

Returns: The same dict, so it can be passed straight to load_state_dict. A dict without the prefix is returned unchanged.

---

source

unwrap_compiled_model

def unwrap_compiled_model(
    model
):

Return model with any torch.compile wrapper removed.

torch.compile(model) returns an OptimizedModule that holds the original module as _orig_mod alongside TorchDynamo state bound to the torch build that created it. A learner exported while compiled loads on any torch, but raises on the first forward pass under a different one. The returned module shares its weights with the wrapper.

Args: model: An nn.Module, or a fastai Learner whose .model is unwrapped in place.

Returns: The innermost nn.Module, or the same Learner that was passed in. A model that was never compiled is returned unchanged.

Note: Handles torch.compile(model). model.compile() compiles in place and leaves no wrapper to remove.

---

source

ModelTrackingCallback

def ModelTrackingCallback(
    model_name:str, loss_function:str, item_tfms:list, size:list, target_spacing:list, apply_reorder:bool,
    experiment_name:str=None, run_name:str=None, auto_start:bool=False, patch_config:dict=None, model_spec:dict=None,
    output_spec:dict=None, inference_config:dict=None, extra_params:dict=None, extra_tags:dict=None,
    dataset_version:str=None, preprocessing_manifest:str | pathlib.Path=None, sample_id_col:str | int=None,
    log_split:bool=True
):

A FastAI callback for comprehensive MLflow experiment tracking.

This callback automatically logs hyperparameters, metrics, model artifacts, and configuration to MLflow during training. If a checkpoint callback (SaveModelCallback, EMACheckpoint, or any TrackerCallback with fname) is present, the best model checkpoint will also be logged as an artifact.

Supports auto-managed runs when created via create_mlflow_callback().

---

source

create_mlflow_callback

def create_mlflow_callback(
    learn, experiment_name:str=None, run_name:str=None, auto_start:bool=True, model_name:str=None,
    model_spec:dict=None, output_spec:dict=None, inference_config:dict=None, extra_params:dict=None,
    extra_tags:dict=None, dataset_version:str=None, preprocessing_manifest:str | pathlib.Path=None,
    sample_id_col:str | int=None, log_split:bool=True
)->ModelTrackingCallback:

Create MLflow tracking callback with auto-extracted configuration.

This factory function automatically extracts configuration from the Learner, eliminating the need to manually specify parameters like size, transforms, loss function, etc.

Auto-extracts from Learner: - Preprocessing: apply_reorder, target_spacing, size/patch_size - Transforms: item_tfms or pre_patch_tfms - Training: loss_func, model architecture

Args: learn: fastai Learner instance experiment_name: MLflow experiment name. If None, uses model name. run_name: MLflow run name. If None, auto-generates with timestamp. auto_start: If True, auto-starts/stops MLflow run in before_fit/after_fit. model_name: Override the auto-extracted model name (used as the experiment name when experiment_name is None). model_spec: Required allow-listed architecture reconstruction specification. output_spec: Required patch-workflow output/decoding contract. inference_config: Required complete config for standard workflows; includes output. extra_params: Additional parameters to log (e.g., {‘dropout’: 0.5}). extra_tags: MLflow tags to set on the run. dataset_version: Optional annotation/dataset fingerprint for tracking. preprocessing_manifest: Optional manifest returned by preprocess_dataset(). Its dataset/cache identities are validated, tagged, and logged. sample_id_col: Optional stable identity column for split_version. By default, the raw image path is used when available. log_split: If True, hashes and logs actual fastai train/val membership.

Returns: ModelTrackingCallback ready to use with learn.fit()

Example: >>> # Instead of this (6 manual params): >>> # mlflow_callback = ModelTrackingCallback( >>> # model_name=f”{task}_{model._get_name()}“, >>> # loss_function=loss_func.loss_func._get_name(), >>> # item_tfms=item_tfms, >>> # size=size, >>> # target_spacing=target_spacing, >>> # apply_reorder=True, >>> # ) >>> # with mlflow.start_run(run_name=”training”): >>> # learn.fit_one_cycle(30, lr, cbs=[mlflow_callback]) >>> >>> # Patch workflow: architecture and decoding remain explicit. >>> model_spec = make_model_spec(“monai.unet”, unet_kwargs) >>> output_spec = make_output_spec(“multiclass_segmentation”, classes=2) >>> callback = create_mlflow_callback( … learn, experiment_name=“Task02_Heart”, … model_spec=model_spec, output_spec=output_spec) >>> learn.fit_one_cycle(30, lr, cbs=[callback, save_best])

Locating Safetensors model artifacts

---

source

find_model_artifacts

def find_model_artifacts(
    run_ids, artifact_role, expected_members:NoneType=None
):

Download explicitly selected Safetensors model artifacts from MLflow.

run_ids must map each stable member identifier to one immutable MLflow run ID. Experiment-name or latest-run discovery is intentionally unsupported: callers must declare exactly which trained models inference will use.

# Test auto-extraction helper functions
from fastcore.test import test_eq, test_fail
from dataclasses import dataclass

# Test _detect_patch_workflow
class MockStandardDls:
    bs = 4
    after_item = None
mock_std = MockStandardDls()
test_eq(_detect_patch_workflow(mock_std), False)

@dataclass
class MockPatchConfig:
    patch_size: list = None
    patch_overlap: float = 0.5
    samples_per_volume: int = 8
    sampler_type: str = 'uniform'
    label_probabilities: dict = None
    queue_length: int = 300
    aggregation_mode: str = 'hann'
    padding_mode: int = 0
    keep_largest_component: bool = False
    apply_reorder: bool = True
    target_spacing: list = None
    
    def __post_init__(self):
        if self.patch_size is None:
            self.patch_size = [96, 96, 96]

class MockPatchDls:
    bs = 4
    patch_config = MockPatchConfig()
mock_patch = MockPatchDls()
test_eq(_detect_patch_workflow(mock_patch), True)

# Test _extract_size_from_transforms with mock transform
class MockPadOrCrop:
    def __init__(self, target_shape):
        self.target_shape = target_shape

class MockTransform:
    def __init__(self, target_shape):
        self.pad_or_crop = MockPadOrCrop(target_shape)

tfms = [MockTransform([128, 128, 64])]
test_eq(_extract_size_from_transforms(tfms), [128, 128, 64])
test_eq(_extract_size_from_transforms(None), None)
test_eq(_extract_size_from_transforms([]), None)

print("All auto-extraction helper tests passed!")

All auto-extraction helper tests passed!

---

source

MLflowUIManager

def MLflowUIManager():

Launch and manage a local mlflow ui server from a notebook.

The UI’s lifetime is tied to the kernel: it is reaped when the interpreter exits, gracefully (via an atexit handler) or on a hard kill (via Linux PR_SET_PDEATHSIG, see [_die_with_parent](https://fastmonai.no/utils.html#_die_with_parent)), so closing/restarting the notebook never leaves an orphaned server holding the port. Call stop() to shut it down sooner. An externally-started UI is reused, not killed.