Nv Segment Ct Finetune

scripts/run_finetune.py

#!/usr/bin/env python3
# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""nv_segment_ct_finetune - auto-configuring VISTA3D continual finetune.

Three presets:
  --smoke   1 iter on bundled spleen_micro fixture (synthetic plumbing).
  --sanity  Real-recipe verification on MSD06 Lung Tumor - mirrors the
            published DFW tutorial config: label mapping [[1, 23]], 5 epochs,
            lr=5e-5, patch [128,128,128], resample 1.5 mm isotropic,
            drop_label_prob=0.0, drop_point_prob=1.0 (automatic segmentation).
            Runs original-spacing evaluate.json before and after finetuning.
            Expected DFW reference scores are pretrained Dice 0.6697,
            finetuned Dice 0.6836, and training-best Dice 0.6905.
  default   user dataset under --dataset-dir, lr=5e-5, 50 epochs.

Add ``--softmax`` to use the upstream fixed-channel, mutually exclusive
segmentation workflow. Its documented defaults are lr=1e-4 and 100 epochs.

The wrapper auto-detects GPU + RAM, picks patch_size and cache_rate, writes
`configs/auto_override.json`, and runs `python -m monai.bundle run` (or
`torchrun --nproc_per_node=N -m monai.bundle run` for multi-GPU) exactly as
NV-Segment-CT's upstream `finetune.md` documents.

