giant/giant/workflow/spec.py

"""Workflow TOML -> frozen dataclasses, validation, and per-task spec hashes.

One spec file is the only place a pipeline is parameterised (see
``configs/workflow_example.toml``):

    [workflow]  name / result_dir / log_dir
    [condor]    accounting group, repo dir, env script, docker images
    [dataset]   steps (training) + reference (rollout seeds & analysis truth)
    [geometry]  geometry-oracle build options
    [[train]]   one per training run      (name, config, epochs, overrides, ...)
    [[rollout]] one per rollout           (name, train = <a [[train]].name>, ...)
    [[analysis]] one per comparison       (name, rollouts = [<[[rollout]].name>, ...])

Every task carries its ``name`` plus a short ``spec_hash`` — 8 hex of the
canonical JSON of its own resolved sub-spec **including its transitive
parents**. That is what makes an edited spec produce a fresh result directory
instead of silently reusing outputs computed under different settings: change
the dataset and every hash downstream of it changes too.

Unknown keys are rejected (with the valid ones listed), in the same spirit as
``giant.config.validate_config_keys`` — a typo in a workflow spec would
otherwise be a silently ignored setting on a multi-day pipeline.
"""

from __future__ import annotations

import hashlib
import json
import tomllib
from dataclasses import MISSING, dataclass, field, fields, is_dataclass
from pathlib import Path
from typing import Any

__all__ = [
    "AnalysisSpec",
    "CondorSpec",
    "DatasetSpec",
    "GeometrySpec",
    "RolloutSpec",
    "TrainSpec",
    "WorkflowSpec",
    "load_spec",
    "spec_hash",
]


class WorkflowSpecError(ValueError):
    """Raised for any malformed workflow spec (unknown key, bad reference, ...)."""


# ---------------------------------------------------------------------------
# sub-specs
# ---------------------------------------------------------------------------


@dataclass(frozen=True)
class CondorSpec:
    """Where and how jobs run — the batch-system half of the spec.

    ``repo_dir`` doubles as b2luigi's ``working_dir`` (jobs ``cd`` there before
    running ``giant/workflow/run.py``), and ``env_script`` is sourced first,
    since submit and worker machines don't share an environment.
    """

    accounting_group: str
    repo_dir: str
    env_script: str = ""
    docker_image_cpu: str = "cverstege/alma9-gridjob"
    docker_image_gpu: str = "mschnepf/slc7-condocker"
    remote: bool = True
    request_cpus: int = 1
    request_memory_mb: int = 8192


@dataclass(frozen=True)
class DatasetSpec:
    """The two datasets every pipeline needs.

    ``steps`` is what training reads; ``reference`` is the held-out file
    rollouts are seeded from and the analysis compares against (the "one
    ground truth" premise of ``giant.analysis``).
    """

    steps: str
    reference: str


@dataclass(frozen=True)
class GeometrySpec:
    """``dwarf build-geometry-oracle`` options (see giant/tools/geometry_oracle.py)."""

    method: str = "slab"
    k: int = 1
    subsample: int = 500_000
    escape_factor: float = 5.0
    seed: int = 0
    depth_axis: int = 2
    n_bins: int = 2000


@dataclass(frozen=True)
class TrainSpec:
    """One training run, fanned out into ``ceil(epochs / epochs_per_job)`` jobs.

    ``overrides`` are ``[train]``/model config keys merged on top of ``config``
    exactly as ``giant train``'s flags are (``giant.config.merge_cli_overrides``),
    so anything expressible on the CLI is expressible here.
    """

    name: str
    config: str | None = None
    epochs: int = 1
    epochs_per_job: int = 1
    overrides: dict[str, Any] = field(default_factory=dict)
    request_gpus: int = 1
    gpu_type: str | None = None
    gpu_memory_mb: int | None = None
    request_memory_mb: int = 16384
    request_cpus: int = 4
    walltime_s: int = 86400
    num_workers: int = 4
    shuffle_buffer: int = 65536
    device: str | None = None


@dataclass(frozen=True)
class RolloutSpec:
    """One ``giant rollout`` run against the checkpoint of ``train``."""

    name: str
    train: str
    n_events: int | None = None
    energy_cutoff: float = 0.1
    max_steps: int = 1000
    steps: int = 10
    batch_size: int = 4096
    max_tracks_per_event: int | None = None
    escape_threshold: float | None = None
    weights: str = "raw"
    seed: int | None = None
    request_gpus: int = 1
    gpu_type: str | None = None
    gpu_memory_mb: int | None = None
    request_memory_mb: int = 16384
    request_cpus: int = 2
    walltime_s: int = 86400
    device: str | None = None


