alnn_rewrite/tasks/__init__.py

"""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 -> PLANNED (placeholder)
    5. evaluate_bayesian-> PLANNED (placeholder)
    6. evaluate_noisy   -> PLANNED (placeholder)
    load_models         -> PLANNED (placeholder); loads saved .pt ensembles from
                           disk into ``state`` so evaluation is decoupled from
                           training (training deletes models from VRAM after saving).
"""

from typing import Any, Dict

from util.progress import ProgressTracker
from util.ui_logger import PipelineLogger

from . import load_data
from . import train_bayesian
from . import train_normal


def _placeholder_task(planned: str):
    """Build a no-op task for a pipeline stage that is not implemented yet.

    It logs a clear warning and returns cleanly (advancing the pipeline) instead
    of crashing, so the scenario wiring can be exercised before the real
    evaluation logic lands.
    """

    def _task(
        tracker: ProgressTracker,
        logger: PipelineLogger,
        config: Dict[str, Any],
        state: Dict[str, Any],
    ) -> None:
        logger.info(f"[NOT IMPLEMENTED] {planned} — skipping (placeholder).")

    return _task


PIPELINE_TASKS = {
    "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": _placeholder_task(
            "Load trained normal/Bayesian ensembles from work_dir into state"
        ),
    },
    "evaluate_regular": {
        "task_name": "Evaluate Regular Models",
        "task_func": _placeholder_task(
            "Evaluate normal ensemble and save netCDF (step 4)"
        ),
    },
    "evaluate_bayesian": {
        "task_name": "Evaluate Bayesian Models",
        "task_func": _placeholder_task(
            "Evaluate Bayesian ensemble with MC uncertainty and save netCDF (step 5)"
        ),
    },
    "evaluate_noisy": {
        "task_name": "Evaluate Models on Noised Data",
        "task_func": _placeholder_task(
            "Evaluate both ensembles across Gaussian noise levels and save netCDF (step 6)"
        ),
    },
}

SCENARIOS = {
    "scen_train_all": {
        "label": "1. Train, Evaluate, & Noise Analysis",
        "tasks": [
            "load_data",
            "train_regular",
            "train_bayesian",
            "load_models",
            "evaluate_regular",
            "evaluate_bayesian",
            "evaluate_noisy",
        ],
    },
    "scen_load_all": {
        "label": "2. Load, Evaluate, & Noise Analysis",
        "tasks": [
            "load_data",
            "load_models",
            "evaluate_regular",
            "evaluate_bayesian",
            "evaluate_noisy",
        ],
    },
    "scen_load_eval": {
        "label": "3. Load & Evaluate (Skip Noise)",
        "tasks": [
            "load_data",
            "load_models",
            "evaluate_regular",
            "evaluate_bayesian",
        ],
    },
}