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

store_patch_variables


def store_patch_variables(
    config_fn:str | pathlib.Path, patch_size:list, patch_overlap:int | float | list, aggregation_mode:str,
    apply_reorder:bool=False, target_spacing:list=None, sampler_type:str='uniform', label_probabilities:dict=None,
    samples_per_volume:int=8, queue_length:int=300, queue_num_workers:int=4, keep_largest_component:bool=False,
    normalization:list=None
):

Save patch-based training and inference configuration to a JSON file.

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

Args: config_fn: Path to save the config file. patch_size: Size of patches [x, y, z]. patch_overlap: Overlap for inference (int, float 0-1, or list). aggregation_mode: GridAggregator mode (‘crop’, ‘average’, ‘hann’). apply_reorder: Whether to reorder to canonical (RAS+) orientation. target_spacing: Target voxel spacing [x, y, z]. sampler_type: Type of sampler used during training. label_probabilities: Label probabilities for LabelSampler. samples_per_volume: Number of patches extracted per volume during training. queue_length: Maximum number of patches in queue buffer. queue_num_workers: Number of workers for parallel patch extraction. keep_largest_component: If True, keep only the largest connected component in binary segmentation predictions during inference. normalization: List of normalization specs (from PatchConfig.normalization) to persist, so inference can reconstruct the same pre-inference normalization.

Example: >>> store_patch_variables( … ‘patch_settings.json’, … patch_size=[96, 96, 96], … patch_overlap=0.5, … aggregation_mode=‘hann’, … apply_reorder=True, … target_spacing=[1.0, 1.0, 1.0], … samples_per_volume=16, … keep_largest_component=True … )

source

load_patch_variables


def load_patch_variables(
    config_fn:str | pathlib.Path
)->dict:

Load patch-based training and inference configuration.

Tries JSON first. Falls back to legacy pickle ONLY for files named .pkl (a non-JSON payload disguised as .json is refused rather than unpickled, since pickle can execute arbitrary code). Because JSON stringifies dict keys, integer label_probabilities keys are restored to ints after a JSON load.

Args: config_fn: Path to the config file.

Returns: Dictionary with patch configuration (patch_size, patch_overlap, aggregation_mode, apply_reorder, target_spacing, sampler_type, label_probabilities, samples_per_volume, queue_length, queue_num_workers, normalization).

Example: >>> config = load_patch_variables(‘patch_settings.json’) >>> from fastMONAI.vision_patch import PatchConfig >>> patch_config = PatchConfig(**config)

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,
    extra_params:dict=None, extra_tags:dict=None, dataset_version:str=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,
    extra_params:dict=None, extra_tags:dict=None, dataset_version:str=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). extra_params: Additional parameters to log (e.g., {‘dropout’: 0.5}). extra_tags: MLflow tags to set on the run. dataset_version: Dataset version hash string for tracking. log_split: If True, auto-logs train/val split CSV when dls has split_df.

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]) >>> >>> # Do this (zero manual params): >>> callback = create_mlflow_callback(learn, experiment_name=“Task02_Heart”) >>> learn.fit_one_cycle(30, lr, cbs=[callback, save_best])

Locating cross-validation fold checkpoints

source

find_fold_learners


def find_fold_learners(
    experiment_name, artifact_path:str='model/best_learner.pkl', n_folds:NoneType=None, fold_tag:str='fold'
):

Discover one trained learner checkpoint per cross-validation fold from an MLflow experiment.

Reads back what [ModelTrackingCallback](https://fastmonai.no/utils.html#modeltrackingcallback) logs per run: takes the most recent run for each fold in experiment_name and downloads artifact_path from it, returning {fold: local_path} sorted by fold. Folds are identified by each run’s tags.<fold_tag>, falling back to a <fold_tag>_<n> run name. Warns if the selected folds span more than one dataset_version tag (a sign of a partial retrain).

Pair it with the soft-vote ensembling in [patch_inference](https://fastmonai.no/vision_patch.html#patch_inference) – but the returned values are PATHS, so load them first::

from fastai.learner import load_learner
paths = find_fold_learners("vs5f_unet")
learners = [load_learner(p) for p in paths.values()]
preds = patch_inference(learner=learners, config=config, file_paths=cases)

Args: experiment_name: MLflow experiment to search (e.g. "vs5f_unet"). artifact_path: run-relative artifact to download per fold (default "model/best_learner.pkl", what [ModelTrackingCallback](https://fastmonai.no/utils.html#modeltrackingcallback) logs). n_folds: optional expected folds for a missing-fold warning – an int n (labels 1..n) or an explicit iterable of labels (use this for 0-indexed folds). None makes no assumption and just returns what was found. fold_tag: run tag naming the fold; also the <fold_tag>_ run-name prefix fallback.

Returns: {fold_int: local_path} sorted by fold, or {} if the experiment is missing or MLflow is unavailable.

# 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.