@dataclass(frozen=True)
class AnalysisSpec:
    """One rollout-vs-reference comparison (N rollout series, one reference)."""

    name: str
    rollouts: tuple[str, ...]
    chunks: int = 1
    energy_bins: int = 4
    bins: int = 50
    top_pdg: int = 6
    gallery: bool = False
    request_memory_mb: int = 8192
    request_cpus: int = 1


@dataclass(frozen=True)
class WorkflowSpec:
    """A whole pipeline: the parsed spec file plus name-keyed lookups."""

    name: str
    result_dir: str
    log_dir: str
    condor: CondorSpec
    dataset: DatasetSpec
    geometry: GeometrySpec
    trains: tuple[TrainSpec, ...]
    rollouts: tuple[RolloutSpec, ...]
    analyses: tuple[AnalysisSpec, ...]
    path: str = ""

    # -- lookups ----------------------------------------------------------
    def train(self, name: str) -> TrainSpec:
        return _lookup(self.trains, name, "train")

    def rollout(self, name: str) -> RolloutSpec:
        return _lookup(self.rollouts, name, "rollout")

    def analysis(self, name: str) -> AnalysisSpec:
        return _lookup(self.analyses, name, "analysis")

    # -- hashes -----------------------------------------------------------
    # Each one folds in everything upstream of it, so a change anywhere in a
    # task's ancestry moves its result directory (and only the affected
    # subtree's).
    def dataset_hash(self) -> str:
        return spec_hash(self.dataset)

    def warm_cache_hash(self, train_name: str) -> str:
        # The setup cache depends on the dataset and on what this training's
        # config asks of it (val split, conditioning, router) — not on how
        # many epochs it runs for, so epochs/resources are deliberately left
        # out and two trainings sharing a config share one warm-cache job.
        t = self.train(train_name)
        return spec_hash(self.dataset, t.config, t.overrides)

    def geometry_hash(self) -> str:
        return spec_hash(self.dataset, self.geometry)

    def train_hash(self, name: str) -> str:
        return spec_hash(self.dataset, self.train(name))

    def rollout_hash(self, name: str) -> str:
        ro = self.rollout(name)
        return spec_hash(self.dataset, self.geometry, self.train(ro.train), ro)

    def analysis_hash(self, name: str) -> str:
        an = self.analysis(name)
        parents = [self.rollout(r) for r in an.rollouts]
        train_parents = [self.train(r.train) for r in parents]
        return spec_hash(self.dataset, self.geometry, train_parents, parents, an)


def _lookup(items, name: str, kind: str):
    for item in items:
        if item.name == name:
            return item
    known = ", ".join(sorted(i.name for i in items)) or "(none defined)"
    raise WorkflowSpecError(f"no [[{kind}]] named {name!r} in this workflow — defined: {known}")


# ---------------------------------------------------------------------------
# hashing
# ---------------------------------------------------------------------------


def spec_hash(*parts: Any, length: int = 8) -> str:
    """Short stable hash of one or more (sub-)specs.

    Canonical JSON (sorted keys, dataclasses expanded) so the value depends
    only on the resolved settings — not on key order in the TOML, nor on
    which defaults were written out explicitly.
    """
    payload = json.dumps([_canonical(p) for p in parts], sort_keys=True, separators=(",", ":"))
    return hashlib.sha256(payload.encode()).hexdigest()[:length]


def _canonical(value: Any) -> Any:
    if is_dataclass(value) and not isinstance(value, type):
        return {f.name: _canonical(getattr(value, f.name)) for f in fields(value)}
    if isinstance(value, dict):
        return {str(k): _canonical(v) for k, v in value.items()}
    if isinstance(value, (list, tuple)):
        return [_canonical(v) for v in value]
    if isinstance(value, Path):
        return str(value)
    return value


# ---------------------------------------------------------------------------
# parsing
# ---------------------------------------------------------------------------


def _build(cls, data: dict, where: str):
    """Instantiate a frozen sub-spec, rejecting unknown/missing keys loudly."""
    valid = {f.name for f in fields(cls)}
    unknown = sorted(set(data) - valid)
    if unknown:
        raise WorkflowSpecError(f"{where}: unknown key(s) {unknown} — valid keys: {sorted(valid)}")
    required = {f.name for f in fields(cls) if f.default is MISSING and f.default_factory is MISSING}
    missing = sorted(required - set(data))
    if missing:
        raise WorkflowSpecError(f"{where}: missing required key(s) {missing}")
    return cls(**data)


def load_spec(path: str | Path) -> WorkflowSpec:
    """Parse and validate a workflow TOML file."""
    path = Path(path)
    try:
        raw = tomllib.loads(path.read_text())
    except tomllib.TOMLDecodeError as exc:
        raise WorkflowSpecError(f"{path}: not valid TOML — {exc}") from exc
    return parse_spec(raw, path=path)


