alnn_rewrite/tasks/__init__.py

import time
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 dummy_task(
    tracker: ProgressTracker,
    logger: PipelineLogger,
    config: Dict[str, Any],
    state: Dict[str, Any],
):
    stop_event = state["events"]["stop"]
    pause_event = state["events"]["pause"]

    steps = 5
    # The title is now set gracefully by the parent injecting the sub_tracker,
    # so we just initialize the total.
    tracker.update(total=steps, advance=0)

    logger.info("Initializing process...")

    for i in range(steps):
        if stop_event.is_set():
            raise InterruptedError("Pipeline execution stopped by user.")

        while pause_event.is_set():
            time.sleep(0.5)
            if stop_event.is_set():
                raise InterruptedError(
                    "Pipeline execution stopped by user while paused."
                )

        # Child Process (Sub Progress Tracker)
        sub_steps = 10
        sub_tracker = tracker.get_sub_tracker(f"Batch {i + 1}")
        sub_tracker.update(total=sub_steps, advance=0)

        for j in range(sub_steps):
            if stop_event.is_set():
                raise InterruptedError("Pipeline execution stopped by user.")
            time.sleep(0.05)
            sub_tracker.update(advance=1)  # Advance sub task

        tracker.update(advance=1)  # Advance main task (clears sub task)

        if i == 2:
            logger.info("Halfway through current task execution...")


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,
    },
    "evaluate_regular": {
        "task_name": "Evaluate Regular Models",
        "task_func": dummy_task,
    },
    "evaluate_bayesian": {
        "task_name": "Evaluate Bayesian Models",
        "task_func": dummy_task,
    },
}

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