Dev (#85)

* Add CLI for five-fold VS training

* Improve vestibular schwannoma parallel training

Author: skaliy

research/vestibular_schwannoma/README.md

@@ -7,6 +7,7 @@ training, inference on new cases, and PACS deployment.
 ## Contents
 
 - `train_5fold.py`: command-line five-fold training and evaluation.
+- `merge_inference_manifests.py`: validate and combine parallel fold subsets for inference.
 - `notebooks/01_five_fold_cross_validation.ipynb`: train and compare UNet, DynUNet, and
   optional SegMamba models.
 - `notebooks/02_inference_new_cases.ipynb`: run one declared model or an explicit ensemble.
@@ -42,9 +43,62 @@ python train_5fold.py --models unet  # One model, all five folds
 python train_5fold.py --skip-unavailable
 ```
 
-The default requests four models across five folds for 500 epochs; run
-`python train_5fold.py --help` before starting. For interactive inspection and visualizations,
-start Jupyter from this directory or `notebooks/`:
+The default requests three models across five folds for 500 epochs; run
+`python train_5fold.py --help` before starting. A launcher processes its requested models
+and folds sequentially. With five GPUs, run one process per fold, assign each process one
+visible GPU, and give it a distinct, previously nonexistent `--results-root`:
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python train_5fold.py --models unet --folds 1 --results-root cv_results/unet_fold_1 &
+CUDA_VISIBLE_DEVICES=1 python train_5fold.py --models unet --folds 2 --results-root cv_results/unet_fold_2 &
+CUDA_VISIBLE_DEVICES=2 python train_5fold.py --models unet --folds 3 --results-root cv_results/unet_fold_3 &
+CUDA_VISIBLE_DEVICES=3 python train_5fold.py --models unet --folds 4 --results-root cv_results/unet_fold_4 &
+CUDA_VISIBLE_DEVICES=4 python train_5fold.py --models unet --folds 5 --results-root cv_results/unet_fold_5 &
+wait
+```
+
+Inside each process, its assigned physical GPU is exposed to PyTorch as CUDA device 0.
+
+Before starting parallel jobs, populate `preprocessed/` once with a single process;
+concurrent first-time cache creation is not supported. When `MLFLOW_TRACKING_URI` is unset,
+processes launched on the same machine from the same fastMONAI checkout automatically share
+fastMONAI's repository-root SQLite tracking store
+(`sqlite:////absolute/path/to/fastMONAI/mlruns.db`). For multiple machines or a central
+tracking service, configure the same remote URI in every shell:
+
+```bash
+export MLFLOW_TRACKING_URI=http://mlflow.example:5000
+```
+
+Do not merge run IDs from independent local MLflow databases: inference must be able to
+resolve every run ID through one tracking URI.
+
+Each launcher atomically updates `completed_run_ids.json` after every successful fold, so
+completed work remains mergeable if a later fold is interrupted. A subset job intentionally
+does not create `inference_run_ids.json`; its completed registry is also rejected by the
+inference loader. Combine disjoint completed-fold registries into a new inference-only
+results root. The merger requires exactly folds 1-5, verifies that dataset/splits,
+preprocessing, model/loss, and training settings match, and rejects missing or overlapping
+folds, duplicate MLflow run IDs, and an existing output root:
+
+```bash
+python merge_inference_manifests.py \
+  cv_results/unet_fold_1/completed_run_ids.json \
+  cv_results/unet_fold_2/completed_run_ids.json \
+  cv_results/unet_fold_3/completed_run_ids.json \
+  cv_results/unet_fold_4/completed_run_ids.json \
+  cv_results/unet_fold_5/completed_run_ids.json \
+  --model unet \
+  --output-root cv_results/unet_5fold_merged
+```
+
+To replace only fold 1, train it into a new results root and merge that new
+`completed_run_ids.json` with registries containing folds 2-5; omit the old fold-1
+registry. The replacement is accepted only when its training contract matches, and the
+merged `--output-root` must also be new.
+
+Use `cv_results/unet_5fold_merged/inference_run_ids.json` in notebook 02. For interactive
+inspection and visualizations, start Jupyter from this directory or `notebooks/`:
 
 ```bash
 jupyter lab notebooks/01_five_fold_cross_validation.ipynb
@@ -65,10 +119,15 @@ to `workflow/`. Training model configs contain the VS-specific architecture and
 model reconstruction, patch inference, metrics, and artifact formats remain fastMONAI
 responsibilities.
 
-Training retains weights-only `.pth` checkpoints for warm-starting or further fitting with a
-newly initialized optimizer and learning-rate schedule; they are not exact training-resume
-checkpoints. Inference and deployment use strict-loaded `.safetensors` artifacts. Generated
-data, results, tracking stores, checkpoints, and model bundles are excluded from Git.
+Training writes selected fold checkpoints below
+`<results-root>/<model>/fold_<n>/checkpoints/`, so different folds and models cannot overwrite
+each other. All-data learners are independently scoped below
+`<results-root>/<model>/all_data/`, but their final artifacts are stored in MLflow rather than as
+a local checkpoint. The fold `.pth` files support warm-starting or further fitting with a newly
+initialized optimizer and learning-rate schedule; they are not exact training-resume
+checkpoints. Final and best inference artifacts remain isolated in their MLflow runs. Inference
+and deployment use strict-loaded `.safetensors` artifacts. Generated data, results, tracking
+stores, checkpoints, and model bundles are excluded from Git.
 
 For container preparation and execution, see
 [deployment/pacs/README.md](deployment/pacs/README.md).

research/vestibular_schwannoma/merge_inference_manifests.py

@@ -0,0 +1,66 @@
+#!/usr/bin/env python3
+"""Merge parallel fold-training manifests into one inference selection."""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+
+PROJECT_ROOT = Path(__file__).resolve().parent
+RESEARCH_ROOT = PROJECT_ROOT.parent
+if str(RESEARCH_ROOT) not in sys.path:
+    sys.path.insert(0, str(RESEARCH_ROOT))
+
+from vestibular_schwannoma.workflow.run_selection import (  # noqa: E402
+    merge_fold_run_selections,
+)
+
+
+def _parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        description=(
+            "Validate and merge disjoint fold inference manifests produced by "
+            "independently launched training jobs."
+        )
+    )
+    parser.add_argument(
+        "selection_files",
+        nargs="+",
+        type=Path,
+        help=(
+            "Source completed_run_ids.json or inference_run_ids.json files to merge."
+        ),
+    )
+    parser.add_argument(
+        "--model",
+        required=True,
+        help="Model key whose best fold runs should be merged, for example unet.",
+    )
+    parser.add_argument(
+        "--output-root",
+        type=Path,
+        required=True,
+        help="New directory in which to write the merged inference_run_ids.json.",
+    )
+    return parser
+
+
+def _project_path(path: Path) -> Path:
+    return path if path.is_absolute() else PROJECT_ROOT / path
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = _parser().parse_args(argv)
+    destination = merge_fold_run_selections(
+        [_project_path(path) for path in args.selection_files],
+        model_key=args.model,
+        output_root=_project_path(args.output_root),
+    )
+    print(f"Merged inference run selection: {destination}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

research/vestibular_schwannoma/notebooks/01_five_fold_cross_validation.ipynb

@@ -83,7 +83,9 @@
     "\n",
     "Patch-loading settings depend on available hardware. More queue workers may improve training throughput, while a larger queue increases RAM usage.\n",
     "\n",
-    "`training_seed` initializes training randomness before each independent run while retaining cuDNN performance optimizations. In an all-data run, every case remains in training and the first case by stable `case_id` order is duplicated only for fastai's validation phase. Its metric is an internal monitor, not held-out evaluation."
+    "`training_seed` initializes training randomness before each independent run while retaining cuDNN performance optimizations. In an all-data run, every case remains in training and the first case by stable `case_id` order is duplicated only for fastai's validation phase. Its metric is an internal monitor, not held-out evaluation.\n",
+    "\n",
+    "Every independently launched job must use a distinct, previously nonexistent `RESULTS_ROOT`. Re-run this configuration cell before starting another sweep. After parallel fold subsets finish (or one is interrupted after completing some folds), use `merge_inference_manifests.py` to validate and combine their `completed_run_ids.json` registries. The merger rejects different dataset, split, preprocessing, model, loss, or training contracts. A partial `completed_run_ids.json` registry cannot be used directly for inference."
    ]
   },
   {
@@ -105,13 +107,13 @@
     "    use_tta=True,\n",
     "    compile_models=True,\n",
     "    target_spacing=(0.4102, 0.4102, 1.5),\n",
-    "    patch_size=(192, 192, 48),\n",
+    "    patch_size=(256, 256, 48),\n",
     "    queue_num_workers=4,\n",
     "    queue_length=300,\n",
     ")\n",
     "\n",
     "DATA_CSV = \"data/ml_dataset.csv\"\n",
-    "RESULTS_RUN = datetime.now(timezone.utc).strftime(\"%Y%m%dT%H%M%SZ\")\n",
+    "RESULTS_RUN = datetime.now(timezone.utc).strftime(\"%Y%m%dT%H%M%S%fZ\")\n",
     "RESULTS_ROOT = Path(\"cv_results\") / RESULTS_RUN\n",
     "\n",
     "torch.backends.cudnn.benchmark = True\n",
@@ -172,7 +174,9 @@
    "source": [
     "## 4. Preprocess once to disk\n",
     "\n",
-    "Preprocessing is fold-independent. `preprocess_dataset()` automatically creates a versioned preprocessing cache and a `preprocessing_manifest.json` file. On later runs, fastMONAI uses the manifest to verify that the source files and preprocessing settings are unchanged before reusing the cache. `PatchConfig(preprocessed=True)` prevents preprocessing from being applied twice during training.\n"
+    "Preprocessing is fold-independent. `preprocess_dataset()` automatically creates a versioned preprocessing cache and a `preprocessing_manifest.json` file. On later runs, fastMONAI uses the manifest to verify that the source files and preprocessing settings are unchanged before reusing the cache. `PatchConfig(preprocessed=True)` prevents preprocessing from being applied twice during training.\n",
+    "\n",
+    "Populate a new preprocessing cache with one process before launching parallel training jobs; concurrent first-time cache creation is not supported.\n"
    ]
   },
   {
@@ -543,7 +547,8 @@
     "- `cv_results/<RESULTS_RUN>/<model>/fold_N/`: metrics and predictions.\n",
     "- `cv_results/<RESULTS_RUN>/<model>/cv_summary.csv`: one complete model summary.\n",
     "- `cv_results/<RESULTS_RUN>/cv_model_comparison.csv`: cross-model summary.\n",
-    "- `cv_results/<RESULTS_RUN>/inference_run_ids.json`: exact completed MLflow runs for notebook 02.\n",
+    "- `cv_results/<RESULTS_RUN>/completed_run_ids.json`: atomically updated MLflow runs for completed folds, including before a later interruption.\n",
+    "- `cv_results/<RESULTS_RUN>/inference_run_ids.json`: exact fully completed MLflow runs for notebook 02.\n",
     "\n",
     "MLflow retains final and selected weights-only checkpoints plus strict-loadable Safetensors artifacts. All-data runs produce only final artifacts.\n",
     "\n",

research/vestibular_schwannoma/tests/workflow/test_config.py

@@ -16,7 +16,7 @@ def test_defaults_preserve_the_notebook_experiment(self):
         self.assertEqual(config.model_keys, ("unet", "dynunet", "segmamba"))
         self.assertEqual(config.folds, (1, 2, 3, 4, 5))
         self.assertEqual(config.target_spacing, (0.4102, 0.4102, 1.5))
-        self.assertEqual(config.patch_size, (192, 192, 48))
+        self.assertEqual(config.patch_size, (256, 256, 48))
         self.assertEqual(config.epochs, 500)
         self.assertEqual(config.batch_size, 4)
         self.assertEqual(config.training_seed, 42)
@@ -44,7 +44,7 @@ def test_invalid_declarations_fail_early(self):
             {"model_keys": ("unet",), "epochs": 0},
             {"model_keys": ("unet",), "training_seed": -1},
             {"model_keys": ("unet",), "training_seed": True},
-            {"model_keys": ("unet",), "patch_size": (192, 192, 0)},
+            {"model_keys": ("unet",), "patch_size": (256, 256, 0)},
             {
                 "model_keys": ("unet",),
                 "foreground_sampling_probability": 0,
@@ -62,7 +62,7 @@ def test_patch_factory_preserves_training_and_inference_contract(self):
         config = ExperimentConfig(model_keys=("unet",))
         normalization = [ZNormalization(masking_method="foreground")]
         patch = make_patch_config(config, normalization)
-        self.assertEqual(patch.patch_size, [192, 192, 48])
+        self.assertEqual(patch.patch_size, [256, 256, 48])
         self.assertEqual(patch.target_spacing, [0.4102, 0.4102, 1.5])
         self.assertEqual(patch.label_probabilities, {0: 0.2, 1: 0.8})
         self.assertTrue(patch.preprocessed)