alnn_rewrite/tasks/load_models.py

"""Discover saved ensemble checkpoints and record their paths in ``state``.

Training deletes each model from VRAM after saving it, so evaluation cannot reuse
in-memory models. This task scans ``<work_dir>/{normal,bayesian}_models/`` for
``model_N.pt`` files and stores ordered path lists; the evaluation tasks then load
one model at a time. Weights are NOT loaded here (that would defeat the memory
decoupling) -- only discovered and ordered.
"""

import pathlib as pl
import re
from typing import Any, Dict, List

from model.factory import KIND_BAYESIAN, KIND_NORMAL
from util.progress import ProgressTracker
from util.ui_logger import PipelineLogger

# Subdirectory (under work_dir) and state key for each model kind.
_MODEL_SOURCES = {
    KIND_NORMAL: ("normal_models", "normal_model_paths"),
    KIND_BAYESIAN: ("bayesian_models", "bayesian_model_paths"),
}

_MODEL_FILE_RE = re.compile(r"^model_(\d+)\.pt$")


def _discover_models(models_dir: pl.Path) -> List[pl.Path]:
    """Return ``model_N.pt`` files in ``models_dir`` sorted by N (natural order).

    Only the directory's top level is scanned, so intermediate checkpoints under
    ``intermediate_models/`` are ignored.
    """
    if not models_dir.is_dir():
        return []

    numbered: List[tuple[int, pl.Path]] = []
    for entry in models_dir.iterdir():
        if not entry.is_file():
            continue
        match = _MODEL_FILE_RE.match(entry.name)
        if match:
            numbered.append((int(match.group(1)), entry))

    numbered.sort(key=lambda pair: pair[0])
    return [path for _, path in numbered]


def load_models_task(
    track: ProgressTracker,
    log: PipelineLogger,
    config: Dict[str, Any],
    state: Dict[str, Any],
) -> None:
    work_dir = pl.Path(config["work_dir"])

    track.update(total=len(_MODEL_SOURCES), advance=0)
    total_found = 0
    for kind, (subdir, state_key) in _MODEL_SOURCES.items():
        paths = _discover_models(work_dir / subdir)
        state[state_key] = paths
        total_found += len(paths)
        if paths:
            log.info(f"Found {len(paths)} {kind} model(s) in {subdir}/.")
        else:
            log.info(f"No {kind} models found in {subdir}/ (evaluation will skip).")
        track.update(advance=1)

    if total_found == 0:
        log.error(
            "No saved models were discovered. Train first, or point the working "
            "directory at an experiment that contains trained models."
        )