Engineering verification only. Output is NOT clinically meaningful.
"""

from __future__ import annotations

import inspect
import json
import os
import random
import re
import shlex
import shutil
import subprocess
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
import venv
from importlib.metadata import PackageNotFoundError
from importlib.metadata import version as package_version
from pathlib import Path
from typing import Optional

import nibabel as nib
import numpy as np
import scipy.ndimage as ndi
import typer

SKILL_DIR = Path(__file__).resolve().parent.parent
BUNDLE_DIR = SKILL_DIR / "bundle"
LABEL_DICT = BUNDLE_DIR / "label_dict.json"
UPSTREAM_CTMR_COMMIT = "cb921f5c58837c0f42a713855d68b32af88e1cdd"
HF_MODEL_REVISION = "afb51518689f71e6abb367ee6301b2cd0225c66a"
LABEL_DICT_URL = (
    f"https://raw.githubusercontent.com/NVIDIA-Medtech/NV-Segment-CTMR/{UPSTREAM_CTMR_COMMIT}/"
    "NV-Segment-CT/configs/label_dict.json"
)
SMOKE_FIXTURE = SKILL_DIR / "fixtures" / "spleen_micro"
# Resolve Medical AI Skills cache root from the script's own location: repo_root/.workbench_data.
# Callers can still override with --dataset-dir when their cache lives elsewhere.
_REPO_ROOT = SKILL_DIR.parent.parent
SANITY_DATASET = _REPO_ROOT / ".workbench_data" / "datasets" / "Task06_Lung"
SANITY_ANATOMY = "lung tumor"  # MSD06 label 1 (cancer) -> vista3d global index 23
VERSION = "0.6.0"
SUPPORTED_MONAI_MAJOR_MINOR = {(1, 4)}
AUTO_VENV_DIR = (
    Path.home() / ".cache" / "nvidia-skills" / "venvs" / "nv-segment-ct-finetune-monai14"
)
_CHILD_ENV_KEYS = (
    "CUDA_DEVICE_ORDER",
    "CUDA_HOME",
    "CUDA_PATH",
    "CUDA_VISIBLE_DEVICES",
    "CUBLAS_WORKSPACE_CONFIG",
    "CURL_CA_BUNDLE",
    "HF_HOME",
    "HF_HUB_CACHE",
    "HF_HUB_OFFLINE",
    "HOME",
    "HUGGINGFACE_HUB_CACHE",
    "LANG",
    "LC_ALL",
    "LD_LIBRARY_PATH",
    "MONAI_DATA_DIRECTORY",
    "NCCL_DEBUG",
    "NCCL_IB_DISABLE",
    "NCCL_P2P_DISABLE",
    "NCCL_SOCKET_IFNAME",
    "NVIDIA_DRIVER_CAPABILITIES",
    "NVIDIA_VISIBLE_DEVICES",
    "NPROC_PER_NODE",
    "NVSEG_FINETUNE_IN_AUTO_VENV",
    "NV_SEGMENT_CTMR_ROOT",
    "NV_SEGMENT_CT_ROOT",
    "OMP_NUM_THREADS",
    "PATH",
    "PYTHONNOUSERSITE",
    "PYTHONPATH",
    "PYTORCH_CUDA_ALLOC_CONF",
    "REQUESTS_CA_BUNDLE",
    "SSL_CERT_FILE",
    "TMPDIR",
    "TORCH_EXTENSIONS_DIR",
    "TORCH_HOME",
    "TRANSFORMERS_CACHE",
    "XDG_CACHE_HOME",
)
_MLFLOW_CHILD_ENV_KEYS = (
    "DATABRICKS_CONFIG_PROFILE",
    "DATABRICKS_HOST",
    "DATABRICKS_TOKEN",
    "MLFLOW_TRACKING_CLIENT_CERT_PATH",
    "MLFLOW_TRACKING_INSECURE_TLS",
    "MLFLOW_TRACKING_PASSWORD",
    "MLFLOW_TRACKING_SERVER_CERT_PATH",
    "MLFLOW_TRACKING_TOKEN",
    "MLFLOW_TRACKING_USERNAME",
)
SANITY_REFERENCE_THRESHOLDS = {
    "formal_pretrained_val_dice_min": 0.65,
    "formal_finetuned_val_dice_min": 0.67,
    "formal_improvement_min": 0.005,
    "training_start_val_dice_min": 0.65,
    "training_best_val_dice_min": 0.68,
    "training_improvement_min": 0.005,
}

# Patch ladder keyed on free GPU MiB; calibrated on RTX 6000 Ada (see SKILL.md).
PATCH_LADDER = [
    (int("8_000"), [int("64"), int("64"), int("64")]),
    (int("16_000"), [int("96"), int("96"), int("96")]),
    (int("32_000"), [int("128"), int("128"), int("128")]),
    (int("48_000"), [int("160"), int("160"), int("160")]),
    (int("10") ** int("9"), [int("192"), int("192"), int("128")]),
]
NIFTI_SUFFIXES = (".nii.gz", ".nii")

# Domain knowledge for input-side anatomy checks (adult ranges; informative
# only - the wrapper records, doesn't hard-fail). Volumes in mL.
ANATOMY_VOLUME_ML = {
    "spleen": (int("50"), int("500")),
    "liver": (int("1000"), int("2500")),
    "pancreas": (int("50"), int("200")),
    "stomach": (int("200"), int("1500")),
    "gallbladder": (int("5"), int("80")),
    "right kidney": (int("100"), int("300")),
    "left kidney": (int("100"), int("300")),
    # Lung tumor (MSD06 cancer label) ranges widely from sub-mL nodules
    # to bulky disease; we set a generous adult-thoracic ceiling and a
    # floor that still flags empty-mask bugs.
    "lung tumor": (float("0.05"), int("500")),
}
ANATOMY_EXPECTED_COMPONENTS = {  # solitary organs; user can override
    "spleen": 1,
    "liver": 1,
    "pancreas": 1,
    "stomach": 1,
    "gallbladder": 1,
    "right kidney": 1,
    "left kidney": 1,
    # Tumors are multifocal in general - leave component count unconstrained.
}

app = typer.Typer(add_completion=False)


def _child_process_env(
    overrides: dict[str, str] | None = None,
    *,
    extra_keys: tuple[str, ...] = (),
) -> dict[str, str]:
    """Build the minimal environment needed by trusted local model processes."""
    env: dict[str, str] = {}
    for name in (*_CHILD_ENV_KEYS, *extra_keys):
        value = os.environ.get(name)
        if value is not None:
            env[name] = value
    env.setdefault("PATH", os.defpath)
    if overrides:
        env.update(overrides)
    return env


def _monai_major_minor(monai_version: str) -> tuple[int, int] | None:
    parts = monai_version.split("+", 1)[0].split(".", 2)
    if len(parts) < 2 or not all(p.isdigit() for p in parts[:2]):
        return None
    return int(parts[0]), int(parts[1])


def require_compatible_runtime() -> None:
    """Fail before launching MONAI when the version is outside the tested range."""
    try:
        monai_version = package_version("monai")
    except PackageNotFoundError as exc:
        raise typer.BadParameter(
            "monai is not installed; install `monai==1.4.0` in the active "
            "environment before running this skill."
        ) from exc

    major_minor = _monai_major_minor(monai_version)
    if major_minor not in SUPPORTED_MONAI_MAJOR_MINOR:
        raise typer.BadParameter(
            f"monai==1.4.0 is required for this bundle; found monai " f"{monai_version}."
        )


def _monai_is_compatible() -> bool:
    try:
        monai_version = package_version("monai")
    except PackageNotFoundError:
        return False
    return _monai_major_minor(monai_version) in SUPPORTED_MONAI_MAJOR_MINOR


def maybe_run_compatible_runtime(*, extra_env_keys: tuple[str, ...] = ()) -> None:
    """Use a cached compatible venv when the caller's MONAI is outside this range.

    A child process keeps the user-facing command simple while preserving the
    active environment's CUDA/Torch via --system-site-packages. The DFW reference
    run used MONAI 1.4.0 on Python 3.10; Python 3.12 environments usually
    should still use MONAI 1.4.0 for this upstream trainer.
    """
    if _monai_is_compatible():
        return
    if os.environ.get("NVSEG_FINETUNE_AUTO_VENV") == "0":
        return
    if os.environ.get("NVSEG_FINETUNE_IN_AUTO_VENV") == "1":
        return

    venv_dir = AUTO_VENV_DIR
    python_bin = venv_dir / "bin" / "python"
    if not python_bin.exists():
        venv.EnvBuilder(system_site_packages=True, with_pip=True).create(venv_dir)

    subprocess.check_call(
        [
            str(python_bin),
            "-m",
            "pip",
            "install",
            "monai==1.4.0",
            "numpy<2",
        ],
        stdout=sys.stderr,
        stderr=sys.stderr,
        env=_child_process_env(),
    )
    sys.stderr.write(f"[nv_segment_ct_finetune] run with {python_bin}\n")
    completed = subprocess.run(
        [str(python_bin), *sys.argv],
        env=_child_process_env(
            {"NVSEG_FINETUNE_IN_AUTO_VENV": "1"},
            extra_keys=extra_env_keys,
        ),
        check=False,
    )
    raise typer.Exit(code=completed.returncode)


def require_bundle_files(*, softmax: bool = False) -> Path | None:
    """Fail with setup instructions before MONAI emits a deep config error."""
    bundle_notes = prepare_bundle_files()
    required = [
        BUNDLE_DIR / "configs" / "train.json",
        BUNDLE_DIR / "configs" / "train_continual.json",
        BUNDLE_DIR / "configs" / "metadata.json",
        LABEL_DICT,
        BUNDLE_DIR / "models" / "model.pt",
    ]
    missing = [p for p in required if not p.exists()]
    if not missing:
        if bundle_notes:
            sys.stderr.write(
                "[nv_segment_ct_finetune] prepared bundle files: " + "; ".join(bundle_notes) + "\n"
            )
        return resolve_softmax_bundle_root() if softmax else None

    rel_missing = [
        str(p.relative_to(SKILL_DIR)) if p.is_relative_to(SKILL_DIR) else str(p) for p in missing
    ]
    raise typer.BadParameter(
        "bundle setup is incomplete; missing: "
        + ", ".join(rel_missing)
        + "\nFrom skills/nv-segment-ct-finetune, run:\n"
        + f"  hf download nvidia/NV-Segment-CT --revision {HF_MODEL_REVISION} --local-dir bundle/\n"
        + '  python -c "import urllib.request; '
        + f"urllib.request.urlretrieve('{LABEL_DICT_URL}', "
        + "'bundle/label_dict.json')\"\n"
        + "  python - <<'PY'\n"
        + "from pathlib import Path\n"
        + "import shutil\n"
        + "for src, dst in [(Path('bundle/metadata.json'), Path('bundle/configs/metadata.json')), (Path('bundle/vista3d_pretrained_model/model.pt'), Path('bundle/models/model.pt'))]:\n"
        + "    dst.parent.mkdir(parents=True, exist_ok=True)\n"
        + "    if dst.is_symlink() or not dst.exists():\n"
        + "        dst.unlink(missing_ok=True)\n"
        + "        shutil.copy2(src, dst)\n"
        + "PY\n"
    )


def _unlink_broken_symlink(path: Path) -> bool:
    if path.is_symlink() and not path.exists():
        path.unlink()
        return True
    return False


def _copy_if_missing_or_broken(src: Path, dst: Path) -> bool:
    if not src.exists():
        return False
    dst.parent.mkdir(parents=True, exist_ok=True)
    if dst.is_symlink() and not dst.exists():
        dst.unlink()
    if dst.exists():
        return False
    shutil.copy2(src, dst)
    return True


def _upstream_config_dirs() -> list[Path]:
    """Local upstream checkouts that can seed missing bundle configs."""
    dirs: list[Path] = []
    env_root = os.environ.get("NV_SEGMENT_CT_ROOT", "").strip()
    if env_root:
        dirs.append(Path(env_root) / "configs")
    ctmr_root = os.environ.get("NV_SEGMENT_CTMR_ROOT", "").strip()
    if ctmr_root:
        root = Path(ctmr_root)
        dirs.extend([root / "configs", root.parent / "NV-Segment-CT" / "configs"])
    dirs.extend(
        [
            Path.home()
            / ".cache"
            / "nvidia-skills"
            / "upstreams"
            / "NV-Segment-CTMR-cb921f5"
            / "NV-Segment-CT"
            / "configs",
            Path.home()
            / ".cache"
            / "nvidia-skills"
            / "upstreams"
            / "NV-Segment-CTMR-cb921f5"
            / "NV-Segment-CTMR"
            / "configs",
            _REPO_ROOT
            / ".workbench_data"
            / "upstreams"
            / "NV-Segment-CTMR"
            / "NV-Segment-CT"
            / "configs",
            _REPO_ROOT
            / ".workbench_data"
            / "upstreams"
            / "NV-Segment-CTMR"
            / "NV-Segment-CTMR"
            / "configs",
        ]
    )
    unique: list[Path] = []
    seen: set[Path] = set()
    for path in dirs:
        resolved = path.expanduser()
        if resolved in seen:
            continue
        seen.add(resolved)
        unique.append(resolved)
    return unique


def _git_commit(path: Path) -> str | None:
    """Return the containing checkout commit when Git metadata is available."""
    try:
        return subprocess.check_output(
            ["git", "-C", str(path), "rev-parse", "HEAD"],
            text=True,
            stderr=subprocess.DEVNULL,
        ).strip()
    except (FileNotFoundError, subprocess.CalledProcessError):
        return None


def resolve_softmax_bundle_root() -> Path:
    """Find the pinned upstream bundle that owns the documented softmax workflow."""
    candidates = sorted(
        (config_dir.parent for config_dir in _upstream_config_dirs()),
        key=lambda path: (path.name != "NV-Segment-CT", str(path)),
    )
    for root in candidates:
        required = (
            root / "configs" / "train_continual_softmax.json",
            root / "configs" / "inference_softmax.json",
            root / "scripts" / "vista3d_softmax.py",
        )
        if not all(path.is_file() for path in required):
            continue
        commit = _git_commit(root)
        if commit is not None and commit != UPSTREAM_CTMR_COMMIT:
            continue
        return root

    raise typer.BadParameter(
        "--softmax requires the pinned NVIDIA-Medtech/NV-Segment-CTMR checkout "
        f"at {UPSTREAM_CTMR_COMMIT}. Clone it once, check out that commit, and "
        "set NV_SEGMENT_CT_ROOT to its NV-Segment-CT directory (or "
        "NV_SEGMENT_CTMR_ROOT to its NV-Segment-CTMR directory)."
    )


def _copy_upstream_config(name: str, *, overwrite_if_different: bool = False) -> bool:
    dst = BUNDLE_DIR / "configs" / name
    for config_dir in _upstream_config_dirs():
        src = config_dir / name
        if not src.exists():
            continue
        if dst.is_symlink() and not dst.exists():
            dst.unlink()
        if not dst.exists():
            dst.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(src, dst)
            return True
        if overwrite_if_different and src.read_bytes() != dst.read_bytes():
            shutil.copy2(src, dst)
            return True
        if dst.exists():
            return False
    return False


class _NoLabelDictRedirects(urllib.request.HTTPRedirectHandler):
    """Keep the pinned label dictionary request on its reviewed HTTPS origin."""

    def redirect_request(self, req, fp, code, msg, headers, newurl):
        return None


def _download_label_dict(dst: Path) -> bool:
    if dst.exists():
        return False
    try:
        source = urllib.parse.urlsplit(LABEL_DICT_URL)
        allowed_source = (
            source.scheme == "https"
            and source.hostname == "raw.githubusercontent.com"
            and source.port in (None, 443)
            and source.username is None
            and source.password is None
            and not source.fragment
        )
    except ValueError:
        return False
    if not allowed_source:
        return False
    dst.parent.mkdir(parents=True, exist_ok=True)
    # Check the scheme/origin before I/O and reject every redirect. In particular,
    # a response cannot downgrade HTTPS or send the request to another host.
    opener = urllib.request.build_opener(_NoLabelDictRedirects())
    try:
        with opener.open(LABEL_DICT_URL, timeout=30) as response:
            payload = response.read()
    except (OSError, urllib.error.URLError):
        return False
    data = json.loads(payload.decode("utf-8"))
    if not isinstance(data, dict) or "lung tumor" not in data:
        return False
    dst.write_text(json.dumps(data, indent=2) + "\n")
    return True


def _fixture_preset(fixture: Path) -> str | None:
    name = fixture.name
    if name == "spleen_micro" or name.startswith("spleen_micro"):
        return "smoke"
    if name in {"Task06", "Task06_Lung"} or name.startswith("Task06_"):
        return "sanity"
    return None


def _resolve_sanity_dataset(fixture: Optional[Path], dataset_dir: Optional[Path]) -> Path:
    if dataset_dir is not None:
        return dataset_dir.resolve()
    if fixture is not None and fixture.is_dir():
        return fixture.resolve()
    return SANITY_DATASET


def prepare_bundle_files() -> list[str]:
    """Make the local downloaded bundle usable in fresh agent commands.

    `hf download --local-dir` can leave old local symlinks untouched when a
    previous checkout used a different skill path. Repairing those files here
    keeps the user-facing command idempotent without requiring shell cleanup.
    """
    notes: list[str] = []
    for rel in (
        "label_dict.json",
        "configs/metadata.json",
        "models/model.pt",
    ):
        if _unlink_broken_symlink(BUNDLE_DIR / rel):
            notes.append(f"removed dangling {rel}")

    sibling_label_dict = SKILL_DIR.parent / "nv-segment-ct" / "bundle" / "label_dict.json"
    if _copy_if_missing_or_broken(sibling_label_dict, LABEL_DICT):
        notes.append("copied label_dict.json from nv-segment-ct cache")
    if _download_label_dict(LABEL_DICT):
        notes.append("downloaded label_dict.json from NVIDIA-Medtech/NV-Segment-CTMR")

    for config_name in (
        "train.json",
        "train_continual.json",
        "multi_gpu_train.json",
        "evaluate.json",
    ):
        if _copy_upstream_config(config_name, overwrite_if_different=True):
            notes.append(f"restored configs/{config_name} from local upstream cache")

    needed_sources = [
        BUNDLE_DIR / "configs" / "train.json",
        BUNDLE_DIR / "configs" / "train_continual.json",
        BUNDLE_DIR / "metadata.json",
        BUNDLE_DIR / "vista3d_pretrained_model" / "model.pt",
    ]
    if not all(p.exists() for p in needed_sources):
        try:
            from huggingface_hub import snapshot_download
        except ImportError:
            return notes
        snapshot_download(
            repo_id="nvidia/NV-Segment-CT",
            revision=HF_MODEL_REVISION,
            local_dir=str(BUNDLE_DIR),
            local_dir_use_symlinks=False,
        )
        notes.append("downloaded nvidia/NV-Segment-CT bundle")

        for config_name in (
            "train.json",
            "train_continual.json",
            "multi_gpu_train.json",
            "evaluate.json",
        ):
            if _copy_upstream_config(config_name, overwrite_if_different=True):
                notes.append(f"restored configs/{config_name} from local upstream cache")

    if _copy_if_missing_or_broken(
        BUNDLE_DIR / "metadata.json",
        BUNDLE_DIR / "configs" / "metadata.json",
    ):
        notes.append("staged configs/metadata.json")
    if _copy_if_missing_or_broken(
        BUNDLE_DIR / "vista3d_pretrained_model" / "model.pt",
        BUNDLE_DIR / "models" / "model.pt",
    ):
        notes.append("staged models/model.pt")
    return notes


def _mean_dice_accepts_num_classes() -> bool:
    try:
        from monai.handlers import MeanDice
    except Exception:
        return True
    return "num_classes" in inspect.signature(MeanDice).parameters


def metric_compat_config_stack() -> list[str]:
    """Return metric-compat config files only when the runtime needs them."""
    if _mean_dice_accepts_num_classes():
        return []
    return [
        write_config(
            "mean_dice_no_num_classes.json",
            {
                "validate#key_metric#val_mean_dice": {
                    "_target_": "MeanDice",
                    "include_background": False,
                    "output_transform": "$monai.handlers.from_engine(['pred', 'label'])",
                }
            },
        )
    ]


def write_config(name: str, payload: dict) -> str:
    """Write a bundle config under configs/ and return its MONAI stack path."""
    cfg = BUNDLE_DIR / "configs" / name
    cfg.parent.mkdir(parents=True, exist_ok=True)
    cfg.write_text(json.dumps(payload, indent=2))
    return f"configs/{name}"


def write_config_path(path: Path, payload: dict) -> str:
    """Write a generated override outside an upstream checkout."""
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(payload, indent=2))
    return str(path)


# --- environment + plan -----------------------------------------------------


def detect_env() -> dict:
    try:
        rows = (
            subprocess.check_output(
                [
                    "nvidia-smi",
                    "--query-gpu=name,memory.total,memory.free",
                    "--format=csv,noheader,nounits",
                ],
                text=True,
                stderr=subprocess.DEVNULL,
            )
            .strip()
            .splitlines()
        )
        n = len(rows)
        name, total, free = (x.strip() for x in rows[0].split(","))
        gpu_name, total_mb, free_mb = name, int(total), int(free)
    except (FileNotFoundError, subprocess.CalledProcessError, IndexError):
        n, gpu_name, total_mb, free_mb = 0, "cpu", 0, 0
    try:
        with open("/proc/meminfo") as f:
            ram_mb = int(next(line for line in f if line.startswith("MemTotal")).split()[1]) // int(
                "1024"
            )
    except (OSError, StopIteration):
        ram_mb = 0
    packages: dict[str, str | None] = {}
    for package in ("monai", "torch", "nibabel", "scipy", "typer", "PyYAML"):
        try:
            packages[package] = package_version(package)
        except PackageNotFoundError:
            packages[package] = None
    try:
        import torch  # type: ignore

        torch_cuda = torch.version.cuda
        torch_cuda_available = bool(torch.cuda.is_available())
    except Exception:
        torch_cuda = None
        torch_cuda_available = False
    return {
        "gpu_count": n,
        "gpu_name": gpu_name,
        "gpu_total_mb": total_mb,
        "gpu_free_mb": free_mb,
        "host_ram_mb": ram_mb,
        "cuda_available": n > 0,
        "python": sys.version.split()[0],
        "packages": packages,
        "torch_cuda": torch_cuda,
        "torch_cuda_available": torch_cuda_available,
    }


def pick_patch(free_mb: int) -> list[int]:
    for ceiling, patch in PATCH_LADDER:
        if free_mb < ceiling:
            return patch
    return PATCH_LADDER[-1][1]


def pick_cache_rate(n_train: int, ram_mb: int) -> float:
    if n_train <= 0 or ram_mb <= 0:
        return float("0.1")
    return round(max(0.0, min(1.0, ram_mb * float("0.25") / (n_train * int("50")))), 2)


def pick_nproc(gpu_count: int) -> int:
    override = os.environ.get("NPROC_PER_NODE")
    if override:
        try:
            return max(1, int(override))
        except ValueError:
            pass
    return max(1, gpu_count)


# --- dataset inspection -----------------------------------------------------


def _strip_nifti(name: str) -> str:
    for s in NIFTI_SUFFIXES:
        if name.endswith(s):
            return name[: -len(s)]
    return name


def _list_nifti(d: Path) -> list[Path]:
    out: list[Path] = []
    if d.is_dir():
        for s in NIFTI_SUFFIXES:
            out.extend(p for p in d.glob(f"*{s}") if not p.name.startswith("."))
    return sorted({p.resolve(): None for p in out})


def _resolve_dataset_path(dataset_dir: Path, raw: str) -> Path:
    path = Path(raw)
    return path if path.is_absolute() else dataset_dir / path


def _audit_volume(img_path: Path, lab_path: Path, user_idx: int, anatomy: Optional[str]) -> dict:
    """Read one image+label pair fully and compute domain-side facts:
    orientation, HU range, spacing, foreground volume in mL, connected-
    component count, plus anatomy-specific bounds checks when known.
    Used on the sampled subset only (cheap per pair, ~1s)."""
    img = nib.load(str(img_path))
    lab = nib.load(str(lab_path))
    img_arr = np.asarray(img.dataobj)
    lab_arr = np.asarray(lab.dataobj).astype(int)
    spacing = tuple(float(z) for z in img.header.get_zooms()[: int("3")])
    vox_mm3 = abs(spacing[0] * spacing[1] * spacing[2])
    fg = lab_arr == user_idx
    fg_vox = int(fg.sum())
    fg_ml = round(fg_vox * vox_mm3 / float("1000.0"), 1)
    n_components = 0
    if fg_vox > 0:
        _, n_components = ndi.label(fg)
    img_min, img_max = float(img_arr.min()), float(img_arr.max())
    out = {
        "case": img_path.name,
        "orientation_code": "".join(nib.orientations.aff2axcodes(img.affine)),
        "spacing_mm": [round(s, int("4")) for s in spacing],
        "voxel_volume_mm3": round(vox_mm3, int("4")),
        "image_dtype": str(img_arr.dtype),
        "image_hu_min": img_min,
        "image_hu_max": img_max,
        "image_hu_looks_like_ct": img_min < -int("500") and img_max > 0,
        "label_dtype": str(lab_arr.dtype),
        "fg_voxels": fg_vox,
        "fg_volume_ml": fg_ml,
        "fg_components": int(n_components),
    }
    if anatomy:
        key = anatomy.strip().lower()
        if key in ANATOMY_VOLUME_ML:
            lo, hi = ANATOMY_VOLUME_ML[key]
            out["anatomy_volume_in_range"] = lo <= fg_ml <= hi
            out["anatomy_volume_expected_ml"] = [lo, hi]
        if key in ANATOMY_EXPECTED_COMPONENTS:
            out["anatomy_components_match"] = n_components == ANATOMY_EXPECTED_COMPONENTS[key]
            out["anatomy_components_expected"] = ANATOMY_EXPECTED_COMPONENTS[key]
    return out


def inspect_and_build_datalist(
    dataset_dir: Path,
    output_dir: Path,
    user_label_idx: int,
    anatomy: Optional[str] = None,
) -> tuple[Path, dict]:
    """Pair imagesTr/* with labelsTr/*, verify every pair, write 5-fold datalist."""
    images = _list_nifti(dataset_dir / "imagesTr")
    if not images:
        raise typer.BadParameter(
            f"expected NIfTI under {dataset_dir}/imagesTr + labelsTr (MSD layout)"
        )
    pairs, bad = [], []
    shapes, spacings, max_drift = set(), [], 0.0
    for img_p in images:
        stem = _strip_nifti(img_p.name)
        lab_p = next(
            (
                dataset_dir / "labelsTr" / f"{stem}{s}"
                for s in NIFTI_SUFFIXES
                if (dataset_dir / "labelsTr" / f"{stem}{s}").exists()
            ),
            None,
        )
        if lab_p is None:
            bad.append({"image": img_p.name, "reason": "no matching label"})
            continue
        try:
            img, lab = nib.load(str(img_p)), nib.load(str(lab_p))
        except Exception as e:
            bad.append({"image": img_p.name, "reason": f"nib.load: {e}"})
            continue
        if tuple(img.shape) != tuple(lab.shape):
            bad.append(
                {
                    "image": img_p.name,
                    "reason": f"shape {tuple(img.shape)} vs {tuple(lab.shape)}",
                }
            )
            continue
        drift = float(np.max(np.abs(np.asarray(img.affine) - np.asarray(lab.affine))))
        if drift > float("1e-3"):
            bad.append({"image": img_p.name, "reason": f"affine drift {drift:.4g}"})
            continue
        max_drift = max(max_drift, drift)
        shapes.add(tuple(img.shape))
        spacings.append(tuple(float(z) for z in img.header.get_zooms()[: int("3")]))
        pairs.append(
            {
                "image": str(img_p.relative_to(dataset_dir)),
                "label": str(lab_p.relative_to(dataset_dir)),
            }
        )
    if not pairs:
        raise typer.BadParameter(f"no valid pairs; first bad: {bad[:3]}")

    # Per-volume domain audit on a sample (orientation, HU range, foreground
    # volume, components, anatomy bounds). Cheap: ~1s per case.
    sampled = []
    seen_labels: set[int] = set()
    for p in pairs[: min(int("5"), len(pairs))]:
        img_p, lab_p = dataset_dir / p["image"], dataset_dir / p["label"]
        seen_labels.update(int(v) for v in np.unique(np.asarray(nib.load(str(lab_p)).dataobj)))
        sampled.append(_audit_volume(img_p, lab_p, user_label_idx, anatomy))
    orient_codes = sorted({s["orientation_code"] for s in sampled})

    split_pairs = list(pairs)
    random.Random(0).shuffle(split_pairs)
    for i, item in enumerate(split_pairs):
        item["fold"] = i % int("5")
    path = output_dir / "auto_datalist.json"
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps({"training": split_pairs, "testing": []}, indent=2))

    audit = {
        "dataset_dir": str(dataset_dir),
        "datalist_source": "auto",
        "datalist_path": str(path),
        "n_pairs": len(split_pairs),
        "n_folds": int("5"),
        "fold_assignment": "random_seed_0_round_robin",
        "shape_consistent": len(shapes) == 1,
        "spacing_range": (
            (
                [round(min(c), int("4")) for c in zip(*spacings)],
                [round(max(c), int("4")) for c in zip(*spacings)],
            )
            if spacings
            else None
        ),
        "affine_max_drift_max": round(max_drift, int("6")),
        "label_uniques_sampled": sorted(seen_labels),
        "user_label_idx": user_label_idx,
        "user_label_idx_present_in_sample": user_label_idx in seen_labels,
        "bad_pairs": bad,
        # Aggregated domain checks over the sample.
        "n_sampled_for_domain_checks": len(sampled),
        "orientation_codes_seen": orient_codes,
        "orientation_consistent": len(orient_codes) == 1,
        "image_dtypes_seen": sorted({s["image_dtype"] for s in sampled}),
        "label_dtypes_seen": sorted({s["label_dtype"] for s in sampled}),
        "image_hu_range_seen": (
            [
                min(s["image_hu_min"] for s in sampled),
                max(s["image_hu_max"] for s in sampled),
            ]
            if sampled
            else None
        ),
        "image_looks_like_ct": all(s["image_hu_looks_like_ct"] for s in sampled),
        "fg_volumes_ml_seen": [s["fg_volume_ml"] for s in sampled],
        "fg_components_seen": [s["fg_components"] for s in sampled],
        "anatomy": anatomy,
        "anatomy_volume_all_in_range": (
            all(s.get("anatomy_volume_in_range") for s in sampled)
            if anatomy and any("anatomy_volume_in_range" in s for s in sampled)
            else None
        ),
        "anatomy_components_all_match": (
            all(s.get("anatomy_components_match") for s in sampled)
            if anatomy and any("anatomy_components_match" in s for s in sampled)
            else None
        ),
        "per_sample": sampled,
    }
    return path, audit


def audit_existing_datalist(
    dataset_dir: Path,
    datalist: Path,
    user_label_idx: int,
    anatomy: Optional[str] = None,
) -> dict:
    """Audit a caller-provided datalist without rewriting its split."""
    data = json.loads(datalist.read_text())
    entries = list(data.get("training", []))
    bad: list[dict] = []
    shapes, spacings, max_drift = set(), [], 0.0
    sampled = []
    seen_labels: set[int] = set()

    for idx, item in enumerate(entries):
        img_raw, lab_raw = item.get("image"), item.get("label")
        if not isinstance(img_raw, str) or not isinstance(lab_raw, str):
            bad.append({"index": idx, "reason": "missing image or label path"})
            continue
        img_p = _resolve_dataset_path(dataset_dir, img_raw)
        lab_p = _resolve_dataset_path(dataset_dir, lab_raw)
        if not img_p.exists() or not lab_p.exists():
            bad.append(
                {
                    "index": idx,
                    "image": img_raw,
                    "label": lab_raw,
                    "reason": "image or label file missing",
                }
            )
            continue
        try:
            img, lab = nib.load(str(img_p)), nib.load(str(lab_p))
        except Exception as e:
            bad.append({"index": idx, "image": img_raw, "reason": f"nib.load: {e}"})
            continue
        if tuple(img.shape) != tuple(lab.shape):
            bad.append(
                {
                    "index": idx,
                    "image": img_raw,
                    "reason": f"shape {tuple(img.shape)} vs {tuple(lab.shape)}",
                }
            )
            continue
        drift = float(np.max(np.abs(np.asarray(img.affine) - np.asarray(lab.affine))))
        if drift > float("1e-3"):
            bad.append({"index": idx, "image": img_raw, "reason": f"affine drift {drift:.4g}"})
            continue
        max_drift = max(max_drift, drift)
        shapes.add(tuple(img.shape))
        spacings.append(tuple(float(z) for z in img.header.get_zooms()[: int("3")]))
        if len(sampled) < int("5"):
            seen_labels.update(int(v) for v in np.unique(np.asarray(nib.load(str(lab_p)).dataobj)))
            sampled.append(_audit_volume(img_p, lab_p, user_label_idx, anatomy))

    orient_codes = sorted({s["orientation_code"] for s in sampled})
    return {
        "dataset_dir": str(dataset_dir),
        "datalist_source": "caller_provided",
        "datalist_path": str(datalist),
        "n_pairs": len(entries),
        "shape_consistent": len(shapes) <= 1,
        "spacing_range": (
            (
                [round(min(c), int("4")) for c in zip(*spacings)],
                [round(max(c), int("4")) for c in zip(*spacings)],
            )
            if spacings
            else None
        ),
        "affine_max_drift_max": round(max_drift, int("6")),
        "label_uniques_sampled": sorted(seen_labels),
        "user_label_idx": user_label_idx,
        "user_label_idx_present_in_sample": user_label_idx in seen_labels,
        "bad_pairs": bad,
        "n_sampled_for_domain_checks": len(sampled),
        "orientation_codes_seen": orient_codes,
        "orientation_consistent": len(orient_codes) <= 1,
        "image_dtypes_seen": sorted({s["image_dtype"] for s in sampled}),
        "label_dtypes_seen": sorted({s["label_dtype"] for s in sampled}),
        "image_hu_range_seen": (
            [
                min(s["image_hu_min"] for s in sampled),
                max(s["image_hu_max"] for s in sampled),
            ]
            if sampled
            else None
        ),
        "image_looks_like_ct": (
            all(s["image_hu_looks_like_ct"] for s in sampled) if sampled else False
        ),
        "fg_volumes_ml_seen": [s["fg_volume_ml"] for s in sampled],
        "fg_components_seen": [s["fg_components"] for s in sampled],
        "anatomy": anatomy,
        "anatomy_volume_all_in_range": (
            all(s.get("anatomy_volume_in_range") for s in sampled)
            if anatomy and any("anatomy_volume_in_range" in s for s in sampled)
            else None
        ),
        "anatomy_components_all_match": (
            all(s.get("anatomy_components_match") for s in sampled)
            if anatomy and any("anatomy_components_match" in s for s in sampled)
            else None
        ),
        "per_sample": sampled,
    }


def ensure_smoke_dataset(
    dataset_dir: Path, datalist: Path, output_dir: Path
) -> tuple[Path, Path, bool]:
    """Materialize synthetic smoke NIfTIs when the fixture ships only a datalist."""
    data = json.loads(datalist.read_text())
    entries = list(data.get("training", []))
    missing = []
    for item in entries:
        image = item.get("image")
        label = item.get("label")
        if not isinstance(image, str) or not isinstance(label, str):
            continue
        if not _resolve_dataset_path(dataset_dir, image).exists():
            missing.append(image)
        if not _resolve_dataset_path(dataset_dir, label).exists():
            missing.append(label)

    if not missing:
        return dataset_dir, datalist, False

    work_dir = output_dir / "smoke_dataset"
    if work_dir.exists():
        shutil.rmtree(work_dir)
    (work_dir / "imagesTr").mkdir(parents=True, exist_ok=True)
    (work_dir / "labelsTr").mkdir(parents=True, exist_ok=True)

    shape = (int("64"), int("64"), int("64"))
    grid = np.indices(shape)
    center = np.array(shape, dtype=float) / float("2")
    radius = float("12")
    affine = np.diag([float("1"), float("1"), float("1"), float("1")])

    for idx, item in enumerate(entries):
        image = item.get("image")
        label = item.get("label")
        if not isinstance(image, str) or not isinstance(label, str):
            continue
        shift = np.array([idx % 2, (idx // 2) % 2, idx % 3], dtype=float) * float("2")
        dist = np.sqrt(
            ((grid - (center[:, None, None, None] + shift[:, None, None, None])) ** 2).sum(axis=0)
        )
        label_arr = (dist <= radius).astype(np.uint8)
        image_arr = np.full(shape, -900.0, dtype=np.float32)
        image_arr[label_arr > 0] = 80.0 + float(idx)
        image_arr += np.random.default_rng(idx).normal(0.0, 5.0, size=shape).astype(np.float32)

        image_path = work_dir / image
        label_path = work_dir / label
        image_path.parent.mkdir(parents=True, exist_ok=True)
        label_path.parent.mkdir(parents=True, exist_ok=True)
        nib.save(nib.Nifti1Image(image_arr, affine), str(image_path))
        nib.save(nib.Nifti1Image(label_arr, affine), str(label_path))

    staged_datalist = work_dir / "datalist.json"
    staged_datalist.write_text(json.dumps(data, indent=2))
    return work_dir, staged_datalist, True


def resolve_mapping(
    target_anatomy: Optional[str], user_idx: int, literal: Optional[str]
) -> tuple[dict, dict]:
    if literal:
        m = json.loads(literal)
        return {"default": m}, {"source": "literal", "value": m}
    if not target_anatomy:
        raise typer.BadParameter("pass --target-anatomy or --label-mapping")
    if not LABEL_DICT.exists():
        raise typer.BadParameter(
            f"label_dict.json missing at {LABEL_DICT}; "
            f"run `hf download nvidia/NV-Segment-CT --revision {HF_MODEL_REVISION}` first"
        )
    d = {
        str(k).strip().lower(): int(v)
        for k, v in json.loads(LABEL_DICT.read_text()).items()
        if isinstance(v, int)
    }
    key = target_anatomy.strip().lower()
    if key not in d:
        raise typer.BadParameter(
            f"{target_anatomy!r} not in label_dict.json; "
            f"closest: {[k for k in d if key in k][:10]}"
        )
    return (
        {"default": [[user_idx, d[key]]]},
        {
            "source": "anatomy_lookup",
            "anatomy": target_anatomy,
            "user_idx": user_idx,
            "vista3d_idx": d[key],
        },
    )


def validate_softmax_mapping(mapping: dict) -> None:
    """Validate the ordered foreground mapping required by the softmax config."""
    entries = mapping.get("default")
    if not isinstance(entries, list) or not entries:
        raise typer.BadParameter("--softmax requires at least one label mapping entry")
    if any(
        not isinstance(pair, list)
        or len(pair) != 2
        or not all(isinstance(value, int) for value in pair)
        for pair in entries
    ):
        raise typer.BadParameter(
            "--softmax label mappings must be ordered [dataset_label, vista_class_id] integer pairs"
        )
    dataset_labels = [pair[0] for pair in entries]
    vista_class_ids = [pair[1] for pair in entries]
    if any(value <= 0 for value in dataset_labels + vista_class_ids):
        raise typer.BadParameter(
            "--softmax mappings contain foreground labels only; dataset labels and "
            "VISTA class IDs must be positive"
        )
    if len(set(dataset_labels)) != len(dataset_labels):
        raise typer.BadParameter("--softmax dataset labels must be unique")
    if len(set(vista_class_ids)) != len(vista_class_ids):
        raise typer.BadParameter("--softmax VISTA class IDs must be unique")


# --- bundle run + log parse -------------------------------------------------


def build_override(
    dataset_dir: Path,
    datalist: Path,
    mapping: dict,
    patch: list[int],
    cache_rate: float,
    epochs: int,
    lr: float,
    ckpt_dir: Path,
    train_output_dir: Path,
    auto_seg: bool = False,
) -> dict:
    """Compose the JSON override layered on top of train.json + train_continual.json.

    `auto_seg=True` mirrors the published MSD06 lung-tumor tutorial:
    `drop_label_prob=0.0, drop_point_prob=1.0` forces automatic segmentation
    (no point prompts during training), and `resample_to_spacing` is pinned
    to the tutorial's 1.5 mm isotropic. Default leaves both prompt
    probabilities at the bundle's mixed-prompt training values.
    """
    override = {
        "dataset_dir": str(dataset_dir),
        "data_list_file_path": str(datalist),
        "image_key": "image",
        "label_key": "label",
        "finetune": True,
        "finetune_model_path": str(BUNDLE_DIR / "models" / "model.pt"),
        "ckpt_dir": str(ckpt_dir),
        "output_dir": str(train_output_dir),
        "patch_size": patch,
        "patch_size_valid": patch,
        "label_mappings": mapping,
        "epochs": epochs,
        "val_interval": 1,
        "val_at_start": True,
        "learning_rate": lr,
        "lr_schedule#activate": False,
        "train_dataset_cache_rate": cache_rate,
        "val_dataset_cache_rate": cache_rate,
    }
    if auto_seg:
        override.update(
            {
                "drop_label_prob": 0.0,
                "drop_point_prob": 1.0,
                "resample_to_spacing": tuple(float(x) for x in ("1.5", "1.5", "1.5")),
            }
        )
    return override


def build_softmax_override(
    dataset_dir: Path,
    datalist: Path,
    mapping: dict,
    patch: list[int],
    cache_rate: float,
    epochs: int,
    lr: float,
    ckpt_dir: Path,
    train_output_dir: Path,
) -> dict:
    """Override the official fixed-channel softmax config without patching it."""
    return {
        "dataset_dir": str(dataset_dir),
        "data_list_file_path": str(datalist),
        "image_key": "image",
        "finetune": True,
        "finetune_model_path": str(BUNDLE_DIR / "models" / "model.pt"),
        "ckpt_dir": str(ckpt_dir),
        "output_dir": str(train_output_dir),
        "patch_size": patch,
        "patch_size_valid": patch,
        "label_mappings": mapping,
        "epochs": epochs,
        "val_interval": 1,
        "val_at_start": True,
        "learning_rate": lr,
        "train_dataset_cache_rate": cache_rate,
        "val_dataset_cache_rate": cache_rate,
    }


def _config_arg(stack: list[str]) -> str:
    cfg_arg = "[" + ",".join(f"'{s}'" for s in stack) + "]"
    return cfg_arg


def _peak_gpu_mb(gpu_csv: Path) -> int:
    peak = 0
    if not gpu_csv.exists():
        return peak
    for line in gpu_csv.read_text().splitlines():
        parts = [p.strip() for p in line.split(",")]
        if not parts:
            continue
        try:
            peak = max(peak, int(parts[-1]))
        except ValueError:
            pass
    return peak


def run_monai_bundle(
    stack: list[str],
    log_path: Path,
    *,
    multi_gpu: bool = False,
    nproc: int = 1,
    extra_args: Optional[list[str]] = None,
    force_single_gpu: bool = False,
    bundle_root: Optional[Path] = None,
    extra_env_keys: tuple[str, ...] = (),
) -> tuple[int, int, list[str]]:
    resolved_bundle_root = bundle_root or BUNDLE_DIR
    cfg_arg = _config_arg(stack)
    if multi_gpu:
        cmd = [
            "torchrun",
            "--nnodes=1",
            f"--nproc_per_node={nproc}",
            "-m",
            "monai.bundle",
            "run",
            "--config_file",
            cfg_arg,
            "--bundle_root",
            str(resolved_bundle_root),
        ]
    else:
        cmd = [
            sys.executable,
            "-m",
            "monai.bundle",
            "run",
            "--config_file",
            cfg_arg,
            "--bundle_root",
            str(resolved_bundle_root),
        ]
    if extra_args:
        cmd.extend(extra_args)

    log_path.parent.mkdir(parents=True, exist_ok=True)
    gpu_csv = log_path.with_suffix(".gpu.csv")
    smi = None
    smi_out = None
    nvidia_smi = shutil.which("nvidia-smi")
    if nvidia_smi is not None:
        smi_cmd = [
            nvidia_smi,
            "--query-gpu=timestamp,index,memory.used",
            "--format=csv,noheader,nounits",
            "-l",
            "1",
        ]
        if force_single_gpu:
            smi_cmd[1:1] = ["-i", "0"]
        smi_out = open(gpu_csv, "w")
        smi = subprocess.Popen(
            smi_cmd,
            stdout=smi_out,
            stderr=subprocess.DEVNULL,
            env=_child_process_env(),
        )
    try:
        env = _child_process_env(extra_keys=extra_env_keys)
        if force_single_gpu and "CUDA_VISIBLE_DEVICES" not in env:
            env["CUDA_VISIBLE_DEVICES"] = "0"
        with open(log_path, "w") as f:
            rc = subprocess.call(
                cmd,
                cwd=resolved_bundle_root,
                stdout=f,
                stderr=subprocess.STDOUT,
                env=env,
            )
    finally:
        if smi is not None:
            smi.terminate()
            try:
                smi.wait(timeout=int("3"))
            except subprocess.TimeoutExpired:
                smi.kill()
        if smi_out is not None:
            smi_out.close()
    return rc, _peak_gpu_mb(gpu_csv), cmd


_LOSS = re.compile(r"train_loss:\s*([0-9.eE+-]+)")
_DICE = re.compile(r"val_mean_dice:\s*([0-9.eE+-]+)")
_OOM = re.compile(r"CUDA out of memory|OutOfMemoryError")


def parse_log(log_path: Path) -> dict:
    text = log_path.read_text() if log_path.exists() else ""
    losses = [float(m.group(1)) for m in _LOSS.finditer(text)]
    dices = [float(m.group(1)) for m in _DICE.finditer(text)]
    best = max(dices) if dices else None
    return {
        "train_loss_first": losses[0] if losses else None,
        "train_loss_last": losses[-1] if losses else None,
        "train_loss_finite": (
            all(loss == loss and abs(loss) != float("inf") for loss in losses) if losses else False
        ),
        "val_dice_per_epoch": dices,
        "baseline_val_dice": dices[0] if dices else None,
        "best_val_dice": best,
        "best_epoch_index": dices.index(best) if best is not None else None,
        "oom": bool(_OOM.search(text)),
        "log_tail": text.splitlines()[-int("25") :],
    }


def build_mlflow_tracking_args(
    output_dir: Path,
    *,
    tracking_uri: str | None,
    experiment_name: str,
    requested_run_name: str | None,
) -> tuple[list[str], dict[str, str]]:
    """Use MONAI's built-in MLflow tracking without changing training config."""
    uri = tracking_uri or (output_dir / "mlruns").resolve().as_uri()
    metadata = {
        "tracking_uri": uri,
        "experiment_name": experiment_name,
    }
    args = [
        "--tracking",
        "mlflow",
        "--tracking_uri",
        uri,
        "--experiment_name",
        experiment_name,
        "--save_execute_config",
        "False",
    ]
    if requested_run_name:
        args.extend(["--run_name", requested_run_name])
        metadata["run_name"] = requested_run_name
    return args, metadata


def read_val_mean_dice(metrics_dir: Path) -> float | None:
    metrics_csv = metrics_dir / "metrics.csv"
    if not metrics_csv.exists():
        return None
    for line in metrics_csv.read_text().splitlines():
        parts = [p.strip() for p in line.split(",")]
        if len(parts) >= 2 and parts[0] == "val_mean_dice":
            try:
                return float(parts[1])
            except ValueError:
                return None
    return None


def _extract_state_dict(obj: object) -> dict | None:
    if not isinstance(obj, dict):
        return None
    values = list(obj.values())
    if values and all(hasattr(v, "shape") for v in values):
        return obj
    for key in ("state_dict", "model", "network"):
        child = obj.get(key)
        if isinstance(child, dict):
            return child
    return None


def compare_checkpoint_weights(reference: Path, candidate: Path) -> dict:
    """Compare checkpoint tensors, not file bytes.

    MONAI/Ignite can reserialize identical weights into a different file, so
    sha256 alone is insufficient for detecting the epoch-0 checkpoint trap.
    """
    out = {
        "reference": str(reference),
        "candidate": str(candidate),
        "compared": False,
        "weights_identical": None,
    }
    if not reference.exists() or not candidate.exists():
        out["error"] = "reference or candidate checkpoint missing"
        return out
    try:
        import torch  # type: ignore

        # Only tensor state dictionaries and safe primitive metadata are needed.
        # Unsupported checkpoint objects fail closed; never retry unsafe pickle.
        ref_obj = torch.load(reference, map_location="cpu", weights_only=True)
        cand_obj = torch.load(candidate, map_location="cpu", weights_only=True)
        ref_state = _extract_state_dict(ref_obj)
        cand_state = _extract_state_dict(cand_obj)
        if ref_state is None or cand_state is None:
            out["error"] = "could not extract tensor state dict"
            return out
        ref_keys = set(ref_state)
        cand_keys = set(cand_state)
        shared_keys = sorted(ref_keys & cand_keys)
        differing_tensors = 0
        shape_or_dtype_mismatches = 0
        tensor_count = 0
        max_abs_diff = 0.0
        total_abs_diff = 0.0
        examples: list[dict] = []
        for key in shared_keys:
            ref_value = ref_state[key]
            cand_value = cand_state[key]
            if not (torch.is_tensor(ref_value) and torch.is_tensor(cand_value)):
                continue
            tensor_count += 1
            if ref_value.shape != cand_value.shape or ref_value.dtype != cand_value.dtype:
                shape_or_dtype_mismatches += 1
                if len(examples) < 5:
                    examples.append(
                        {
                            "key": key,
                            "reference_shape": list(ref_value.shape),
                            "candidate_shape": list(cand_value.shape),
                            "reference_dtype": str(ref_value.dtype),
                            "candidate_dtype": str(cand_value.dtype),
                        }
                    )
                continue
            if torch.equal(ref_value, cand_value):
                continue
            differing_tensors += 1
            diff = (ref_value.float() - cand_value.float()).abs()
            tensor_max = float(diff.max().item()) if diff.numel() else 0.0
            tensor_sum = float(diff.sum().item()) if diff.numel() else 0.0
            max_abs_diff = max(max_abs_diff, tensor_max)
            total_abs_diff += tensor_sum
            if len(examples) < 5:
                examples.append(
                    {
                        "key": key,
                        "shape": list(ref_value.shape),
                        "max_abs_diff": tensor_max,
                        "sum_abs_diff": tensor_sum,
                    }
                )
        missing = sorted(ref_keys - cand_keys)
        extra = sorted(cand_keys - ref_keys)
        weights_identical = (
            not missing and not extra and shape_or_dtype_mismatches == 0 and differing_tensors == 0
        )
        out.update(
            {
                "compared": True,
                "same_keys": not missing and not extra,
                "missing_keys_count": len(missing),
                "extra_keys_count": len(extra),
                "tensor_count": tensor_count,
                "differing_tensors": differing_tensors,
                "shape_or_dtype_mismatches": shape_or_dtype_mismatches,
                "max_abs_diff": max_abs_diff,
                "total_abs_diff": total_abs_diff,
                "weights_identical": weights_identical,
                "examples": examples,
            }
        )
    except Exception as exc:
        out["error"] = str(exc)
    return out


def sanity_reference_checks(
    *,
    formal_pretrained: float | None,
    formal_finetuned: float | None,
    formal_improvement: float | None,
    training_start: float | None,
    training_best: float | None,
    training_improvement: float | None,
    best_checkpoint_changed: bool | None,
    overall_rc: int,
) -> dict:
    thresholds = SANITY_REFERENCE_THRESHOLDS
    checks = {
        "return_code_ok": overall_rc == 0,
        "formal_pretrained_val_dice_ok": (
            formal_pretrained is not None
            and formal_pretrained >= thresholds["formal_pretrained_val_dice_min"]
        ),
        "formal_finetuned_val_dice_ok": (
            formal_finetuned is not None
            and formal_finetuned >= thresholds["formal_finetuned_val_dice_min"]
        ),
        "formal_improvement_ok": (
            formal_improvement is not None
            and formal_improvement >= thresholds["formal_improvement_min"]
        ),
        "training_start_val_dice_ok": (
            training_start is not None
            and training_start >= thresholds["training_start_val_dice_min"]
        ),
        "training_best_val_dice_ok": (
            training_best is not None and training_best >= thresholds["training_best_val_dice_min"]
        ),
        "training_improvement_ok": (
            training_improvement is not None
            and training_improvement >= thresholds["training_improvement_min"]
        ),
        "best_checkpoint_changed_ok": best_checkpoint_changed is True,
    }
    failed = [name for name, ok in checks.items() if not ok]
    return {
        "thresholds": thresholds,
        "checks": checks,
        "failed_checks": failed,
        "passed": not failed,
    }


# --- CLI --------------------------------------------------------------------


@app.command()
def main(
    fixture: Optional[Path] = typer.Argument(
        None,
        help=(
            "Optional positional fixture path. The eval_engine harness calls "
            "the script as `python run_finetune.py <fixture>`; this argument "
            "lets the wrapper auto-pick the preset from the fixture's "
            "basename: `spleen_micro` -> --smoke, `Task06_Lung` -> --sanity, "
            "any other directory -> treated as --dataset-dir. Explicit flags "
            "(--smoke / --sanity / --dataset-dir) still win when given."
        ),
    ),
    dataset_dir: Optional[Path] = typer.Option(
        None, "--dataset-dir", help="Root containing imagesTr/ and labelsTr/."
    ),
    datalist: Optional[Path] = typer.Option(
        None,
        "--datalist",
        help="MONAI-bundle datalist JSON. Optional; auto-built when omitted.",
    ),
    target_anatomy: Optional[str] = typer.Option(
        None,
        "--target-anatomy",
        help="Anatomy name resolved against bundle/label_dict.json.",
    ),
    user_label_idx: int = typer.Option(
        1,
        "--user-label-idx",
        help="Label index that --target-anatomy occupies in the user's datalist.",
    ),
    label_mapping: Optional[str] = typer.Option(
        None,
        "--label-mapping",
        help="Literal `[[user_idx, vista3d_idx], ...]`. Overrides --target-anatomy.",
    ),
    epochs: Optional[int] = typer.Option(
        None,
        "--epochs",
        help="Override preset epochs (standard=50, softmax=100, sanity=5, smoke=2).",
    ),
    patch_size: Optional[str] = typer.Option(
        None, "--patch-size", help="JSON list. Overrides auto-derived patch size."
    ),
    cache_rate: Optional[float] = typer.Option(None, "--cache-rate"),
    learning_rate: Optional[float] = typer.Option(
        None,
        "--learning-rate",
        help="Standard/--sanity default: 5e-5; --softmax default: 1e-4.",
    ),
    output_dir: Path = typer.Option(
        Path("runs") / time.strftime("finetune_%Y%m%d_%H%M%S"), "--output-dir"
    ),
    smoke: bool = typer.Option(
        False,
        "--smoke",
        help="1 iter on bundled spleen_micro fixture (plumbing oracle).",
    ),
    sanity: bool = typer.Option(
        False,
        "--sanity",
        help="Tutorial-recipe verification on cached MSD06 Lung Tumor.",
    ),
    auto_seg: bool = typer.Option(
        False,
        "--auto-seg",
        help="Use automatic class-prompt training: drop_label_prob=0.0, drop_point_prob=1.0.",
    ),
    softmax: bool = typer.Option(
        False,
        "--softmax",
        help=(
            "Use the upstream fixed-channel softmax workflow for mutually exclusive "
            "semantic-segmentation labels."
        ),
    ),
    skip_formal_eval: bool = typer.Option(
        False,
        "--skip-formal-eval",
        help="Skip evaluate.json before/after scoring. Smoke always skips it.",
    ),
    mlflow_tracking_uri: Optional[str] = typer.Option(None, "--mlflow-tracking-uri"),
    mlflow_experiment_name: Optional[str] = typer.Option(
        None,
        "--mlflow-experiment-name",
        help="Enable MONAI-native MLflow tracking for this experiment.",
    ),
    mlflow_run_name: Optional[str] = typer.Option(None, "--mlflow-run-name"),
) -> None:
    """Auto-configure and run a documented VISTA3D finetuning workflow.

    \b
    Presets:
      --smoke   synthetic plumbing, 4 cases x 1 iter.
      --sanity  Real-recipe verification on MSD06 Lung Tumor - mirrors the
                published DFW tutorial: label mapping [[1, 23]], 5 epochs,
                lr=5e-5, patch [128,128,128], resample 1.5 mm isotropic,
                drop_label_prob=0.0, drop_point_prob=1.0, single GPU, and
                original-spacing evaluate.json scores before/after finetune.
      default   user dataset under --dataset-dir, lr=5e-5, 50 epochs.
      --softmax official fixed-channel training for mutually exclusive labels.

    The skill is built for "user brings their own dataset" (MSD layout:
    `imagesTr/` + `labelsTr/` with matching basenames). MSD06 lung tumor is
    the canonical sanity dataset.
    """
    t0 = time.perf_counter()
    if mlflow_experiment_name is None and (mlflow_tracking_uri or mlflow_run_name):
        raise typer.BadParameter(
            "--mlflow-experiment-name is required when other MLflow options are set."
        )
    output_dir = output_dir.resolve()
    output_dir.mkdir(parents=True, exist_ok=True)
    timings: dict[str, float] = {}
    t_phase = time.perf_counter()

    if softmax and auto_seg:
        raise typer.BadParameter(
            "--softmax cannot be combined with --auto-seg; softmax is a fixed-channel "
            "automatic-only architecture, not class-prompt training."
        )
    if softmax and sanity:
        raise typer.BadParameter(
            "--softmax cannot be combined with --sanity; the Task06 reference "
            "scores belong to the standard VISTA3D continual-learning workflow."
        )

    mlflow_env_keys = _MLFLOW_CHILD_ENV_KEYS if mlflow_experiment_name is not None else ()
    maybe_run_compatible_runtime(extra_env_keys=mlflow_env_keys)
    require_compatible_runtime()
    softmax_bundle_root = require_bundle_files(softmax=softmax)

    # Fixture-driven preset detection. Only consulted when no explicit
    # mode flag was passed, so callers retain full control. eval_engine's
    # default args template is [python, script, fixture], which hits this
    # path; humans tend to call the script with --smoke / --sanity directly.
    if fixture is not None and not smoke and not sanity and dataset_dir is None:
        fixture = fixture.resolve()
        preset = _fixture_preset(fixture)
        if preset == "smoke":
            smoke = True
        elif preset == "sanity":
            sanity = True
        elif fixture.is_dir():
            dataset_dir = fixture

    if softmax and sanity:
        raise typer.BadParameter(
            "--softmax cannot be combined with the Task06 sanity preset; use a "
            "caller-provided dataset and evaluate the softmax checkpoint separately."
        )

    # Preset selection - fill in dataset + defaults.
    smoke_generated_dataset = False
    if smoke:
        if fixture is not None and fixture.is_dir():
            dataset_dir = fixture.resolve()
        else:
            dataset_dir = SMOKE_FIXTURE
        datalist = dataset_dir / "datalist.json"
        dataset_dir, datalist, smoke_generated_dataset = ensure_smoke_dataset(
            dataset_dir, datalist, output_dir
        )
        target_anatomy = target_anatomy or "spleen"
    elif sanity:
        dataset_dir = _resolve_sanity_dataset(fixture, dataset_dir)
        if not dataset_dir.is_dir():
            raise typer.BadParameter(
                f"--sanity needs an MSD06 Lung Tumor dataset directory; tried {dataset_dir}\n"
                f"Pass the DFW/MSD Task06 path positionally, pass --dataset-dir, "
                f"or populate {SANITY_DATASET}."
            )
        target_anatomy = target_anatomy or SANITY_ANATOMY
        if epochs is None:
            epochs = int("5")
        if learning_rate is None:
            learning_rate = float("5e-5")
        if patch_size is None:
            patch_size = "[128,128,128]"
        if cache_rate is None:
            cache_rate = 1.0
        auto_seg = True
    if dataset_dir is None:
        raise typer.BadParameter("--dataset-dir required (or use --sanity / --smoke).")
    if not dataset_dir.is_dir():
        raise typer.BadParameter(f"dataset_dir does not exist: {dataset_dir}")
    if learning_rate is None:
        learning_rate = float("1e-4") if softmax else float("5e-5")

    env = detect_env()
    mapping, mapping_src = resolve_mapping(target_anatomy, user_label_idx, label_mapping)
    if softmax:
        validate_softmax_mapping(mapping)
    timings["env_detect"] = time.perf_counter() - t_phase
    t_phase = time.perf_counter()

    # Build or load the datalist. --sanity uses the same MSD-layout
    # auto-build as user datasets (the tutorial's seed-0 5-fold split is
    # what inspect_and_build_datalist already produces for MSD06 Lung Tumor).
    if datalist is None:
        datalist, dataset_audit = inspect_and_build_datalist(
            dataset_dir,
            output_dir,
            user_label_idx=user_label_idx,
            anatomy=target_anatomy,
        )
    else:
        if not datalist.is_file():
            raise typer.BadParameter(f"datalist not found: {datalist}")
        if smoke:
            dataset_audit = {
                "dataset_dir": str(dataset_dir),
                "datalist_source": "caller_provided",
                "datalist_path": str(datalist),
                "smoke_generated_dataset": smoke_generated_dataset,
            }
        else:
            dataset_audit = audit_existing_datalist(
                dataset_dir,
                datalist,
                user_label_idx=user_label_idx,
                anatomy=target_anatomy,
            )

    timings["dataset_audit"] = time.perf_counter() - t_phase
    t_phase = time.perf_counter()

    n_train = len(json.loads(datalist.read_text()).get("training", []))
    dataset_audit.setdefault("n_pairs", n_train)
    plan_patch = json.loads(patch_size) if patch_size else pick_patch(env["gpu_free_mb"])
    plan_cache = (
        cache_rate if cache_rate is not None else pick_cache_rate(n_train, env["host_ram_mb"])
    )
    plan_epochs = (
        epochs
        if epochs is not None
        else 2 if smoke else int("5") if sanity else int("100") if softmax else int("50")
    )
    workflow = "softmax" if softmax else "standard"
    formal_eval = bool(not softmax and not smoke and not skip_formal_eval)
    force_single_gpu = bool(smoke or sanity)
    nproc = 1 if force_single_gpu else pick_nproc(env["gpu_count"])
    multi_gpu = (not force_single_gpu) and nproc >= 2 and env["cuda_available"]
    ckpt_dir = output_dir / "checkpoints"
    train_output_dir = output_dir / "val_during_train"
    ckpt_dir.mkdir(parents=True, exist_ok=True)
    train_output_dir.mkdir(parents=True, exist_ok=True)

    plan = {
        "patch_size": plan_patch,
        "train_dataset_cache_rate": plan_cache,
        "epochs": plan_epochs,
        "learning_rate": learning_rate,
        "nproc_per_node": nproc,
        "multi_gpu": multi_gpu,
        "formal_eval": formal_eval,
        "auto_seg": auto_seg,
        "workflow": workflow,
        "force_single_gpu": force_single_gpu,
        "preset": "smoke" if smoke else "sanity" if sanity else "finetune",
        "rationale": [
            f"patch_size={plan_patch} (chose for free GPU={env['gpu_free_mb']} MiB)",
            f"cache_rate={plan_cache} (RAM={env['host_ram_mb']} MiB, n_train={n_train})",
            f"epochs={plan_epochs}",
            f"learning_rate={learning_rate}",
            f"nproc_per_node={nproc}",
            (
                "fixed-channel mutually exclusive softmax training"
                if softmax
                else "automatic class-prompt training" if auto_seg else "bundle prompt-mix training"
            ),
            (
                "single-gpu DFW Task06 recipe"
                if sanity
                else "single-gpu smoke preset" if smoke else "host GPU policy"
            ),
        ],
    }

    if softmax:
        if softmax_bundle_root is None:
            raise RuntimeError("softmax bundle resolution unexpectedly returned no path")
        override = build_softmax_override(
            dataset_dir,
            datalist,
            mapping,
            plan_patch,
            plan_cache,
            plan_epochs,
            learning_rate,
            ckpt_dir,
            train_output_dir,
        )
        generated_config_dir = output_dir / "configs"
        override_file = write_config_path(generated_config_dir / "softmax_override.json", override)
        no_logging_file = write_config_path(
            generated_config_dir / "softmax_no_logging.json",
            {"use_mlflow": False, "use_tensorboard": False},
        )
        train_stack = ["configs/train_continual_softmax.json"]
        if multi_gpu:
            train_stack.append("configs/multi_gpu_train.json")
        train_stack.extend([override_file, no_logging_file])
        eval_stack: list[str] = []
        run_bundle_root = softmax_bundle_root
    else:
        override = build_override(
            dataset_dir,
            datalist,
            mapping,
            plan_patch,
            plan_cache,
            plan_epochs,
            learning_rate,
            ckpt_dir,
            train_output_dir,
            auto_seg=auto_seg,
        )
        override_file = (
            write_config("train_continual_task06_lung.json", override)
            if sanity
            else write_config("auto_override.json", override)
        )
        no_logging_file = write_config(
            "dfw_no_logging.json",
            {"use_mlflow": False, "use_tensorboard": False},
        )
        metric_compat_files = metric_compat_config_stack()
        train_stack = ["configs/train.json", "configs/train_continual.json"]
        if multi_gpu:
            train_stack.append("configs/multi_gpu_train.json")
        train_stack.extend([override_file, no_logging_file, *metric_compat_files])
        eval_stack = [
            "configs/train.json",
            "configs/train_continual.json",
            "configs/evaluate.json",
            override_file,
            no_logging_file,
            *metric_compat_files,
        ]
        run_bundle_root = BUNDLE_DIR

    timings["plan"] = time.perf_counter() - t_phase
    t_phase = time.perf_counter()

    pretrained = BUNDLE_DIR / "models" / "model.pt"
    formal_pretrained = None
    formal_finetuned = None
    formal_pre_rc = None
    formal_post_rc = None
    formal_pre_cmd: list[str] | None = None
    formal_post_cmd: list[str] | None = None
    phase_peaks: dict[str, int] = {}
    mlflow_tracking: dict[str, str] | None = None

    if formal_eval:
        pre_eval_dir = output_dir / "eval_pretrained"
        pre_eval_dir.mkdir(parents=True, exist_ok=True)
        pre_log = output_dir / "eval_pretrained.log"
        formal_pre_rc, pre_peak, formal_pre_cmd = run_monai_bundle(
            eval_stack,
            pre_log,
            force_single_gpu=True,
            bundle_root=run_bundle_root,
            extra_args=[
                "--ckpt_path",
                str(pretrained),
                "--output_dir",
                str(pre_eval_dir),
            ],
        )
        phase_peaks["eval_pretrained"] = pre_peak
        timings["eval_pretrained"] = time.perf_counter() - t_phase
        formal_pretrained = read_val_mean_dice(pre_eval_dir)
        t_phase = time.perf_counter()

    train_extra_args: list[str] | None = None
    if mlflow_experiment_name is not None:
        train_extra_args, mlflow_tracking = build_mlflow_tracking_args(
            output_dir,
            tracking_uri=mlflow_tracking_uri,
            experiment_name=mlflow_experiment_name,
            requested_run_name=mlflow_run_name,
        )

    log_path = output_dir / "finetune.log"
    rc, train_peak, cmd = run_monai_bundle(
        train_stack,
        log_path,
        multi_gpu=multi_gpu,
        nproc=nproc,
        extra_args=train_extra_args,
        force_single_gpu=force_single_gpu,
        bundle_root=run_bundle_root,
        extra_env_keys=mlflow_env_keys,
    )
    phase_peaks["finetune"] = train_peak
    metrics = parse_log(log_path)
    timings["bundle_run"] = time.perf_counter() - t_phase

    finetune_ckpt = ckpt_dir / ("model_softmax.pt" if softmax else "model_finetune.pt")

    if formal_eval and rc == 0 and finetune_ckpt.exists():
        t_phase = time.perf_counter()
        post_eval_dir = output_dir / "eval_finetuned"
        post_eval_dir.mkdir(parents=True, exist_ok=True)
        post_log = output_dir / "eval_finetuned.log"
        formal_post_rc, post_peak, formal_post_cmd = run_monai_bundle(
            eval_stack,
            post_log,
            force_single_gpu=True,
            bundle_root=run_bundle_root,
            extra_args=[
                "--ckpt_path",
                str(finetune_ckpt),
                "--output_dir",
                str(post_eval_dir),
            ],
        )
        phase_peaks["eval_finetuned"] = post_peak
        timings["eval_finetuned"] = time.perf_counter() - t_phase
        formal_finetuned = read_val_mean_dice(post_eval_dir)

    checkpoint_comparisons = {
        "best": (
            compare_checkpoint_weights(pretrained, finetune_ckpt)
            if finetune_ckpt.exists()
            else None
        ),
    }
    best_checkpoint_changed = (
        checkpoint_comparisons["best"] is not None
        and checkpoint_comparisons["best"].get("weights_identical") is False
    )

    # Regression gate.
    baseline, best = metrics["baseline_val_dice"], metrics["best_val_dice"]
    formal_improvement = (
        round(formal_finetuned - formal_pretrained, int("4"))
        if formal_finetuned is not None and formal_pretrained is not None
        else None
    )
    formal_regressed = (
        formal_finetuned < formal_pretrained - float("1e-3")
        if formal_finetuned is not None and formal_pretrained is not None
        else None
    )
    formal_improved = (
        formal_finetuned > formal_pretrained + float("1e-3")
        if formal_finetuned is not None and formal_pretrained is not None
        else None
    )
    if baseline is None or best is None:
        regressed = improved = improvement = recommended = None
    else:
        improvement = round(best - baseline, int("4"))
        regressed = best < baseline - float("1e-3")
        improved = best > baseline + float("1e-3")
        recommended = (
            str(finetune_ckpt)
            if (improved and best_checkpoint_changed and finetune_ckpt.exists())
            else str(pretrained)
        )
    if formal_pretrained is not None:
        candidates: list[tuple[float, Path]] = []
        if (
            formal_finetuned is not None
            and formal_finetuned > formal_pretrained + float("1e-3")
            and best_checkpoint_changed
            and finetune_ckpt.exists()
        ):
            candidates.append((formal_finetuned, finetune_ckpt))
        recommended = str(max(candidates)[1]) if candidates else str(pretrained)
    if softmax:
        # The source checkpoint initializes the softmax network but is not
        # loadable by inference_softmax.json because the architectures differ.
        recommended = str(finetune_ckpt) if rc == 0 and finetune_ckpt.exists() else None

    peak_mb = max(phase_peaks.values()) if phase_peaks else 0
    phase_return_codes = {
        "eval_pretrained": formal_pre_rc,
        "finetune": rc,
        "eval_finetuned": formal_post_rc,
    }
    overall_rc = max((v for v in phase_return_codes.values() if v is not None), default=0)
    if sanity and formal_eval:
        sanity_checks = sanity_reference_checks(
            formal_pretrained=formal_pretrained,
            formal_finetuned=formal_finetuned,
            formal_improvement=formal_improvement,
            training_start=baseline,
            training_best=best,
            training_improvement=improvement,
            best_checkpoint_changed=best_checkpoint_changed,
            overall_rc=overall_rc,
        )
        sanity_ok = bool(sanity_checks["passed"])
    else:
        sanity_checks = None
        sanity_ok = (
            bool(baseline is not None and baseline >= float("0.5") and regressed is False)
            if sanity
            else None
        )

    result = {
        "skill": "nv_segment_ct_finetune",
        "model": "NVIDIA-Medtech/NV-Segment-CT (VISTA3D)",
        "model_repo": "https://huggingface.co/nvidia/NV-Segment-CT",
        "version": VERSION,
        "input": {
            "dataset_dir": str(dataset_dir),
            "datalist": str(datalist),
            "n_train_cases": n_train,
            "label_mappings": mapping,
            "label_mapping_resolution": mapping_src,
            "dataset_audit": dataset_audit,
            "smoke": smoke,
            "sanity": sanity,
            "auto_seg": auto_seg,
            "softmax": softmax,
            "workflow": workflow,
            "formal_eval": formal_eval,
        },
        "environment": env,
        "plan": plan,
        "invocation": {
            "command": " ".join(shlex.quote(c) for c in cmd),
            "commands": {
                "eval_pretrained": (
                    " ".join(shlex.quote(c) for c in formal_pre_cmd) if formal_pre_cmd else None
                ),
                "finetune": " ".join(shlex.quote(c) for c in cmd),
                "eval_finetuned": (
                    " ".join(shlex.quote(c) for c in formal_post_cmd) if formal_post_cmd else None
                ),
            },
            "command_prefix": (" ".join(cmd[: cmd.index("run") + 1]) if "run" in cmd else cmd[0]),
            "config_stack": train_stack,
            "eval_config_stack": eval_stack if formal_eval else None,
            "phase_return_codes": phase_return_codes,
            "multi_gpu": multi_gpu,
            "cwd": str(run_bundle_root),
            "bundle_root": str(run_bundle_root),
            "upstream_commit": (_git_commit(run_bundle_root) if softmax else UPSTREAM_CTMR_COMMIT),
            "override_file": override_file,
            "no_logging_file": no_logging_file,
            "mlflow_tracking": mlflow_tracking,
        },
        "output": {
            "finetuned_ckpt": str(finetune_ckpt) if finetune_ckpt.exists() else None,
            "finetuned_ckpt_exists": finetune_ckpt.exists(),
            "pretrained_ckpt": str(pretrained),
            "recommended_ckpt": recommended,
            "checkpoint_architecture": (
                "fixed_channel_softmax" if softmax else "vista3d_promptable"
            ),
            "checkpoint_comparisons_to_pretrained": checkpoint_comparisons,
            "finetuned_ckpt_matches_pretrained_weights": (
                checkpoint_comparisons["best"].get("weights_identical")
                if checkpoint_comparisons["best"] is not None
                else None
            ),
            "baseline_val_dice": baseline,
            "best_val_dice": best,
            "best_epoch_index": metrics["best_epoch_index"],
            "improvement_over_baseline": improvement,
            "regressed": regressed,
            "improved": improved,
            "training_start_val_dice": baseline,
            "training_best_val_dice": best,
            "training_best_epoch_index": metrics["best_epoch_index"],
            "formal_eval_enabled": formal_eval,
            "formal_pretrained_val_dice": formal_pretrained,
            "formal_finetuned_val_dice": formal_finetuned,
            "formal_improvement_over_pretrained": formal_improvement,
            "formal_regressed": formal_regressed,
            "formal_improved": formal_improved,
            "val_dice_per_epoch": metrics["val_dice_per_epoch"],
            "train_loss_first": metrics["train_loss_first"],
            "train_loss_last": metrics["train_loss_last"],
            "train_loss_finite": metrics["train_loss_finite"],
            "oom": metrics["oom"],
            "sanity_reference_checks": sanity_checks,
            "sanity_recovery_demonstrated": sanity_ok,
        },
        "runtime": {
            "wall_seconds": round(time.perf_counter() - t0, int("3")),
            "peak_gpu_mb": peak_mb,
            "phase_peak_gpu_mb": phase_peaks,
            "return_code": overall_rc,
            "log_path": str(log_path),
            "log_tail": metrics["log_tail"],
        },
        "cost": {
            "steps": [
                {
                    "step": "env_detect",
                    "label": "Step 2: detect GPU + RAM",
                    "seconds": round(timings.get("env_detect", 0.0), int("3")),
                },
                {
                    "step": "dataset_audit",
                    "label": "Step 0: audit inputs + build datalist",
                    "seconds": round(timings.get("dataset_audit", 0.0), int("3")),
                },
                {
                    "step": "plan",
                    "label": "Step 2: compose plan + override",
                    "seconds": round(timings.get("plan", 0.0), int("3")),
                },
                {
                    "step": "eval_pretrained",
                    "label": "Step 4a: evaluate pretrained checkpoint",
                    "seconds": round(timings.get("eval_pretrained", 0.0), int("3")),
                    "peak_gpu_mb": phase_peaks.get("eval_pretrained", 0),
                },
                {
                    "step": "bundle_run",
                    "label": "Step 5: monai.bundle run + log parse",
                    "seconds": round(timings.get("bundle_run", 0.0), int("3")),
                    "peak_gpu_mb": phase_peaks.get("finetune", 0),
                },
                {
                    "step": "eval_finetuned",
                    "label": "Step 6: evaluate fine-tuned checkpoint",
                    "seconds": round(timings.get("eval_finetuned", 0.0), int("3")),
                    "peak_gpu_mb": phase_peaks.get("eval_finetuned", 0),
                },
            ],
            "total_seconds": round(sum(timings.values()), int("3")),
        },
        "intended_use_disclaimer": (
            "Engineering verification only. Output is NOT clinically meaningful. "
            "This wrapper invokes the upstream `monai.bundle run` finetune entry "
            + (
                "described in NV-Segment-CT's softmax_finetune.md; it does not "
                "modify the upstream softmax implementation."
                if softmax
                else "described in NV-Segment-CT's finetune.md; it does not modify training."
            )
        ),
    }
    output_dir.mkdir(parents=True, exist_ok=True)
    (output_dir / "output.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))

    if overall_rc != 0 or metrics["oom"]:
        sys.exit(2)
    if sanity and not result["output"]["sanity_recovery_demonstrated"]:
        print(
            f"\n[SANITY FAIL] training_start={baseline} training_best={best} "
            f"formal_pretrained={formal_pretrained} "
            f"formal_finetuned={formal_finetuned} "
            f"best_checkpoint_changed={best_checkpoint_changed}. Need Task06 "
            f"formal eval recovery and a best checkpoint whose tensors differ "
            f"from the pretrained checkpoint.",
            file=sys.stderr,
        )
        sys.exit(2)


if __name__ == "__main__":
    app()