diff --git a/research/vestibular_schwannoma/README.md b/research/vestibular_schwannoma/README.md index c108e79..16db028 100644 --- a/research/vestibular_schwannoma/README.md +++ b/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 +`//fold_/checkpoints/`, so different folds and models cannot overwrite +each other. All-data learners are independently scoped below +`//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). diff --git a/research/vestibular_schwannoma/merge_inference_manifests.py b/research/vestibular_schwannoma/merge_inference_manifests.py new file mode 100644 index 0000000..2688b2b --- /dev/null +++ b/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()) diff --git a/research/vestibular_schwannoma/notebooks/01_five_fold_cross_validation.ipynb b/research/vestibular_schwannoma/notebooks/01_five_fold_cross_validation.ipynb index b9d54b5..d5e43ba 100644 --- a/research/vestibular_schwannoma/notebooks/01_five_fold_cross_validation.ipynb +++ b/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///fold_N/`: metrics and predictions.\n", "- `cv_results///cv_summary.csv`: one complete model summary.\n", "- `cv_results//cv_model_comparison.csv`: cross-model summary.\n", - "- `cv_results//inference_run_ids.json`: exact completed MLflow runs for notebook 02.\n", + "- `cv_results//completed_run_ids.json`: atomically updated MLflow runs for completed folds, including before a later interruption.\n", + "- `cv_results//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", diff --git a/research/vestibular_schwannoma/tests/workflow/test_config.py b/research/vestibular_schwannoma/tests/workflow/test_config.py index d07499e..5374888 100644 --- a/research/vestibular_schwannoma/tests/workflow/test_config.py +++ b/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) diff --git a/research/vestibular_schwannoma/tests/workflow/test_inference.py b/research/vestibular_schwannoma/tests/workflow/test_inference.py index 502b6bd..e4e7fb8 100644 --- a/research/vestibular_schwannoma/tests/workflow/test_inference.py +++ b/research/vestibular_schwannoma/tests/workflow/test_inference.py @@ -12,6 +12,10 @@ from vestibular_schwannoma.workflow import inference from vestibular_schwannoma.workflow.config import VS_OUTPUT_SPEC +from vestibular_schwannoma.workflow.run_selection import ( + make_training_contract, + read_inference_run_ids, +) class InferenceArtifactTests(unittest.TestCase): @@ -44,17 +48,224 @@ def metadata( }, } - def write_selection(self, path, *, model_key="unet", role="best", run_ids=None): + def write_selection( + self, + path, + *, + model_key="unet", + role="best", + run_ids=None, + contract_payload=None, + manifest_kind="inference_selection", + include_contract=True, + ): manifest = { "schema_version": 1, + "manifest_kind": manifest_kind, "run_group": "test-group", "models": { - model_key: {role: run_ids or {"fold_1": "run-1"}}, + model_key: { + role: run_ids if run_ids is not None else {"fold_1": "run-1"} + }, }, } + if include_contract: + manifest["training_contracts"] = { + model_key: make_training_contract( + contract_payload or {"campaign": "test"} + ) + } path.write_text(json.dumps(manifest), encoding="utf-8") return path + def test_merges_disjoint_fold_run_selections(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + first = self.write_selection( + root / "folds_123.json", + manifest_kind="completed_registry", + run_ids={ + "fold_1": "run-1", + "fold_2": "run-2", + "fold_3": "run-3", + }, + ) + second = self.write_selection( + root / "folds_45.json", + manifest_kind="completed_registry", + run_ids={"fold_4": "run-4", "fold_5": "run-5"}, + ) + destination = inference.merge_fold_run_selections( + [first, second], + model_key="unet", + output_root=root / "merged", + ) + manifest = json.loads(destination.read_text(encoding="utf-8")) + + self.assertEqual(destination.name, "inference_run_ids.json") + self.assertEqual(manifest["manifest_kind"], "inference_selection") + self.assertEqual(manifest["run_group"], "merged") + self.assertEqual( + manifest["training_contracts"]["unet"]["payload"], + {"campaign": "test"}, + ) + self.assertEqual( + manifest["models"]["unet"]["best"], + { + "fold_1": "run-1", + "fold_2": "run-2", + "fold_3": "run-3", + "fold_4": "run-4", + "fold_5": "run-5", + }, + ) + + def test_completed_registry_cannot_be_loaded_directly_for_inference(self): + with tempfile.TemporaryDirectory() as directory: + selection = self.write_selection( + Path(directory) / "completed_run_ids.json", + manifest_kind="completed_registry", + ) + with self.assertRaisesRegex(ValueError, "partial registry"): + inference.load_inference_models( + run_selection_file=selection, + model_key="unet", + artifact_role="best", + device="cpu", + ) + + def test_legacy_schema_one_selection_without_contract_remains_readable(self): + with tempfile.TemporaryDirectory() as directory: + selection = self.write_selection( + Path(directory) / "legacy.json", + manifest_kind="inference_selection", + include_contract=False, + ) + manifest = json.loads(selection.read_text()) + manifest.pop("manifest_kind") + selection.write_text(json.dumps(manifest), encoding="utf-8") + + self.assertEqual( + read_inference_run_ids( + selection, + model_key="unet", + artifact_role="best", + ), + {"fold_1": "run-1"}, + ) + + def test_reader_rejects_invalid_or_duplicate_run_ids(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + empty = self.write_selection( + root / "empty.json", + run_ids={"fold_1": ""}, + ) + duplicate = self.write_selection( + root / "duplicate.json", + run_ids={"fold_1": "same-run", "fold_2": "same-run"}, + ) + with self.assertRaisesRegex(ValueError, "non-empty string"): + read_inference_run_ids( + empty, + model_key="unet", + artifact_role="best", + ) + with self.assertRaisesRegex(ValueError, "duplicate run IDs"): + read_inference_run_ids( + duplicate, + model_key="unet", + artifact_role="best", + ) + + def test_merge_rejects_training_contract_mismatch(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + first = self.write_selection( + root / "first.json", + run_ids={"fold_1": "run-1"}, + ) + second = self.write_selection( + root / "second.json", + run_ids={"fold_2": "run-2"}, + contract_payload={"campaign": "different"}, + ) + with self.assertRaisesRegex(ValueError, "Training contract mismatch"): + inference.merge_fold_run_selections( + [first, second], + model_key="unet", + output_root=root / "mismatch-output", + ) + self.assertFalse((root / "mismatch-output").exists()) + + def test_merge_rejects_overlapping_or_incomplete_fold_selections(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + first = self.write_selection( + root / "first.json", + run_ids={"fold_1": "run-1"}, + ) + overlap = self.write_selection( + root / "overlap.json", + run_ids={"fold_1": "other-run"}, + ) + with self.assertRaisesRegex(ValueError, "Duplicate model member 'fold_1'"): + inference.merge_fold_run_selections( + [first, overlap], + model_key="unet", + output_root=root / "overlap-output", + ) + with self.assertRaisesRegex(ValueError, r"missing=\['fold_2'"): + inference.merge_fold_run_selections( + [first], + model_key="unet", + output_root=root / "incomplete-output", + ) + + self.assertFalse((root / "overlap-output").exists()) + self.assertFalse((root / "incomplete-output").exists()) + + def test_merge_rejects_duplicate_run_ids_and_existing_output_root(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + first = self.write_selection( + root / "first.json", + run_ids={"fold_1": "same-run"}, + ) + second = self.write_selection( + root / "second.json", + run_ids={ + "fold_2": "same-run", + "fold_3": "run-3", + "fold_4": "run-4", + "fold_5": "run-5", + }, + ) + with self.assertRaisesRegex(ValueError, "duplicate MLflow run IDs"): + inference.merge_fold_run_selections( + [first, second], + model_key="unet", + output_root=root / "duplicate-output", + ) + + existing = root / "existing" + existing.mkdir() + self.write_selection( + second, + run_ids={ + "fold_2": "different-run", + "fold_3": "run-3", + "fold_4": "run-4", + "fold_5": "run-5", + }, + ) + with self.assertRaisesRegex(FileExistsError, "new --output-root"): + inference.merge_fold_run_selections( + [first, second], + model_key="unet", + output_root=existing, + ) + def test_loader_requires_exactly_one_source(self): for selection, local in ((None, {}), ("selection.json", {"a": "model"})): with ( @@ -263,7 +474,6 @@ def test_rejects_incompatible_ensemble_before_loading_models(self): ) load.assert_not_called() - def test_rejects_artifacts_that_remove_disconnected_predictions(self): with tempfile.TemporaryDirectory() as directory: path = Path(directory) / "fold_1.safetensors" @@ -272,9 +482,7 @@ def test_rejects_artifacts_that_remove_disconnected_predictions(self): patch.object( inference, "read_safetensors_metadata", - return_value=self.metadata( - "run-1", keep_largest_component=True - ), + return_value=self.metadata("run-1", keep_largest_component=True), ), patch.object(inference, "load_safetensors_model") as load, self.assertRaisesRegex( diff --git a/research/vestibular_schwannoma/tests/workflow/test_merge_inference_manifests.py b/research/vestibular_schwannoma/tests/workflow/test_merge_inference_manifests.py new file mode 100644 index 0000000..8037cf1 --- /dev/null +++ b/research/vestibular_schwannoma/tests/workflow/test_merge_inference_manifests.py @@ -0,0 +1,64 @@ +import tempfile +import unittest +from contextlib import redirect_stderr, redirect_stdout +from io import StringIO +from pathlib import Path +from unittest.mock import patch + +from vestibular_schwannoma import merge_inference_manifests + + +class MergeInferenceManifestsCliTests(unittest.TestCase): + def test_main_passes_resolved_inputs_to_fixed_fold_merger(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + first = root / "folds_123" / "inference_run_ids.json" + second = root / "folds_45" / "inference_run_ids.json" + output_root = root / "merged" + destination = output_root / "inference_run_ids.json" + with ( + patch.object( + merge_inference_manifests, + "merge_fold_run_selections", + return_value=destination, + ) as merge, + redirect_stdout(StringIO()), + ): + status = merge_inference_manifests.main( + [ + str(first), + str(second), + "--model", + "unet", + "--output-root", + str(output_root), + ] + ) + + self.assertEqual(status, 0) + merge.assert_called_once_with( + [first, second], + model_key="unet", + output_root=output_root, + ) + + def test_custom_fold_set_is_rejected(self): + with ( + redirect_stderr(StringIO()), + self.assertRaises(SystemExit), + ): + merge_inference_manifests._parser().parse_args( + [ + "fold_1/completed_run_ids.json", + "--model", + "unet", + "--folds", + "1", + "--output-root", + "merged", + ] + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/research/vestibular_schwannoma/tests/workflow/test_models.py b/research/vestibular_schwannoma/tests/workflow/test_models.py index 9a1d6e7..3c6d04d 100644 --- a/research/vestibular_schwannoma/tests/workflow/test_models.py +++ b/research/vestibular_schwannoma/tests/workflow/test_models.py @@ -8,17 +8,48 @@ class TrainingModelConfigTests(unittest.TestCase): def test_specs_are_the_single_architecture_declaration(self): self.assertEqual(models.UNET_SPEC["arch_id"], "monai.unet") self.assertEqual(models.DYNUNET_SPEC["arch_id"], "monai.dynunet") - self.assertEqual(models.DYNUNET_SMALL_SPEC["arch_id"], "monai.dynunet") self.assertEqual( models.DYNUNET_SPEC["wrapper_spec"][0]["wrapper_id"], "fastmonai.dynunet_ds_adapter", ) self.assertEqual( - models.DYNUNET_SMALL_SPEC["arch_kwargs"]["filters"], - [32, 64, 128, 256, 512], + models.UNET_SPEC["arch_kwargs"]["channels"], + [32, 64, 128, 256, 320], + ) + self.assertEqual( + models.DYNUNET_SPEC["arch_kwargs"]["filters"], + [32, 64, 128, 256, 320], ) self.assertEqual(models.SEGMAMBA_SPEC["arch_id"], "segmamba.v2") + def test_registry_contains_only_supported_architectures(self): + self.assertEqual( + set(models.TRAINING_MODEL_CONFIGS), + {"unet", "dynunet", "segmamba"}, + ) + + def test_loss_specs_declare_scientifically_relevant_parameters(self): + self.assertEqual( + models.TRAINING_MODEL_CONFIGS["unet"].loss_spec, + { + "loss_id": "monai.dice_ce", + "kwargs": { + "to_onehot_y": True, + "softmax": True, + "include_background": False, + "batch": True, + }, + }, + ) + self.assertEqual( + models.TRAINING_MODEL_CONFIGS["dynunet"].loss_spec, + { + "loss_id": "monai.deep_supervision", + "kwargs": {"weight_mode": "exp"}, + "base_loss": models.DICE_CE_LOSS_SPEC, + }, + ) + def test_declared_order_is_preserved(self): configs = models.get_training_model_configs(("dynunet", "unet")) self.assertEqual(list(configs), ["dynunet", "unet"]) @@ -39,9 +70,7 @@ def test_missing_optional_segmamba_can_skip_or_fail(self): configs = models.get_training_model_configs(("unet", "segmamba")) self.assertEqual(list(configs), ["unet"]) with self.assertRaisesRegex(ImportError, "SegMamba was requested"): - models.get_training_model_configs( - ("segmamba",), skip_unavailable=False - ) + models.get_training_model_configs(("segmamba",), skip_unavailable=False) def test_compilation_is_applied_after_spec_construction(self): sentinel_model = object() @@ -51,7 +80,9 @@ def test_compilation_is_applied_after_spec_construction(self): patch.object( models, "build_model_from_spec", return_value=sentinel_model ) as build, - patch.object(models.torch, "compile", return_value=compiled_model) as compile, + patch.object( + models.torch, "compile", return_value=compiled_model + ) as compile, ): result = models.build_training_model(config, compile_model=True) build.assert_called_once_with(config.model_spec) @@ -62,9 +93,7 @@ def test_segmamba_preserves_the_uncompiled_training_path(self): sentinel_model = object() config = models.TRAINING_MODEL_CONFIGS["segmamba"] with ( - patch.object( - models, "build_model_from_spec", return_value=sentinel_model - ), + patch.object(models, "build_model_from_spec", return_value=sentinel_model), patch.object(models.torch, "compile") as compile, ): result = models.build_training_model(config, compile_model=True) diff --git a/research/vestibular_schwannoma/tests/workflow/test_train_5fold.py b/research/vestibular_schwannoma/tests/workflow/test_train_5fold.py index 7ed2edf..8f3ba80 100644 --- a/research/vestibular_schwannoma/tests/workflow/test_train_5fold.py +++ b/research/vestibular_schwannoma/tests/workflow/test_train_5fold.py @@ -6,6 +6,7 @@ class FiveFoldLauncherTests(unittest.TestCase): def test_defaults_select_all_models_and_preserve_80_20_control(self): args = train_5fold._parser().parse_args([]) + self.assertEqual(train_5fold.DEFAULT_MODELS, ("unet", "dynunet", "segmamba")) self.assertEqual(tuple(args.models), train_5fold.DEFAULT_MODELS) self.assertEqual(args.folds, [1, 2, 3, 4, 5]) @@ -17,11 +18,6 @@ def test_70_30_sampling_remains_an_explicit_comparison(self): self.assertEqual(args.foreground_probability, 0.7) - def test_small_dynunet_uses_the_intermediate_512_bottleneck(self): - spec = train_5fold.TRAINING_MODEL_CONFIGS["dynunet_small"].model_spec - - self.assertEqual(spec["arch_kwargs"]["filters"], [32, 64, 128, 256, 512]) - if __name__ == "__main__": unittest.main() diff --git a/research/vestibular_schwannoma/tests/workflow/test_training.py b/research/vestibular_schwannoma/tests/workflow/test_training.py index c58212b..499aaa8 100644 --- a/research/vestibular_schwannoma/tests/workflow/test_training.py +++ b/research/vestibular_schwannoma/tests/workflow/test_training.py @@ -1,6 +1,7 @@ import json import tempfile import unittest +from dataclasses import replace from contextlib import redirect_stderr, redirect_stdout from io import StringIO from pathlib import Path @@ -69,7 +70,9 @@ def test_fold_evaluation_reads_normalization_from_patch_config(self): ) with tempfile.TemporaryDirectory() as directory: with ( - patch.object(training, "patch_inference", return_value=[object()]) as infer, + patch.object( + training, "patch_inference", return_value=[object()] + ) as infer, patch.object(training, "evaluate_segmentations", return_value=metrics), patch.object( training, @@ -142,13 +145,16 @@ def test_fold_dataloaders_close_when_model_construction_fails(self): dls = MagicMock() dls.train.subjects_dataset = [object()] dls.valid.subjects_dataset = [object()] + results_directory = tempfile.TemporaryDirectory() + self.addCleanup(results_directory.cleanup) + results_dir = Path(results_directory.name) / "unet" with ( - patch.object( - training.MedPatchDataLoaders, "from_df", return_value=dls - ), + patch.object(training.MedPatchDataLoaders, "from_df", return_value=dls), patch.object(training, "make_gpu_augmentation", return_value=object()), patch.object( - training, "build_training_model", side_effect=RuntimeError("build failed") + training, + "build_training_model", + side_effect=RuntimeError("build failed"), ), redirect_stdout(StringIO()), ): @@ -160,10 +166,34 @@ def test_fold_dataloaders_close_when_model_construction_fails(self): experiment=experiment, patch_config=SimpleNamespace(), preprocessing_manifest="manifest.json", - results_dir="results/unet", + results_dir=results_dir, ) dls.close.assert_called_once_with() + def test_fold_training_rejects_an_existing_fold_directory(self): + experiment = ExperimentConfig( + model_keys=("unet",), + folds=(1,), + epochs=1, + compile_models=False, + ) + with tempfile.TemporaryDirectory() as directory: + results_dir = Path(directory) / "unet" + (results_dir / "fold_1").mkdir(parents=True) + with ( + patch.object(training, "_set_training_seed"), + self.assertRaisesRegex(FileExistsError, "fold_1"), + ): + training.train_one_fold( + TRAINING_MODEL_CONFIGS["unet"], + 1, + pd.DataFrame(), + experiment=experiment, + patch_config=SimpleNamespace(), + preprocessing_manifest="manifest.json", + results_dir=results_dir, + ) + def test_fold_tracking_uses_fixed_model_and_output_contract(self): experiment = ExperimentConfig( model_keys=("unet",), @@ -191,6 +221,10 @@ def test_fold_tracking_uses_fixed_model_and_output_contract(self): flip={"axes": (0, 1, 2), "p": 0.5}, ) + results_directory = tempfile.TemporaryDirectory() + self.addCleanup(results_directory.cleanup) + results_dir = Path(results_directory.name) / "run-group" / "unet" + with ( patch.object(training, "_set_training_seed") as set_training_seed, patch.object(training.MedPatchDataLoaders, "from_df", return_value=dls), @@ -209,17 +243,21 @@ def test_fold_tracking_uses_fixed_model_and_output_contract(self): redirect_stdout(StringIO()), ): learner_factory.return_value.to_bf16.return_value = learner - training.train_one_fold( + fold_run = training.train_one_fold( model_config, 1, train_df, experiment=experiment, patch_config=SimpleNamespace(), preprocessing_manifest="manifest.json", - results_dir="results/unet", + results_dir=results_dir, ) set_training_seed.assert_called_once_with(42) + learner_kwargs = learner_factory.call_args.kwargs + self.assertEqual(learner_kwargs["path"], results_dir / "fold_1") + self.assertEqual(learner_kwargs["model_dir"], "checkpoints") + self.assertEqual(fold_run.results_dir, results_dir / "fold_1") self.assertEqual( create_callback.call_args.kwargs["experiment_name"], "vestibular_schwannoma_unet", @@ -227,9 +265,7 @@ def test_fold_tracking_uses_fixed_model_and_output_contract(self): self.assertEqual( create_callback.call_args.kwargs["model_spec"], model_config.model_spec ) - self.assertIs( - create_callback.call_args.kwargs["output_spec"], VS_OUTPUT_SPEC - ) + self.assertIs(create_callback.call_args.kwargs["output_spec"], VS_OUTPUT_SPEC) extra_params = create_callback.call_args.kwargs["extra_params"] self.assertEqual(extra_params["training_seed"], 42) self.assertEqual(extra_params["foreground_sampling_probability"], 0.8) @@ -274,6 +310,9 @@ def test_all_data_training_uses_one_stable_duplicated_monitor_case(self): flip={"axes": (0, 1, 2), "p": 0.5}, ) events = [] + results_directory = tempfile.TemporaryDirectory() + self.addCleanup(results_directory.cleanup) + results_dir = Path(results_directory.name) / "run-group" / "unet" def seed_run(seed): events.append(("seed", seed)) @@ -312,9 +351,12 @@ def build_model(*args, **kwargs): experiment=experiment, patch_config=SimpleNamespace(), preprocessing_manifest="manifest.json", - run_group="run-group", + results_dir=results_dir, ) + learner_kwargs = learner_factory.call_args.kwargs + self.assertEqual(learner_kwargs["path"], results_dir / "all_data") + self.assertEqual(learner_kwargs["model_dir"], "checkpoints") fit_df = from_df.call_args.kwargs["df"] training_rows = fit_df.loc[~fit_df["is_val"]] monitor_rows = fit_df.loc[fit_df["is_val"]] @@ -371,6 +413,8 @@ def fake_train(model_config, fold, *args, **kwargs): ) with tempfile.TemporaryDirectory() as directory: + preprocessing_manifest = Path(directory) / "preprocessing_manifest.json" + preprocessing_manifest.write_text("{}", encoding="utf-8") with ( patch.object(training, "train_one_fold", side_effect=fake_train), redirect_stdout(StringIO()), @@ -381,17 +425,181 @@ def fake_train(model_config, fold, *args, **kwargs): train_df, experiment=experiment, patch_config=SimpleNamespace(), - preprocessing_manifest="manifest.json", - results_root=directory, + preprocessing_manifest=preprocessing_manifest, + results_root=Path(directory) / "run", ) - selection = json.loads(sweep.inference_run_ids_path.read_text()) + inference_selection_exists = ( + Path(directory) / "run" / "inference_run_ids.json" + ).exists() + completed = json.loads(sweep.completed_run_ids_path.read_text()) self.assertEqual(list(sweep.fold_runs["unet"]), [2]) self.assertEqual(len(sweep.failures), 1) self.assertEqual(sweep.failures[0].fold, 1) self.assertEqual(sweep.failures[0].error_type, "RuntimeError") - self.assertEqual(selection["models"], {}) + self.assertIsNone(sweep.inference_run_ids_path) + self.assertFalse(inference_selection_exists) + self.assertEqual(completed["manifest_kind"], "completed_registry") + self.assertEqual( + completed["models"]["unet"]["best"], + {"fold_2": "run-2"}, + ) + self.assertIn("unet", completed["training_contracts"]) - def test_sweep_writes_complete_fold_and_all_data_run_selection(self): + def test_completed_fold_manifest_survives_an_interrupted_subset(self): + experiment = ExperimentConfig( + model_keys=("unet",), + folds=(1, 2), + epochs=1, + compile_models=False, + ) + train_df = pd.DataFrame({"case_id": ["a", "b"], "fold": [1, 2]}) + + def fake_train(model_config, fold, *args, **kwargs): + if fold == 2: + raise KeyboardInterrupt + return training.FoldRun( + model_key=model_config.key, + fold=fold, + run_id="run-1", + results_dir=Path("fold_1"), + results=pd.DataFrame({"case_id": ["a"], "dsc": [0.8]}), + ) + + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + preprocessing_manifest = root / "preprocessing_manifest.json" + preprocessing_manifest.write_text("{}", encoding="utf-8") + results_root = root / "run" + with ( + patch.object(training, "train_one_fold", side_effect=fake_train), + redirect_stdout(StringIO()), + self.assertRaises(KeyboardInterrupt), + ): + training.run_training_sweep( + {"unet": TRAINING_MODEL_CONFIGS["unet"]}, + train_df, + experiment=experiment, + patch_config=SimpleNamespace(), + preprocessing_manifest=preprocessing_manifest, + results_root=results_root, + ) + completed = json.loads( + (results_root / "completed_run_ids.json").read_text() + ) + + self.assertEqual( + completed["models"]["unet"]["best"], + {"fold_1": "run-1"}, + ) + self.assertEqual(completed["manifest_kind"], "completed_registry") + self.assertFalse((results_root / "inference_run_ids.json").exists()) + + def test_training_contract_covers_scientific_settings_but_not_fold_subset(self): + train_df = pd.DataFrame({"case_id": ["a", "b", "c"], "fold": [1, 2, 3]}) + patch_config = SimpleNamespace(patch_overlap=0.5) + model_config = TRAINING_MODEL_CONFIGS["unet"] + + with tempfile.TemporaryDirectory() as directory: + preprocessing_manifest = Path(directory) / "preprocessing_manifest.json" + preprocessing_manifest.write_text( + '{"dataset_version": "test"}', + encoding="utf-8", + ) + + def contract( + folds, + *, + dataframe=train_df, + patch=patch_config, + model=model_config, + epochs=1, + ): + return training._make_training_contracts( + {"unet": model}, + dataframe, + ExperimentConfig( + model_keys=("unet",), + folds=folds, + epochs=epochs, + compile_models=False, + ), + patch, + preprocessing_manifest, + )["unet"] + + first = contract((1, 2)) + second = contract((3,)) + changed_epoch = contract((3,), epochs=2) + changed_patch = contract( + (3,), + patch=SimpleNamespace(patch_overlap=0.25), + ) + changed_assignment = contract( + (3,), + dataframe=train_df.assign(fold=[1, 1, 3]), + ) + changed_model = contract( + (3,), + model=replace( + model_config, + model_spec={"arch_id": "different"}, + ), + ) + changed_loss = contract( + (3,), + model=replace( + model_config, + loss_spec={"loss_id": "different"}, + ), + ) + preprocessing_manifest.write_text( + '{"dataset_version": "different"}', + encoding="utf-8", + ) + changed_preprocessing = contract((3,)) + + self.assertEqual(first["sha256"], second["sha256"]) + for changed in ( + changed_epoch, + changed_patch, + changed_assignment, + changed_model, + changed_loss, + changed_preprocessing, + ): + self.assertNotEqual(first["sha256"], changed["sha256"]) + self.assertEqual( + first["payload"]["training_workflow_revision"], + training.VS_TRAINING_WORKFLOW_REVISION, + ) + self.assertEqual( + first["payload"]["loss_spec"], + model_config.loss_spec, + ) + + def test_sweep_rejects_an_existing_results_root(self): + experiment = ExperimentConfig( + model_keys=("unet",), + folds=(1,), + epochs=1, + compile_models=False, + ) + with tempfile.TemporaryDirectory() as directory: + preprocessing_manifest = Path(directory) / "preprocessing_manifest.json" + preprocessing_manifest.write_text("{}", encoding="utf-8") + results_root = Path(directory) / "existing" + results_root.mkdir() + with self.assertRaisesRegex(FileExistsError, "new --results-root"): + training.run_training_sweep( + {"unet": TRAINING_MODEL_CONFIGS["unet"]}, + pd.DataFrame({"case_id": ["a"], "fold": [1]}), + experiment=experiment, + patch_config=SimpleNamespace(), + preprocessing_manifest=preprocessing_manifest, + results_root=results_root, + ) + + def test_sweep_writes_subset_registry_and_all_data_inference_selection(self): experiment = ExperimentConfig( model_keys=("unet",), folds=(2, 4), @@ -412,11 +620,13 @@ def fake_train(model_config, fold, *args, **kwargs): ) with tempfile.TemporaryDirectory() as directory: + preprocessing_manifest = Path(directory) / "preprocessing_manifest.json" + preprocessing_manifest.write_text("{}", encoding="utf-8") with ( patch.object(training, "train_one_fold", side_effect=fake_train), patch.object( training, "train_all_data_model", return_value="run-final" - ), + ) as train_all_data_model, redirect_stdout(StringIO()), ): sweep = training.run_training_sweep( @@ -424,20 +634,71 @@ def fake_train(model_config, fold, *args, **kwargs): train_df, experiment=experiment, patch_config=SimpleNamespace(), - preprocessing_manifest="manifest.json", - results_root=directory, + preprocessing_manifest=preprocessing_manifest, + results_root=Path(directory) / "run", ) selection = json.loads(sweep.inference_run_ids_path.read_text()) + completed = json.loads(sweep.completed_run_ids_path.read_text()) + self.assertEqual( + selection["training_contracts"], + completed["training_contracts"], + ) + self.assertEqual( + train_all_data_model.call_args.kwargs["results_dir"], + Path(directory) / "run" / "unet", + ) self.assertEqual(selection["schema_version"], 1) self.assertEqual( - selection["models"], - { - "unet": { - "best": {"fold_2": "run-2", "fold_4": "run-4"}, - "final": {"all_data": "run-final"}, - } - }, + selection["models"]["unet"], {"final": {"all_data": "run-final"}} + ) + self.assertEqual( + completed["models"]["unet"]["best"], {"fold_2": "run-2", "fold_4": "run-4"} + ) + + def test_inference_selection_requires_canonical_five_folds(self): + contracts = {"unet": {"sha256": "contract"}} + subset_experiment = ExperimentConfig( + model_keys=("unet",), + folds=(1,), + ) + subset_sweep = training.TrainingSweep( + fold_runs={ + "unet": {1: SimpleNamespace(run_id="run-1")}, + } + ) + + folds = training.CROSS_VALIDATION_FOLDS + full_sweep = training.TrainingSweep( + fold_runs={ + "unet": {fold: SimpleNamespace(run_id=f"run-{fold}") for fold in folds} + } + ) + full_experiment = ExperimentConfig( + model_keys=("unet",), + folds=tuple(reversed(folds)), + ) + + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + subset_root = root / "subset" + subset_root.mkdir() + subset_path = training._write_inference_run_ids( + subset_sweep, subset_experiment, subset_root, contracts + ) + self.assertIsNone(subset_path) + self.assertFalse((subset_root / "inference_run_ids.json").exists()) + + full_root = root / "full" + full_root.mkdir() + full_path = training._write_inference_run_ids( + full_sweep, full_experiment, full_root, contracts + ) + manifest = json.loads(full_path.read_text(encoding="utf-8")) + + self.assertEqual( + list(manifest["models"]["unet"]["best"]), + [f"fold_{fold}" for fold in folds], ) diff --git a/research/vestibular_schwannoma/train_5fold.py b/research/vestibular_schwannoma/train_5fold.py index b44c2dc..d72758c 100644 --- a/research/vestibular_schwannoma/train_5fold.py +++ b/research/vestibular_schwannoma/train_5fold.py @@ -27,7 +27,7 @@ python train_5fold.py --skip-unavailable -The defaults request four models, folds 1-5, and 500 epochs. Models and folds run +The defaults request three models, folds 1-5, and 500 epochs. Models and folds run sequentially so only one model occupies GPU memory at a time. Each held-out fold is evaluated with TTA. @@ -41,8 +41,10 @@ then ``--queue-length``; a larger queue consumes more RAM. Preprocessing is cached in ``preprocessed/`` and outputs go below -``cv_results//`` unless overridden. Generated data, caches, MLflow -state, predictions, and weights stay outside Git through the project ``.gitignore``. +``cv_results//`` unless overridden. Every independently launched job must +use a new results root. Warm the preprocessing cache with one process before launching +multiple training jobs. Generated data, caches, MLflow state, predictions, and weights +stay outside Git through the project ``.gitignore``. This launcher performs cross-validation only. Use the shared workflow directly (or extend the CLI explicitly) if an all-data final/deployment model is required. @@ -63,19 +65,27 @@ PROJECT_ROOT = Path(__file__).resolve().parent -if str(PROJECT_ROOT) not in sys.path: - sys.path.insert(0, str(PROJECT_ROOT)) - -from workflow.config import ExperimentConfig, make_patch_config # noqa: E402 -from workflow.models import ( # noqa: E402 +RESEARCH_ROOT = PROJECT_ROOT.parent +if str(RESEARCH_ROOT) not in sys.path: + sys.path.insert(0, str(RESEARCH_ROOT)) + +from vestibular_schwannoma.workflow.config import ( # noqa: E402 + CROSS_VALIDATION_FOLDS, + ExperimentConfig, + make_patch_config, +) +from vestibular_schwannoma.workflow.models import ( # noqa: E402 TRAINING_MODEL_CONFIGS, get_training_model_configs, ) -from workflow.results import aggregate_results, build_model_comparison # noqa: E402 -from workflow.training import run_training_sweep # noqa: E402 +from vestibular_schwannoma.workflow.results import ( # noqa: E402 + aggregate_results, + build_model_comparison, +) +from vestibular_schwannoma.workflow.training import run_training_sweep # noqa: E402 -DEFAULT_MODELS = ("unet", "dynunet", "dynunet_small", "segmamba") +DEFAULT_MODELS = ("unet", "dynunet", "segmamba") def _parser() -> argparse.ArgumentParser: @@ -90,13 +100,13 @@ def _parser() -> argparse.ArgumentParser: nargs="+", choices=tuple(TRAINING_MODEL_CONFIGS), default=list(DEFAULT_MODELS), - help="Model keys to train in order (default: all four).", + help="Model keys to train in order (default: all three).", ) parser.add_argument( "--folds", nargs="+", type=int, - default=[1, 2, 3, 4, 5], + default=list(CROSS_VALIDATION_FOLDS), help="Held-out folds to run (default: 1 2 3 4 5).", ) parser.add_argument("--epochs", type=int, default=500) @@ -152,7 +162,10 @@ def _parser() -> argparse.ArgumentParser: parser.add_argument( "--results-root", type=Path, - help="Output directory (default: cv_results/).", + help=( + "New output directory; it must not already exist " + "(default: cv_results/)." + ), ) parser.add_argument( "--no-compile", @@ -252,7 +265,7 @@ def main(argv: list[str] | None = None) -> int: use_tta=True, compile_models=not args.no_compile, target_spacing=(0.4102, 0.4102, 1.5), - patch_size=(192, 192, 48), + patch_size=(256, 256, 48), preprocess_workers=args.preprocess_workers, samples_per_volume=args.samples_per_volume, queue_num_workers=args.queue_workers, @@ -263,7 +276,7 @@ def main(argv: list[str] | None = None) -> int: data_csv = _project_path(args.data_csv) preprocessed_dir = _project_path(args.preprocessed_dir) - timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S%fZ") results_root = _project_path(args.results_root or Path("cv_results") / timestamp) _print_plan(experiment, data_csv, results_root) diff --git a/research/vestibular_schwannoma/workflow/config.py b/research/vestibular_schwannoma/workflow/config.py index 821eeef..a7286ef 100644 --- a/research/vestibular_schwannoma/workflow/config.py +++ b/research/vestibular_schwannoma/workflow/config.py @@ -11,10 +11,13 @@ make_output_spec, ) from .models import TRAINING_MODEL_CONFIGS +from .run_selection import ( + CROSS_VALIDATION_FOLDS, + INFERENCE_RUN_IDS_FILENAME as INFERENCE_RUN_IDS_FILENAME, + INFERENCE_RUN_IDS_SCHEMA as INFERENCE_RUN_IDS_SCHEMA, +) -INFERENCE_RUN_IDS_SCHEMA = 1 -INFERENCE_RUN_IDS_FILENAME = "inference_run_ids.json" VS_OUTPUT_SPEC = make_output_spec("multiclass_segmentation", classes=2) @@ -23,7 +26,7 @@ class ExperimentConfig: """Immutable settings shared by cross-validation and all-data training.""" model_keys: tuple[str, ...] = ("unet", "dynunet", "segmamba") - folds: tuple[int, ...] = (1, 2, 3, 4, 5) + folds: tuple[int, ...] = CROSS_VALIDATION_FOLDS run_cross_validation: bool = True train_all_data: bool = False training_seed: int = 42 @@ -33,7 +36,7 @@ class ExperimentConfig: use_tta: bool = True compile_models: bool = True target_spacing: tuple[float, float, float] = (0.4102, 0.4102, 1.5) - patch_size: tuple[int, int, int] = (192, 192, 48) + patch_size: tuple[int, int, int] = (256, 256, 48) preprocess_workers: int = min(32, os.cpu_count() or 1) samples_per_volume: int = 4 queue_num_workers: int = 4 @@ -54,7 +57,9 @@ def __post_init__(self) -> None: raise ValueError("Enable cross-validation, all-data training, or both") if self.run_cross_validation: if not self.folds: - raise ValueError("folds must not be empty when cross-validation is enabled") + raise ValueError( + "folds must not be empty when cross-validation is enabled" + ) if len(set(self.folds)) != len(self.folds): raise ValueError("folds contains duplicates") @@ -78,7 +83,9 @@ def __post_init__(self) -> None: if invalid: raise ValueError(f"These settings must be positive: {invalid}") - if len(self.target_spacing) != 3 or any(value <= 0 for value in self.target_spacing): + if len(self.target_spacing) != 3 or any( + value <= 0 for value in self.target_spacing + ): raise ValueError("target_spacing must contain three positive values") if len(self.patch_size) != 3 or any(value <= 0 for value in self.patch_size): raise ValueError("patch_size must contain three positive values") diff --git a/research/vestibular_schwannoma/workflow/inference.py b/research/vestibular_schwannoma/workflow/inference.py index f3ae993..44384b2 100644 --- a/research/vestibular_schwannoma/workflow/inference.py +++ b/research/vestibular_schwannoma/workflow/inference.py @@ -2,7 +2,6 @@ from __future__ import annotations -import json from collections.abc import Mapping from dataclasses import dataclass from pathlib import Path @@ -14,7 +13,11 @@ patch_config_to_dict, read_safetensors_metadata, ) -from .config import INFERENCE_RUN_IDS_SCHEMA, VS_OUTPUT_SPEC +from .config import VS_OUTPUT_SPEC +from .run_selection import ( + merge_fold_run_selections as merge_fold_run_selections, + read_inference_run_ids as _read_inference_run_ids, +) @dataclass(frozen=True) @@ -43,49 +46,6 @@ def _validate_member_mapping(name: str, values: Mapping | None) -> dict: return resolved -def _read_inference_run_ids( - selection_file: str | Path, - *, - model_key: str, - artifact_role: str, -) -> dict[str, str]: - path = Path(selection_file).expanduser() - if not path.is_file(): - raise FileNotFoundError(f"Inference run selection not found: {path}") - try: - manifest = json.loads(path.read_text(encoding="utf-8")) - except json.JSONDecodeError as exc: - raise ValueError(f"Invalid inference run selection JSON: {path}") from exc - if not isinstance(manifest, dict): - raise ValueError("Inference run selection must be a JSON object") - if manifest.get("schema_version") != INFERENCE_RUN_IDS_SCHEMA: - raise ValueError( - f"Unsupported inference run selection schema: " - f"{manifest.get('schema_version')!r}" - ) - if not isinstance(manifest.get("run_group"), str) or not manifest["run_group"]: - raise ValueError("Inference run selection has an invalid run_group") - models = manifest.get("models") - if not isinstance(models, dict): - raise ValueError("Inference run selection has an invalid models mapping") - if model_key not in models: - raise ValueError( - f"Model {model_key!r} is not ready in {path}; " - f"available models: {sorted(models)}" - ) - roles = models[model_key] - if not isinstance(roles, dict) or artifact_role not in roles: - available = sorted(roles) if isinstance(roles, dict) else [] - raise ValueError( - f"Role {artifact_role!r} is not ready for model {model_key!r}; " - f"available roles: {available}" - ) - run_ids = _validate_member_mapping("inference run selection", roles[artifact_role]) - if not run_ids: - raise ValueError("Inference run selection contains no model members") - return run_ids - - def load_inference_models( *, run_selection_file: str | Path | None = None, @@ -113,10 +73,6 @@ def load_inference_models( model_key=model_key, artifact_role=artifact_role, ) - if any(not isinstance(run_id, str) or not run_id for run_id in run_ids.values()): - raise ValueError("Every MLflow run ID must be a non-empty string") - if len(set(run_ids.values())) != len(run_ids): - raise ValueError("Inference run selection contains duplicate run IDs") resolved = find_model_artifacts( run_ids=run_ids, artifact_role=artifact_role, @@ -134,7 +90,9 @@ def load_inference_models( missing = [str(path) for path in artifacts.values() if not path.is_file()] if missing: raise FileNotFoundError(f"Model artifacts not found: {missing}") - invalid = [str(path) for path in artifacts.values() if path.suffix != ".safetensors"] + invalid = [ + str(path) for path in artifacts.values() if path.suffix != ".safetensors" + ] if invalid: raise ValueError(f"All model artifacts must be Safetensors files: {invalid}") canonical = [path.resolve() for path in artifacts.values()] diff --git a/research/vestibular_schwannoma/workflow/models.py b/research/vestibular_schwannoma/workflow/models.py index 9014f87..2ebbdd7 100644 --- a/research/vestibular_schwannoma/workflow/models.py +++ b/research/vestibular_schwannoma/workflow/models.py @@ -36,6 +36,7 @@ class TrainingModelConfig: key: str display_name: str model_spec: dict + loss_spec: dict make_loss: Callable[[], object] experiment_name: str supports_compile: bool = True @@ -45,6 +46,22 @@ def checkpoint_name(self) -> str: return f"best_{self.key}" +DICE_CE_LOSS_SPEC = { + "loss_id": "monai.dice_ce", + "kwargs": { + "to_onehot_y": True, + "softmax": True, + "include_background": False, + "batch": True, + }, +} +DYNUNET_LOSS_SPEC = { + "loss_id": "monai.deep_supervision", + "kwargs": {"weight_mode": "exp"}, + "base_loss": DICE_CE_LOSS_SPEC, +} + + def _make_dice_ce_loss() -> CustomLoss: return CustomLoss( loss_func=DiceCELoss( @@ -72,7 +89,7 @@ def _make_dynunet_loss() -> CustomLoss: "spatial_dims": 3, "in_channels": 1, "out_channels": 2, - "channels": (64, 128, 256, 512, 1024), + "channels": (32, 64, 128, 256, 320), "strides": (2, 2, 2, 2), "num_res_units": 4, "norm": "INSTANCE", @@ -103,9 +120,7 @@ def _make_dynunet_spec(filters: list[int]) -> dict: ) -DYNUNET_SPEC = _make_dynunet_spec([64, 128, 256, 512, 1024]) - -DYNUNET_SMALL_SPEC = _make_dynunet_spec([32, 64, 128, 256, 512]) +DYNUNET_SPEC = _make_dynunet_spec([32, 64, 128, 256, 320]) SEGMAMBA_SPEC = make_model_spec( "segmamba.v2", @@ -125,27 +140,23 @@ def _make_dynunet_spec(filters: list[int]) -> dict: key="unet", display_name="UNet", model_spec=UNET_SPEC, + loss_spec=DICE_CE_LOSS_SPEC, make_loss=_make_dice_ce_loss, experiment_name="vestibular_schwannoma_unet", ), "dynunet": TrainingModelConfig( key="dynunet", - display_name="DynUNet", + display_name="DynUNet Small (32-320)", model_spec=DYNUNET_SPEC, + loss_spec=DYNUNET_LOSS_SPEC, make_loss=_make_dynunet_loss, experiment_name="vestibular_schwannoma_dynunet", ), - "dynunet_small": TrainingModelConfig( - key="dynunet_small", - display_name="DynUNet Small (32-512)", - model_spec=DYNUNET_SMALL_SPEC, - make_loss=_make_dynunet_loss, - experiment_name="vestibular_schwannoma_dynunet_small", - ), "segmamba": TrainingModelConfig( key="segmamba", display_name="SegMamba V2", model_spec=SEGMAMBA_SPEC, + loss_spec=DICE_CE_LOSS_SPEC, make_loss=_make_dice_ce_loss, experiment_name="vestibular_schwannoma_segmamba", supports_compile=False, diff --git a/research/vestibular_schwannoma/workflow/run_selection.py b/research/vestibular_schwannoma/workflow/run_selection.py new file mode 100644 index 0000000..d10f066 --- /dev/null +++ b/research/vestibular_schwannoma/workflow/run_selection.py @@ -0,0 +1,255 @@ +"""Read, validate, and combine MLflow inference run selections.""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Mapping, Sequence +from pathlib import Path + + +INFERENCE_RUN_IDS_SCHEMA = 1 +INFERENCE_RUN_IDS_FILENAME = "inference_run_ids.json" +COMPLETED_RUN_IDS_FILENAME = "completed_run_ids.json" +INFERENCE_SELECTION_KIND = "inference_selection" +COMPLETED_REGISTRY_KIND = "completed_registry" +TRAINING_CONTRACT_SCHEMA = 1 +CROSS_VALIDATION_FOLDS = (1, 2, 3, 4, 5) + + +def _payload_sha256(payload: Mapping) -> str: + canonical = json.dumps( + payload, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + ) + return hashlib.sha256(canonical.encode("utf-8")).hexdigest() + + +def make_training_contract(payload: Mapping) -> dict: + """Create a checksummed, JSON-serializable training merge contract.""" + + resolved = dict(payload) + return { + "schema_version": TRAINING_CONTRACT_SCHEMA, + "sha256": _payload_sha256(resolved), + "payload": resolved, + } + + +def read_inference_manifest(selection_file: str | Path) -> tuple[Path, dict]: + """Read and validate the common structure of an inference run manifest.""" + + path = Path(selection_file).expanduser() + if not path.is_file(): + raise FileNotFoundError(f"Inference run selection not found: {path}") + try: + manifest = json.loads(path.read_text(encoding="utf-8")) + except json.JSONDecodeError as error: + raise ValueError(f"Invalid inference run selection JSON: {path}") from error + if not isinstance(manifest, dict): + raise ValueError("Inference run selection must be a JSON object") + if manifest.get("schema_version") != INFERENCE_RUN_IDS_SCHEMA: + raise ValueError( + f"Unsupported inference run selection schema: " + f"{manifest.get('schema_version')!r}" + ) + if not isinstance(manifest.get("run_group"), str) or not manifest["run_group"]: + raise ValueError("Inference run selection has an invalid run_group") + manifest_kind = manifest.get("manifest_kind", INFERENCE_SELECTION_KIND) + if manifest_kind not in {INFERENCE_SELECTION_KIND, COMPLETED_REGISTRY_KIND}: + raise ValueError(f"Unsupported inference manifest kind: {manifest_kind!r}") + if not isinstance(manifest.get("models"), dict): + raise ValueError("Inference run selection has an invalid models mapping") + return path, manifest + + +def _run_ids_from_manifest( + manifest: Mapping, + path: Path, + *, + model_key: str, + artifact_role: str, +) -> dict[str, str]: + models = manifest["models"] + if model_key not in models: + raise ValueError( + f"Model {model_key!r} is not ready in {path}; " + f"available models: {sorted(models)}" + ) + roles = models[model_key] + if not isinstance(roles, dict) or artifact_role not in roles: + available = sorted(roles) if isinstance(roles, dict) else [] + raise ValueError( + f"Role {artifact_role!r} is not ready for model {model_key!r}; " + f"available roles: {available}" + ) + values = roles[artifact_role] + if not isinstance(values, Mapping): + raise TypeError( + "inference run selection must be a mapping of member ID to value" + ) + run_ids = dict(values) + if any(not isinstance(member, str) or not member for member in run_ids): + raise ValueError( + "Every inference run selection member ID must be a non-empty string" + ) + if not run_ids: + raise ValueError("Inference run selection contains no model members") + if any(not isinstance(run_id, str) or not run_id for run_id in run_ids.values()): + raise ValueError("Every MLflow run ID must be a non-empty string") + if len(set(run_ids.values())) != len(run_ids): + raise ValueError("Inference run selection contains duplicate run IDs") + return run_ids + + +def read_inference_run_ids( + selection_file: str | Path, + *, + model_key: str, + artifact_role: str, +) -> dict[str, str]: + """Read one model role from a validated inference run selection.""" + + path, manifest = read_inference_manifest(selection_file) + if manifest.get("manifest_kind") == COMPLETED_REGISTRY_KIND: + raise ValueError( + "completed_run_ids.json is a partial registry and cannot be used " + "directly for inference; merge it into inference_run_ids.json first" + ) + return _run_ids_from_manifest( + manifest, + path, + model_key=model_key, + artifact_role=artifact_role, + ) + + +def _training_contract( + manifest: Mapping, + path: Path, + *, + model_key: str, +) -> dict: + contracts = manifest.get("training_contracts") + if not isinstance(contracts, Mapping) or model_key not in contracts: + raise ValueError( + f"Inference run selection has no training contract for {model_key!r}: " + f"{path}" + ) + contract = contracts[model_key] + if not isinstance(contract, Mapping): + raise ValueError(f"Invalid training contract for {model_key!r}: {path}") + contract = dict(contract) + if contract.get("schema_version") != TRAINING_CONTRACT_SCHEMA: + raise ValueError( + f"Unsupported training contract schema for {model_key!r}: " + f"{contract.get('schema_version')!r}" + ) + payload = contract.get("payload") + digest = contract.get("sha256") + if not isinstance(payload, Mapping) or not isinstance(digest, str) or not digest: + raise ValueError(f"Invalid training contract for {model_key!r}: {path}") + if _payload_sha256(payload) != digest: + raise ValueError( + f"Training contract checksum mismatch for {model_key!r}: {path}" + ) + return contract + + +def merge_fold_run_selections( + selection_files: Sequence[str | Path], + *, + model_key: str, + output_root: str | Path, +) -> Path: + """Merge disjoint, contract-compatible folds into one inference manifest.""" + + if isinstance(selection_files, (str, Path)): + raise TypeError("selection_files must be a sequence of manifest paths") + paths = [Path(path).expanduser() for path in selection_files] + if not paths: + raise ValueError("At least one inference run selection is required") + if not isinstance(model_key, str) or not model_key: + raise ValueError("model_key must be a non-empty string") + + folds = CROSS_VALIDATION_FOLDS + expected_members = tuple(f"fold_{fold}" for fold in folds) + + merged = {} + shared_contract = None + for path in paths: + resolved_path, manifest = read_inference_manifest(path) + run_ids = _run_ids_from_manifest( + manifest, + resolved_path, + model_key=model_key, + artifact_role="best", + ) + contract = _training_contract( + manifest, + resolved_path, + model_key=model_key, + ) + if shared_contract is None: + shared_contract = contract + elif contract["sha256"] != shared_contract["sha256"]: + raise ValueError( + f"Training contract mismatch for {model_key!r}: {resolved_path}" + ) + for member, run_id in run_ids.items(): + if member in merged: + raise ValueError( + f"Duplicate model member {member!r} across inference selections" + ) + merged[member] = run_id + + actual_members = set(merged) + expected_member_set = set(expected_members) + if actual_members != expected_member_set: + missing = [ + member for member in expected_members if member not in actual_members + ] + unexpected = sorted(actual_members - expected_member_set) + raise ValueError( + "Merged fold members do not match the expected folds; " + f"missing={missing}, unexpected={unexpected}" + ) + if len(set(merged.values())) != len(merged): + raise ValueError("Merged inference selection contains duplicate MLflow run IDs") + + destination_root = Path(output_root).expanduser() + try: + destination_root.mkdir(parents=True, exist_ok=False) + except FileExistsError as error: + raise FileExistsError( + f"Merged results root already exists: {destination_root}. " + "Choose a new --output-root." + ) from error + + manifest = { + "schema_version": INFERENCE_RUN_IDS_SCHEMA, + "manifest_kind": INFERENCE_SELECTION_KIND, + "run_group": destination_root.name, + "training_contracts": {model_key: shared_contract}, + "models": { + model_key: {"best": {member: merged[member] for member in expected_members}} + }, + } + destination = destination_root / INFERENCE_RUN_IDS_FILENAME + temporary = destination.with_suffix(".json.tmp") + try: + temporary.write_text( + json.dumps(manifest, indent=2) + "\n", + encoding="utf-8", + ) + temporary.replace(destination) + except Exception: + temporary.unlink(missing_ok=True) + try: + destination_root.rmdir() + except OSError: + pass + raise + return destination diff --git a/research/vestibular_schwannoma/workflow/training.py b/research/vestibular_schwannoma/workflow/training.py index e5a91e7..29ec362 100644 --- a/research/vestibular_schwannoma/workflow/training.py +++ b/research/vestibular_schwannoma/workflow/training.py @@ -3,9 +3,10 @@ from __future__ import annotations import gc +import hashlib import json import traceback -from dataclasses import dataclass, field +from dataclasses import asdict, dataclass, field from pathlib import Path import numpy as np @@ -23,16 +24,23 @@ PatchConfig, create_mlflow_callback, evaluate_segmentations, + patch_config_to_dict, patch_inference, ) -from .config import ( +from .config import VS_OUTPUT_SPEC, ExperimentConfig, make_gpu_augmentation +from .models import TrainingModelConfig, build_training_model +from .run_selection import ( + COMPLETED_REGISTRY_KIND, + COMPLETED_RUN_IDS_FILENAME, + CROSS_VALIDATION_FOLDS, INFERENCE_RUN_IDS_FILENAME, INFERENCE_RUN_IDS_SCHEMA, - VS_OUTPUT_SPEC, - ExperimentConfig, - make_gpu_augmentation, + INFERENCE_SELECTION_KIND, + make_training_contract, ) -from .models import TrainingModelConfig, build_training_model + + +VS_TRAINING_WORKFLOW_REVISION = 1 @dataclass(frozen=True) @@ -64,43 +72,186 @@ class TrainingSweep: fold_runs: dict[str, dict[int, FoldRun]] = field(default_factory=dict) all_data_run_ids: dict[str, str] = field(default_factory=dict) failures: list[RunFailure] = field(default_factory=list) + completed_run_ids_path: Path | None = None inference_run_ids_path: Path | None = None -def _write_inference_run_ids( - sweep: TrainingSweep, +def _claim_results_root(results_root: Path) -> None: + """Atomically reserve a new output root for exactly one training sweep.""" + + try: + results_root.mkdir(parents=True, exist_ok=False) + except FileExistsError as error: + raise FileExistsError( + f"Results root already exists: {results_root}. " + "Use a new --results-root for every independently launched training job." + ) from error + + +def _json_sha256(value) -> str: + canonical = json.dumps( + value, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + ) + return hashlib.sha256(canonical.encode("utf-8")).hexdigest() + + +def _make_training_contracts( + model_configs: dict[str, TrainingModelConfig], + train_df: pd.DataFrame, experiment: ExperimentConfig, - results_root: Path, -) -> Path: - """Persist exact MLflow runs for each fully completed inference role.""" + patch_config: PatchConfig, + preprocessing_manifest: str | Path, +) -> dict[str, dict]: + """Build fold-independent scientific contracts for safe cross-job merging.""" + + required = {"case_id", "fold"} + missing = sorted(required - set(train_df.columns)) + if missing: + raise ValueError(f"Training data is missing contract columns: {missing}") + if train_df[["case_id", "fold"]].isna().any().any(): + raise ValueError("Training contract case_id and fold values must be present") + + assignments = [ + {"case_id": str(case_id), "fold": int(fold)} + for case_id, fold in ( + train_df[["case_id", "fold"]] + .sort_values(["case_id", "fold"], kind="stable") + .itertuples(index=False, name=None) + ) + ] + manifest_path = Path(preprocessing_manifest).expanduser() + if not manifest_path.is_file(): + raise FileNotFoundError(f"Preprocessing manifest not found: {manifest_path}") + + experiment_payload = asdict(experiment) + for field_name in ("model_keys", "folds", "continue_on_error"): + experiment_payload.pop(field_name) + shared_payload = { + "experiment": experiment_payload, + "fold_assignments_sha256": _json_sha256(assignments), + "patch_config": patch_config_to_dict(patch_config), + "preprocessing_manifest_sha256": hashlib.sha256( + manifest_path.read_bytes() + ).hexdigest(), + "training_case_count": len(assignments), + "training_workflow_revision": VS_TRAINING_WORKFLOW_REVISION, + } + + contracts = {} + for model_key, model_config in model_configs.items(): + contracts[model_key] = make_training_contract( + { + **shared_payload, + "loss_factory": ( + f"{model_config.make_loss.__module__}." + f"{model_config.make_loss.__qualname__}" + ), + "loss_spec": model_config.loss_spec, + "model_key": model_key, + "model_spec": model_config.model_spec, + "supports_compile": model_config.supports_compile, + } + ) + return contracts + +def _selection_models( + sweep: TrainingSweep, + experiment: ExperimentConfig, + *, + complete_only: bool, +) -> dict: models = {} - requested_folds = set(experiment.folds) + canonical_folds = set(CROSS_VALIDATION_FOLDS) for model_key in experiment.model_keys: roles = {} fold_runs = sweep.fold_runs.get(model_key, {}) - if experiment.run_cross_validation and set(fold_runs) == requested_folds: + folds_complete = set(fold_runs) == canonical_folds + fold_order = ( + CROSS_VALIDATION_FOLDS + if complete_only and folds_complete + else experiment.folds + ) + completed_folds = [fold for fold in fold_order if fold in fold_runs] + if ( + experiment.run_cross_validation + and completed_folds + and (folds_complete or not complete_only) + ): roles["best"] = { - f"fold_{fold}": fold_runs[fold].run_id for fold in experiment.folds + f"fold_{fold}": fold_runs[fold].run_id for fold in completed_folds } if experiment.train_all_data and model_key in sweep.all_data_run_ids: roles["final"] = {"all_data": sweep.all_data_run_ids[model_key]} if roles: models[model_key] = roles + return models + +def _write_run_ids( + *, + models: dict, + training_contracts: dict[str, dict], + results_root: Path, + filename: str, + manifest_kind: str, +) -> Path: manifest = { "schema_version": INFERENCE_RUN_IDS_SCHEMA, + "manifest_kind": manifest_kind, "run_group": results_root.name, + "training_contracts": { + model_key: training_contracts[model_key] for model_key in models + }, "models": models, } - results_root.mkdir(parents=True, exist_ok=True) - destination = results_root / INFERENCE_RUN_IDS_FILENAME + destination = results_root / filename temporary = destination.with_suffix(".json.tmp") temporary.write_text(json.dumps(manifest, indent=2) + "\n", encoding="utf-8") temporary.replace(destination) return destination +def _write_completed_run_ids( + sweep: TrainingSweep, + experiment: ExperimentConfig, + results_root: Path, + training_contracts: dict[str, dict], +) -> Path: + """Persist every completed fold so interrupted subset jobs remain mergeable.""" + + return _write_run_ids( + models=_selection_models(sweep, experiment, complete_only=False), + training_contracts=training_contracts, + results_root=results_root, + filename=COMPLETED_RUN_IDS_FILENAME, + manifest_kind=COMPLETED_REGISTRY_KIND, + ) + + +def _write_inference_run_ids( + sweep: TrainingSweep, + experiment: ExperimentConfig, + results_root: Path, + training_contracts: dict[str, dict], +) -> Path | None: + """Persist only canonical five-fold or all-data inference roles.""" + + models = _selection_models(sweep, experiment, complete_only=True) + if not models: + return None + return _write_run_ids( + models=models, + training_contracts=training_contracts, + results_root=results_root, + filename=INFERENCE_RUN_IDS_FILENAME, + manifest_kind=INFERENCE_SELECTION_KIND, + ) + + def _mark_failed(callback) -> None: """Mark a callback-owned run failed through its public lifecycle API.""" @@ -138,11 +289,15 @@ def _gpu_augmentations_json(gpu_augmentation) -> str: def _component_counts(prediction, ground_truth_path: str | Path) -> dict[str, int]: """Count predicted regions and regions with no ground-truth overlap.""" - prediction_array = np.asarray( - prediction.detach().cpu() - if isinstance(prediction, torch.Tensor) - else prediction - ).squeeze().astype(bool) + prediction_array = ( + np.asarray( + prediction.detach().cpu() + if isinstance(prediction, torch.Tensor) + else prediction + ) + .squeeze() + .astype(bool) + ) ground_truth = ( np.asarray(tio.LabelMap(ground_truth_path).data).squeeze().astype(bool) ) @@ -203,15 +358,11 @@ def evaluate_fold( for prediction, mask_path in zip(predictions, mask_paths) ] ) - results = pd.concat( - [results.reset_index(drop=True), component_metrics], axis=1 - ) + results = pd.concat([results.reset_index(drop=True), component_metrics], axis=1) results.insert(1, "image", [Path(path).name for path in image_paths]) results.to_csv(fold_dir / "results.csv", index=False) - numeric = results.select_dtypes(include="number").replace( - [np.inf, -np.inf], np.nan - ) + numeric = results.select_dtypes(include="number").replace([np.inf, -np.inf], np.nan) mlflow_callback.log_metrics_table(results, display=False) mlflow_callback.log_metrics( {f"val_{metric}": numeric[metric].mean() for metric in numeric.columns} @@ -219,7 +370,9 @@ def evaluate_fold( mlflow_callback.log_dataframe(results) print(f"[Fold {fold}] Results saved to {fold_dir / 'results.csv'}") - print(f"[Fold {fold}] DSC: {results['dsc'].mean():.4f} +/- {results['dsc'].std():.4f}") + print( + f"[Fold {fold}] DSC: {results['dsc'].mean():.4f} +/- {results['dsc'].std():.4f}" + ) return results @@ -236,6 +389,8 @@ def train_one_fold( """Train a fresh model and evaluate it on exactly one held-out fold.""" _set_training_seed(experiment.training_seed) + fold_dir = Path(results_dir) / f"fold_{fold}" + fold_dir.mkdir(parents=True, exist_ok=False) fold_df = train_df.copy() fold_df["is_val"] = fold_df["fold"] == fold gpu_augmentation = make_gpu_augmentation(experiment) @@ -267,6 +422,8 @@ def train_one_fold( model, loss_func=model_config.make_loss(), metrics=[AccumulatedDice(n_classes=2)], + path=fold_dir, + model_dir="checkpoints", ).to_bf16() save_best = EMACheckpoint( monitor="accumulated_dice", @@ -311,7 +468,7 @@ def train_one_fold( model_key=model_config.key, fold=fold, run_id=mlflow_callback.run_id, - results_dir=Path(results_dir) / f"fold_{fold}", + results_dir=fold_dir, results=results, ) except Exception: @@ -328,7 +485,7 @@ def train_all_data_model( experiment: ExperimentConfig, patch_config: PatchConfig, preprocessing_manifest: str | Path, - run_group: str, + results_dir: str | Path, ) -> str: """Train every case, duplicating one stable case only for internal monitoring.""" @@ -337,6 +494,8 @@ def train_all_data_model( raise ValueError("No cases are available for all-data training") if "case_id" not in train_df.columns or train_df["case_id"].isna().any(): raise ValueError("case_id is required for stable all-data monitoring") + all_data_dir = Path(results_dir) / "all_data" + all_data_dir.mkdir(parents=True, exist_ok=False) monitor_df = train_df.sort_values("case_id", kind="stable").iloc[[0]].copy() monitor_case_id = str(monitor_df.iloc[0]["case_id"]) print( @@ -369,6 +528,8 @@ def train_all_data_model( model, loss_func=model_config.make_loss(), metrics=[AccumulatedDice(n_classes=2)], + path=all_data_dir, + model_dir="checkpoints", ).to_bf16() mlflow_callback = create_mlflow_callback( learn, @@ -376,7 +537,7 @@ def train_all_data_model( run_name="all_data", extra_tags={ "training_scope": "all_data", - "run_group": run_group, + "run_group": Path(results_dir).parent.name, "monitor_case_id": monitor_case_id, "monitor_is_training_duplicate": "true", }, @@ -417,6 +578,14 @@ def run_training_sweep( """Run every requested fold and/or all-data model with structured failures.""" results_root = Path(results_root) + training_contracts = _make_training_contracts( + model_configs, + train_df, + experiment, + patch_config, + preprocessing_manifest, + ) + _claim_results_root(results_root) sweep = TrainingSweep() for model_key, model_config in model_configs.items(): @@ -436,6 +605,12 @@ def run_training_sweep( results_dir=model_results_dir, ) sweep.fold_runs[model_key][fold] = run + sweep.completed_run_ids_path = _write_completed_run_ids( + sweep, + experiment, + results_root, + training_contracts, + ) except Exception as error: sweep.failures.append( RunFailure( @@ -462,7 +637,13 @@ def run_training_sweep( experiment=experiment, patch_config=patch_config, preprocessing_manifest=preprocessing_manifest, - run_group=results_root.name, + results_dir=model_results_dir, + ) + sweep.completed_run_ids_path = _write_completed_run_ids( + sweep, + experiment, + results_root, + training_contracts, ) print( f"[{model_key}] all-data MLflow run: " @@ -477,17 +658,29 @@ def run_training_sweep( message=str(error), ) ) - print( - f"[FAILED] {model_key} all_data: " - f"{type(error).__name__}: {error}" - ) + print(f"[FAILED] {model_key} all_data: {type(error).__name__}: {error}") traceback.print_exc() if not experiment.continue_on_error: raise + sweep.completed_run_ids_path = _write_completed_run_ids( + sweep, + experiment, + results_root, + training_contracts, + ) sweep.inference_run_ids_path = _write_inference_run_ids( - sweep, experiment, results_root + sweep, + experiment, + results_root, + training_contracts, ) - print(f"Inference run selection: {sweep.inference_run_ids_path}") + if sweep.inference_run_ids_path is None: + print( + "Inference run selection not written; merge completed fold registries " + "to create one." + ) + else: + print(f"Inference run selection: {sweep.inference_run_ids_path}") print("\nSweep complete.") return sweep