def parse_spec(raw: dict, path: str | Path = "") -> WorkflowSpec:
    """Validate an already-parsed workflow spec mapping."""
    top_valid = {"workflow", "condor", "dataset", "geometry", "train", "rollout", "analysis"}
    unknown = sorted(set(raw) - top_valid)
    if unknown:
        raise WorkflowSpecError(
            f"{path or '<spec>'}: unknown top-level table(s) {unknown} — valid: {sorted(top_valid)}"
        )

    for required in ("workflow", "condor", "dataset"):
        if required not in raw:
            raise WorkflowSpecError(f"{path or '<spec>'}: missing required [{required}] table")

    wf = dict(raw["workflow"])
    wf_valid = {"name", "result_dir", "log_dir"}
    wf_unknown = sorted(set(wf) - wf_valid)
    if wf_unknown:
        raise WorkflowSpecError(f"[workflow]: unknown key(s) {wf_unknown} — valid keys: {sorted(wf_valid)}")
    if "name" not in wf or "result_dir" not in wf:
        raise WorkflowSpecError("[workflow]: 'name' and 'result_dir' are required")
    result_dir = str(Path(wf["result_dir"]).expanduser())
    log_dir = str(Path(wf.get("log_dir", Path(result_dir) / "logs")).expanduser())

    condor = _build(CondorSpec, dict(raw["condor"]), "[condor]")
    dataset = _build(DatasetSpec, dict(raw["dataset"]), "[dataset]")
    geometry = _build(GeometrySpec, dict(raw.get("geometry", {})), "[geometry]")

    trains = tuple(_build(TrainSpec, dict(t), f"[[train]] #{i}") for i, t in enumerate(raw.get("train", [])))
    rollouts = tuple(_build(RolloutSpec, dict(r), f"[[rollout]] #{i}") for i, r in enumerate(raw.get("rollout", [])))
    analyses = tuple(
        _build(AnalysisSpec, {**a, "rollouts": tuple(a.get("rollouts", ()))}, f"[[analysis]] #{i}")
        for i, a in enumerate(raw.get("analysis", []))
    )

    _check_unique(trains, "train")
    _check_unique(rollouts, "rollout")
    _check_unique(analyses, "analysis")

    train_names = {t.name for t in trains}
    for ro in rollouts:
        if ro.train not in train_names:
            raise WorkflowSpecError(
                f"[[rollout]] {ro.name!r}: train={ro.train!r} names no [[train]] — defined: {sorted(train_names)}"
            )
    rollout_names = {r.name for r in rollouts}
    for an in analyses:
        if not an.rollouts:
            raise WorkflowSpecError(f"[[analysis]] {an.name!r}: 'rollouts' must name at least one [[rollout]]")
        for r in an.rollouts:
            if r not in rollout_names:
                raise WorkflowSpecError(
                    f"[[analysis]] {an.name!r}: rollout {r!r} is not defined — "
                    f"defined: {sorted(rollout_names) or '(none)'}"
                )
        if len(set(an.rollouts)) != len(an.rollouts):
            raise WorkflowSpecError(f"[[analysis]] {an.name!r}: repeated rollout name(s) in 'rollouts'")
        if an.chunks < 1:
            raise WorkflowSpecError(f"[[analysis]] {an.name!r}: chunks must be >= 1, got {an.chunks}")

    for t in trains:
        if t.epochs < 1:
            raise WorkflowSpecError(f"[[train]] {t.name!r}: epochs must be >= 1, got {t.epochs}")
        if t.epochs_per_job < 1:
            raise WorkflowSpecError(f"[[train]] {t.name!r}: epochs_per_job must be >= 1, got {t.epochs_per_job}")

    return WorkflowSpec(
        name=wf["name"],
        result_dir=result_dir,
        log_dir=log_dir,
        condor=condor,
        dataset=dataset,
        geometry=geometry,
        trains=trains,
        rollouts=rollouts,
        analyses=analyses,
        path=str(path),
    )


def _check_unique(items, kind: str) -> None:
    names = [i.name for i in items]
    dupes = sorted({n for n in names if names.count(n) > 1})
    if dupes:
        raise WorkflowSpecError(f"[[{kind}]] names must be unique — repeated: {dupes}")


def epoch_milestones(train: TrainSpec) -> list[int]:
    """Cumulative epoch counts, one per chained ``TrainEpochTask``.

    ``epochs_per_job`` trades queue waits against job length: with
    ``epochs=10, epochs_per_job=3`` this is ``[3, 6, 9, 10]``, i.e. job *k*
    resumes job *k-1*'s ``last.pt`` and trains up to its own milestone.
    """
    step = train.epochs_per_job
    milestones = list(range(step, train.epochs + 1, step))
    if not milestones or milestones[-1] != train.epochs:
        milestones.append(train.epochs)
    return milestones