"""Pipeline task registry and scenario definitions. A *task* is a callable ``task(tracker, logger, config, state) -> None``. A *scenario* is an ordered list of task ids the pipeline runs. Tasks communicate through the mutable ``state`` dict (dataloaders, model paths, control events). Pipeline stages (see ai/ARCHITECTURE.md): 1. load_data -> implemented 2. train_regular -> implemented 3. train_bayesian -> implemented 4. evaluate_regular -> implemented (TEST split; *_val variants do VALIDATION) 5. evaluate_bayesian-> implemented (TEST split; *_val variants do VALIDATION) 6. evaluate_noisy -> implemented (TEST split; *_val variant does VALIDATION) combine_evaluations -> pools the per-split evaluations into combined_*.nc 7. run_analysis -> reads evaluations (prefers combined_*), writes analysis/. load_models -> implemented; discovers saved .pt ensembles on disk and records their paths in ``state`` so evaluation is decoupled from training (training frees models from VRAM after saving). Evaluation loads one model at a time. """ import os from typing import Any, Callable, Dict, List, Tuple, TypedDict from . import analyze from . import evaluate from . import load_data from . import load_models from . import train_bayesian from . import train_normal from . import kfold # imported last: depends on the sibling task modules above class TaskEntry(TypedDict): task_name: str task_func: Callable[..., None] class ScenarioEntry(TypedDict): label: str tasks: List[str] PIPELINE_TASKS: Dict[str, TaskEntry] = { "load_data": { "task_name": "Load Image and ADNIMERGE", "task_func": load_data.load_data_task, }, "train_regular": { "task_name": "Train Regular Models", "task_func": train_normal.train_normal_task, }, "train_bayesian": { "task_name": "Train Bayesian Models", "task_func": train_bayesian.train_bayesian_task, }, "load_models": { "task_name": "Load Saved Models", "task_func": load_models.load_models_task, }, "evaluate_regular": { "task_name": "Evaluate Regular Models", "task_func": evaluate.evaluate_normal_task, }, "evaluate_bayesian": { "task_name": "Evaluate Bayesian Models", "task_func": evaluate.evaluate_bayesian_task, }, "evaluate_noisy": { "task_name": "Evaluate Models on Noised Data", "task_func": evaluate.evaluate_noisy_task, }, "evaluate_regular_val": { "task_name": "Evaluate Regular Models (Validation)", "task_func": evaluate.evaluate_normal_val_task, }, "evaluate_bayesian_val": { "task_name": "Evaluate Bayesian Models (Validation)", "task_func": evaluate.evaluate_bayesian_val_task, }, "evaluate_noisy_val": { "task_name": "Evaluate Models on Noised Data (Validation)", "task_func": evaluate.evaluate_noisy_val_task, }, "combine_evaluations": { "task_name": "Combine Split Evaluations", "task_func": evaluate.combine_evaluations_task, }, "run_analysis": { "task_name": "Analyze Evaluations", "task_func": analyze.run_analysis_task, }, "run_kfold": { "task_name": "K-Fold Cross-Validation", "task_func": kfold.run_kfold_task, }, } # -------------------------------------------------------------------------- # Output artifact categories and overwrite protection # -------------------------------------------------------------------------- # Each category maps to the work_dir subdirectories it owns and whether it is # "protected" -- i.e. expensive/slow to regenerate, so overwriting it warrants a # confirmation prompt. Analysis artifacts are cheap to regenerate and therefore # unprotected: reruns of an analysis-only scenario proceed without a warning. class ArtifactCategory(TypedDict): dirs: Tuple[str, ...] protected: bool desc: str ARTIFACT_CATEGORIES: Dict[str, ArtifactCategory] = { "models": { "dirs": ("normal_models", "bayesian_models"), "protected": True, "desc": "trained model checkpoints", }, "evaluations": { "dirs": ("evaluations",), "protected": True, "desc": "model evaluation outputs (netCDF)", }, "kfold": { "dirs": ("folds",), "protected": True, "desc": "per-fold k-fold models & evaluations", }, "analysis": { "dirs": (analyze.ANALYSIS_SUBDIR,), "protected": False, "desc": "analysis artifacts (regeneratable)", }, } # Which artifact category each task writes to disk. Tasks not listed here (e.g. # load_data, load_models) produce no persistent, overwriteable output. TASK_WRITES = { "train_regular": ("models",), "train_bayesian": ("models",), "evaluate_regular": ("evaluations",), "evaluate_bayesian": ("evaluations",), "evaluate_noisy": ("evaluations",), "evaluate_regular_val": ("evaluations",), "evaluate_bayesian_val": ("evaluations",), "evaluate_noisy_val": ("evaluations",), "combine_evaluations": ("evaluations",), "run_analysis": ("analysis",), # k-fold writes per-fold models+evals under folds/ and the aggregated OOF # files under evaluations/; both are protected (slow to regenerate). "run_kfold": ("kfold", "evaluations"), } def scenario_output_categories(scenario_id: str) -> List[str]: """Return the artifact categories a scenario writes, in first-seen order.""" categories: List[str] = [] for task_id in SCENARIOS[scenario_id]["tasks"]: for category in TASK_WRITES.get(task_id, ()): if category not in categories: categories.append(category) return categories def protected_output_conflicts( work_dir: str, scenario_id: str ) -> List[Tuple[str, str]]: """Return ``(category, subdir)`` for PROTECTED outputs the scenario writes that already exist and are non-empty. Empty result => safe to run without a destructive-overwrite prompt (only unprotected/regeneratable outputs, if any, would be replaced).""" conflicts: List[Tuple[str, str]] = [] for category in scenario_output_categories(scenario_id): meta = ARTIFACT_CATEGORIES.get(category) if not meta or not meta["protected"]: continue for subdir in meta["dirs"]: path = os.path.join(work_dir, subdir) if os.path.isdir(path) and os.listdir(path): conflicts.append((category, subdir)) return conflicts def describe_protected_conflicts(work_dir: str, scenario_id: str) -> List[str]: """Human-readable one-liners for each protected overwrite conflict.""" lines: List[str] = [] for category, subdir in protected_output_conflicts(work_dir, scenario_id): desc = ARTIFACT_CATEGORIES[category]["desc"] lines.append(f"{subdir}/ — {desc}") return lines # Evaluation of both held-out splits (test + val), pooled for analysis. Shared by # the train and load scenarios so they stay in lockstep. _EVAL_ALL_SPLITS = [ "evaluate_regular", "evaluate_bayesian", "evaluate_regular_val", "evaluate_bayesian_val", ] _EVAL_ALL_SPLITS_NOISY = ["evaluate_noisy", "evaluate_noisy_val"] SCENARIOS: Dict[str, ScenarioEntry] = { "scen_train_all": { "label": "1. Train, Evaluate (test+val), Noise, & Analyze", "tasks": [ "load_data", "train_regular", "train_bayesian", "load_models", *_EVAL_ALL_SPLITS, *_EVAL_ALL_SPLITS_NOISY, "combine_evaluations", "run_analysis", ], }, "scen_load_all": { "label": "2. Load, Evaluate (test+val), Noise, & Analyze", "tasks": [ "load_data", "load_models", *_EVAL_ALL_SPLITS, *_EVAL_ALL_SPLITS_NOISY, "combine_evaluations", "run_analysis", ], }, "scen_load_eval": { "label": "3. Load & Evaluate (test+val, Skip Noise) & Analyze", "tasks": [ "load_data", "load_models", *_EVAL_ALL_SPLITS, "combine_evaluations", "run_analysis", ], }, "scen_analyze": { "label": "4. Analyze Existing Evaluations", "tasks": [ "run_analysis", ], }, # Set aside for now (kept working, not part of the default flow): full # patient-grouped cross-validation. Costs ~k x a normal train+eval run. "scen_kfold": { "label": "5. K-Fold Cross-Validation (train + OOF eval + analyze)", "tasks": [ "run_kfold", "run_analysis", ], }, }