From a482b04761688aeb6bf831ac9fa667bccdacaa54 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Wed, 26 Aug 2026 11:52:17 +0200 Subject: [PATCH] Add the giant/workflow b2luigi task graph (gitea #83) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit One workflow TOML now parameterises a whole experiment and `giant workflow run ` turns it into a b2luigi DAG whose targets are files on /ceph: nothing already produced is recomputed, every step waits for its inputs, and HTCondor submission/polling is b2luigi's job. - spec.py: workflow TOML -> frozen dataclasses with name-uniqueness and cross-reference validation, unknown keys rejected the way giant.config rejects them, and a short spec_hash per task that folds in its transitive parents — so an edited spec re-runs exactly the affected subtree. - htcondor.py: the CPU/GPU submit settings. The GPU requirement strings (ProvidesEtpCeph + optional device/memory pins) are ported from the condor-gpu-train-rollout branch rather than rewritten. - tasks.py: DatasetTask, WarmCacheTask, GeometryOracleTask, TrainEpochTask (one short GPU job per epoch, chained via --resume, which the training loop already supports unchanged), TrainTask (publishes best.pt/last.pt and a concatenated metrics.csv so downstream never sees the epoch fan-out), RolloutTask, AnalysisPrepTask, AnalysisComputeTask (one job per plot x chunk, walltime sized from run_meta.json at submit time), AnalysisRenderTask (always local — the only step importing plotstyle/LaTeX), WorkflowTask. Task bodies call the existing entry points; none of them reimplement anything. - run.py + `giant workflow run`: settings wiring and the script b2luigi re-executes on workers. add_filename_to_cmd is off because b2luigi passes only the script's basename, and --spec is forwarded via task_cmd_additional_args so a worker resolves the identical task graph. configs/workflow_example.toml is the documented starting point. Co-Authored-By: Claude Opus 5 --- configs/workflow_example.toml | 68 ++++ giant/cli.py | 44 +++ giant/workflow/__init__.py | 40 +++ giant/workflow/htcondor.py | 82 +++++ giant/workflow/run.py | 103 ++++++ giant/workflow/spec.py | 383 ++++++++++++++++++++ giant/workflow/tasks.py | 647 ++++++++++++++++++++++++++++++++++ pyproject.toml | 9 + tests/test_workflow_spec.py | 148 ++++++++ tests/test_workflow_tasks.py | 172 +++++++++ 10 files changed, 1696 insertions(+) create mode 100644 configs/workflow_example.toml create mode 100644 giant/workflow/__init__.py create mode 100644 giant/workflow/htcondor.py create mode 100644 giant/workflow/run.py create mode 100644 giant/workflow/spec.py create mode 100644 giant/workflow/tasks.py create mode 100644 tests/test_workflow_spec.py create mode 100644 tests/test_workflow_tasks.py diff --git a/configs/workflow_example.toml b/configs/workflow_example.toml new file mode 100644 index 0000000..488be23 --- /dev/null +++ b/configs/workflow_example.toml @@ -0,0 +1,68 @@ +# Example GIANT workflow spec — `giant workflow run configs/workflow_example.toml`. +# +# One file parameterises a whole experiment: the datasets, the geometry oracle, +# N trainings, N rollouts, and the analyses comparing them. Every task's output +# directory carries a hash of its resolved sub-spec (plus its parents), so +# editing anything here re-runs exactly the affected subtree and nothing else. +# +# result_dir/log_dir must be visible from both the submit host and the workers +# (i.e. on /ceph) — there is deliberately no HTCondor file transfer. + +[workflow] +name = "baseline-vs-router" +result_dir = "/ceph/lbogner/workflows/baseline-vs-router" +log_dir = "/ceph/lbogner/workflows/baseline-vs-router/logs" + +[condor] +accounting_group = "cms" +repo_dir = "/work/lbogner/giant" # also b2luigi's working_dir +env_script = "/work/lbogner/giant/condor_env.sh" +docker_image_cpu = "cverstege/alma9-gridjob" +docker_image_gpu = "mschnepf/slc7-condocker" +remote = true + +[dataset] +steps = "/ceph/lbogner/geant_steps/train/" # training data +reference = "/ceph/lbogner/geant_steps/holdout/" # rollout seeds + analysis truth + +[geometry] +method = "slab" +subsample = 500_000 + +[[train]] +name = "baseline" +config = "configs/baseline.toml" +epochs = 200 +epochs_per_job = 1 # one short GPU job per epoch, chained +request_gpus = 1 +gpu_memory_mb = 20000 +overrides = { lr = 3e-4 } # `giant train` flag names + +[[train]] +name = "router-balanced" +config = "configs/router.toml" +epochs = 200 +epochs_per_job = 1 +request_gpus = 1 +gpu_memory_mb = 20000 + +[[rollout]] +name = "baseline" +train = "baseline" # -> [[train]].name +n_events = 2000 +energy_cutoff = 0.1 + +[[rollout]] +name = "router-balanced" +train = "router-balanced" +n_events = 2000 +energy_cutoff = 0.1 + +[[analysis]] +name = "baseline-vs-router" +rollouts = ["baseline", "router-balanced"] +chunks = 32 +energy_bins = 4 +bins = 50 +top_pdg = 6 +gallery = true diff --git a/giant/cli.py b/giant/cli.py index c3058e3..3c47d8f 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -4,6 +4,7 @@ from enum import Enum import math from pathlib import Path import re +import sys from typing import Optional import uuid as uuid_mod @@ -1558,6 +1559,49 @@ def rollout( typer.echo(f"reference: {ref_path}") +workflow_app = typer.Typer( + no_args_is_help=True, + help="b2luigi pipeline orchestration: one spec file -> cache-warm, train, rollout, analysis.", +) +app.add_typer(workflow_app, name="workflow") + + +@workflow_app.command("run") +def workflow_run( + spec: Annotated[Path, typer.Argument(help="Workflow TOML (see configs/workflow_example.toml)")], + batch: Annotated[ + bool, + typer.Option("--batch/--local", help="Submit batch-system tasks to HTCondor, or run everything locally"), + ] = False, + workers: Annotated[int, typer.Option("--workers", help="Concurrent luigi workers")] = 1, + mode: Annotated[ + str, + typer.Option( + "--mode", + help="run | dry-run (print pending tasks) | show-output (print every target) | remove (delete outputs)", + ), + ] = "run", + scheduler_host: Annotated[Optional[str], typer.Option("--scheduler-host", help="luigid host")] = None, + scheduler_port: Annotated[Optional[int], typer.Option("--scheduler-port", help="luigid port")] = None, +) -> None: + """Run a workflow spec end to end (the only sanctioned multi-step entry point). + + A thin exec of `giant/workflow/run.py`, which b2luigi also re-executes on + every worker — so there is one documented entry point and one code path. + """ + import subprocess + + script = Path(__file__).resolve().parent / "workflow" / "run.py" + cmd = [sys.executable, str(script), "--spec", str(spec), "--workers", str(workers), "--mode", mode] + if batch: + cmd.append("--batch") + if scheduler_host: + cmd += ["--scheduler-host", scheduler_host] + if scheduler_port: + cmd += ["--scheduler-port", str(scheduler_port)] + raise typer.Exit(subprocess.run(cmd).returncode) + + analyze_app = typer.Typer( no_args_is_help=True, help="Rollout-vs-reference analysis: parallel compute on HTCondor + local render.", diff --git a/giant/workflow/__init__.py b/giant/workflow/__init__.py new file mode 100644 index 0000000..cfc57b3 --- /dev/null +++ b/giant/workflow/__init__.py @@ -0,0 +1,40 @@ +"""b2luigi orchestration of the full GIANT pipeline. + +One workflow TOML (``spec.py``) parameterises an entire experiment — dataset, +geometry oracle, N trainings, N rollouts, N analyses — and ``giant workflow +run `` turns it into a b2luigi task graph (``tasks.py``) whose +targets are files on ``/ceph``: nothing is recomputed that already exists, +every step waits for its inputs, and HTCondor submission/polling is b2luigi's +job rather than a hand-rolled submit-file generator. + +This is the only sanctioned way to run a multi-step pipeline; ``giant`` and +``dwarf`` stay single-step primitives that these tasks invoke. + +``tasks``/``run`` import b2luigi, so they are *not* imported here — a plain +``import giant.workflow`` (or ``giant.workflow.spec``) works without the +``workflow`` extra installed. +""" + +from giant.workflow.spec import ( + AnalysisSpec, + CondorSpec, + DatasetSpec, + GeometrySpec, + RolloutSpec, + TrainSpec, + WorkflowSpec, + load_spec, + spec_hash, +) + +__all__ = [ + "AnalysisSpec", + "CondorSpec", + "DatasetSpec", + "GeometrySpec", + "RolloutSpec", + "TrainSpec", + "WorkflowSpec", + "load_spec", + "spec_hash", +] diff --git a/giant/workflow/htcondor.py b/giant/workflow/htcondor.py new file mode 100644 index 0000000..e0cde5c --- /dev/null +++ b/giant/workflow/htcondor.py @@ -0,0 +1,82 @@ +"""HTCondor job descriptions for the workflow tasks. + +b2luigi writes every key of a task's ``htcondor_settings`` dict straight into +that job's submit description, so these helpers are just the ETP-specific +resource/requirement conventions in one place: + +* **CPU jobs** (setup cache, geometry oracle, analysis compute) keep what + ``giant analyze submit`` used: ``+RemoteJob`` for grid I/O, or + ``TARGET.ProvidesETPResources`` when the files are local to the cluster. +* **GPU jobs** (training epochs, rollout) are remote-only, so they always + carry ``+RemoteJob`` and reach ``/ceph`` through + ``TARGET.ProvidesEtpCeph`` — the requirement strings are ported from the + ``condor-gpu-train-rollout`` branch's ``giant/condor.py`` rather than + rewritten, since they encode what the ETP HTCondor wiki documents for + TOpAS/NEMO2 GPU workers. +""" + +from __future__ import annotations + +from giant.workflow.spec import CondorSpec + +__all__ = ["cpu_settings", "gpu_settings", "gpu_requirements"] + + +def cpu_settings( + condor: CondorSpec, + *, + request_memory_mb: int | None = None, + request_cpus: int | None = None, + walltime_s: int | None = None, +) -> dict: + settings: dict = { + "universe": "docker", + "docker_image": condor.docker_image_cpu, + "request_memory": request_memory_mb if request_memory_mb is not None else condor.request_memory_mb, + "request_cpus": request_cpus if request_cpus is not None else condor.request_cpus, + "accounting_group": condor.accounting_group, + "should_transfer_files": "YES", + "when_to_transfer_output": "ON_EXIT", + } + if condor.remote: + settings["+RemoteJob"] = "True" + else: + settings["requirements"] = "TARGET.ProvidesETPResources" + if walltime_s is not None: + settings["+RequestWalltime"] = int(walltime_s) + return settings + + +def gpu_requirements(gpu_type: str | None = None, gpu_memory_mb: int | None = None) -> str: + """``TARGET.ProvidesEtpCeph`` (remote /ceph access) ANDed with any GPU pin.""" + clauses = ["TARGET.ProvidesEtpCeph =?= True"] + if gpu_type is not None: + clauses.append(f'TARGET.GPUs_DeviceName =?= "{gpu_type}"') + if gpu_memory_mb is not None: + clauses.append(f"TARGET.GPUs_GlobalMemoryMb >= {gpu_memory_mb}") + return " && ".join(clauses) + + +def gpu_settings( + condor: CondorSpec, + *, + 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, +) -> dict: + return { + "universe": "docker", + "docker_image": condor.docker_image_gpu, + "request_memory": request_memory_mb, + "request_cpus": request_cpus, + "RequestGPUs": request_gpus, + "+RequestWalltime": int(walltime_s), + "accounting_group": condor.accounting_group, + "should_transfer_files": "YES", + "when_to_transfer_output": "ON_EXIT", + "+RemoteJob": "True", + "requirements": f"({gpu_requirements(gpu_type, gpu_memory_mb)})", + } diff --git a/giant/workflow/run.py b/giant/workflow/run.py new file mode 100644 index 0000000..e6d82ec --- /dev/null +++ b/giant/workflow/run.py @@ -0,0 +1,103 @@ +#!/usr/bin/env python +"""Entry point b2luigi re-executes on every worker. + +Locally this is what ``giant workflow run `` execs; on a batch +worker it is what the generated wrapper script runs (after ``cd repo_dir`` and +sourcing ``env_script``), with ``--spec`` forwarded via the +``task_cmd_additional_args`` setting so the worker resolves exactly the same +spec — and therefore the same task graph and output paths — as the submitter. + +b2luigi needs a real script path for that re-execution, which is why this is a +script rather than a ``python -m`` module. +""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +# Allow `python giant/workflow/run.py` from a checkout that isn't installed. +sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + +import b2luigi # noqa: E402 + +from giant.workflow.spec import WorkflowSpec, load_spec # noqa: E402 +from giant.workflow.tasks import WorkflowTask, set_spec # noqa: E402 + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description="Run a GIANT workflow spec with b2luigi.") + parser.add_argument("--spec", required=True, help="Workflow TOML (see configs/workflow_example.toml)") + parser.add_argument("--workers", type=int, default=1, help="Concurrent luigi workers") + parser.add_argument( + "--batch", + action="store_true", + help="Submit batch-system tasks to HTCondor (otherwise everything runs locally)", + ) + parser.add_argument( + "--mode", + choices=("run", "dry-run", "show-output", "remove"), + default="run", + help="run (default), dry-run (print pending tasks), show-output (print every target), remove (delete outputs)", + ) + parser.add_argument("--scheduler-host", default=None, help="luigid host (default: local scheduler)") + parser.add_argument("--scheduler-port", type=int, default=None, help="luigid port") + return parser + + +def configure(spec: WorkflowSpec, spec_path: Path, batch: bool) -> None: + """Wire b2luigi's settings from the spec. + + ``/ceph`` is shared between submit host and workers, so there is + deliberately no ``transfer_files``: ``result_dir``/``log_dir`` must live + somewhere both sides can see. + """ + set_spec(spec) + + b2luigi.set_setting("result_dir", spec.result_dir) + b2luigi.set_setting("log_dir", spec.log_dir) + b2luigi.set_setting("task_file_dir", str(Path(spec.result_dir) / "task_files")) + b2luigi.set_setting("use_parameter_name_in_output", True) + b2luigi.set_setting("batch_system", "htcondor" if batch else "local") + b2luigi.set_setting("working_dir", spec.condor.repo_dir) + b2luigi.set_setting("job_name", spec.name) + if spec.condor.env_script: + b2luigi.set_setting("env_script", spec.condor.env_script) + # The worker command is ` [] --batch-runner + # --task-id ...`, run after `cd working_dir`. Only the *basename* would be + # used, so the filename is dropped and the repo-relative script path is + # made part of the executable instead. + b2luigi.set_setting("add_filename_to_cmd", False) + b2luigi.set_setting("executable", [".venv/bin/python", "giant/workflow/run.py"]) + b2luigi.set_setting("task_cmd_additional_args", ["--spec", str(spec_path)]) + + +def main(argv: list[str] | None = None) -> None: + args, _ = build_parser().parse_known_args(argv) + spec_path = Path(args.spec).resolve() + spec = load_spec(spec_path) + configure(spec, spec_path, batch=args.batch) + + kwargs: dict = {} + if args.scheduler_host: + kwargs["scheduler_host"] = args.scheduler_host + if args.scheduler_port: + kwargs["scheduler_port"] = args.scheduler_port + + b2luigi.process( + WorkflowTask(workflow_name=spec.name), + workers=args.workers, + batch=args.batch, + dry_run=args.mode == "dry-run", + show_output=args.mode == "show-output", + remove=args.mode == "remove", + auto_confirm=args.mode == "remove", + # run.py owns --spec/--mode/...; b2luigi must not choke on them. + ignore_additional_command_line_args=True, + **kwargs, + ) + + +if __name__ == "__main__": + main() diff --git a/giant/workflow/spec.py b/giant/workflow/spec.py new file mode 100644 index 0000000..48bb461 --- /dev/null +++ b/giant/workflow/spec.py @@ -0,0 +1,383 @@ +"""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 = , ...) + [[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 ''}: 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 ''}: 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 diff --git a/giant/workflow/tasks.py b/giant/workflow/tasks.py new file mode 100644 index 0000000..8115500 --- /dev/null +++ b/giant/workflow/tasks.py @@ -0,0 +1,647 @@ +"""The b2luigi task graph: cache-warm -> train -> rollout -> analysis. + + DatasetTask (external) ─┬─> WarmCacheTask(train) ─> TrainEpochTask(train, 1..N) ─> TrainTask(train) ─┐ + └─> GeometryOracleTask ──┐ │ + └──> RolloutTask(rollout) <────────────────────────┘ + │ + AnalysisPrepTask(analysis) ─> AnalysisComputeTask(analysis, plot, chunk) ─> AnalysisRenderTask(analysis) + ^ + WorkflowTask (wrapper) ─────────────────────────────────────────────────────────────┘ + +Every task's output directory is ``//name=/spec_hash= +/…`` — the hash covers the task's resolved sub-spec *and its transitive +parents* (``giant/workflow/spec.py``), so editing the spec produces a fresh +directory for exactly the affected subtree instead of silently reusing stale +outputs. + +Task bodies never reimplement anything: they call the same entry points the +CLIs do (``run_warm_setup_cache``, ``run_build_geometry_oracle``, +``run_train_job``, ``giant.analysis.prep``/``compute_one``/``merge_all``, +``render_run``), or shell out to ``giant rollout``, which has no library-level +entry point of its own. + +Training is fanned out into **one short GPU job per epoch** (or per +``epochs_per_job`` epochs): job *k* runs ``run_train_job`` with ``epochs = k`` +and ``resume = /last.pt``, which the training loop already handles +(``giant/training/loop.py`` sets ``start_epoch = ckpt["epoch"] + 1`` and +returns early when the checkpoint already covers ``epochs``). A 200-epoch run +then becomes 200 schedulable jobs that survive preemption and give luigi a +real progress signal, at the cost of one (cache-warmed) setup scan and one +queue wait per job. +""" + +from __future__ import annotations + +import json +import shutil +import subprocess +import sys +from pathlib import Path + +import b2luigi + +from giant.workflow.htcondor import cpu_settings, gpu_settings +from giant.workflow.spec import WorkflowSpec, epoch_milestones + +__all__ = [ + "AnalysisComputeTask", + "AnalysisPrepTask", + "AnalysisRenderTask", + "DatasetTask", + "GeometryOracleTask", + "RolloutTask", + "TrainEpochTask", + "TrainTask", + "WarmCacheTask", + "WorkflowTask", + "analysis_dir", + "analysis_jobs", + "get_spec", + "set_spec", +] + + +# --------------------------------------------------------------------------- +# the active spec +# --------------------------------------------------------------------------- +# luigi parameters must be simple scalars, so tasks carry only `name` + +# `spec_hash` and read the rest out of the one spec this process was started +# with. Batch workers re-execute `run.py --spec ` (see +# `task_cmd_additional_args` there), so they resolve the identical spec. + +_SPEC: WorkflowSpec | None = None + + +def set_spec(spec: WorkflowSpec) -> None: + global _SPEC + _SPEC = spec + + +def get_spec() -> WorkflowSpec: + if _SPEC is None: + raise RuntimeError("no workflow spec loaded — call giant.workflow.tasks.set_spec() first") + return _SPEC + + +def _result_dir(*parts: str) -> Path: + return Path(get_spec().result_dir).joinpath(*parts) + + +def _task_dir(kind: str, name: str, spec_hash: str) -> Path: + """``//name=/spec_hash=``.""" + return _result_dir(kind, f"name={name}", f"spec_hash={spec_hash}") + + +def analysis_dir(spec: WorkflowSpec, name: str) -> Path: + """The analysis run directory — what ``prep`` lays out and every later step reads.""" + return Path(spec.result_dir) / "analysis" / f"name={name}" / f"spec_hash={spec.analysis_hash(name)}" + + +def analysis_jobs(spec: WorkflowSpec, name: str) -> list[tuple[str, int]]: + """Every ``(plot_id, chunk)`` compute job of one analysis. + + ``chunkable=False`` specs (the checkpoint-bound diagnostics, already + bounded/subsampled) always run as a single chunk — the same rule the + deleted ``_job_walltimes`` applied. + """ + from giant.analysis.catalog import catalog_ids, get_spec as get_plot_spec + + chunks = spec.analysis(name).chunks + jobs: list[tuple[str, int]] = [] + for plot_id in catalog_ids(): + n = chunks if get_plot_spec(plot_id).chunkable else 1 + jobs.extend((plot_id, chunk) for chunk in range(n)) + return jobs + + +def _giant_cmd() -> list[str]: + """How to invoke the ``giant`` CLI from inside a task (worker or locally).""" + return [sys.executable, "-m", "giant.cli"] + + +# --------------------------------------------------------------------------- +# inputs +# --------------------------------------------------------------------------- + + +class DatasetTask(b2luigi.ExternalTask): + """A steps parquet file or directory that must already exist. + + Nothing produces it, so a missing path is a hard, immediate error rather + than a job that fails hours later — the usual cause being ``/ceph`` not + mounted on the machine the workflow was started from. + """ + + path = b2luigi.Parameter() + + def output(self): + return b2luigi.LocalTarget(str(self.path)) + + def complete(self): + if not Path(str(self.path)).exists(): + raise FileNotFoundError( + f"dataset {self.path!r} does not exist — is /ceph mounted on this machine? " + "(see CLAUDE.md's Compute environment section)" + ) + return True + + +# --------------------------------------------------------------------------- +# setup stage +# --------------------------------------------------------------------------- + + +class WarmCacheTask(b2luigi.Task): + """Precompute one training's setup-stage sidecar (vocab maps, event split, + normalizer stats) so every per-epoch job is a cache hit instead of a + full rescan. + + The real product (``.giant_train_cache.json``) lives next to the + dataset, not under ``result_dir``, so the target here is a small stamp + recording that sidecar's path/mtime/size. + """ + + name = b2luigi.Parameter() + spec_hash = b2luigi.Parameter() + + @property + def htcondor_settings(self): + spec = get_spec() + return cpu_settings(spec.condor, request_memory_mb=32768, request_cpus=4, walltime_s=21600) + + def requires(self): + yield DatasetTask(path=get_spec().dataset.steps) + + def output(self): + return b2luigi.LocalTarget(str(_task_dir("warm_cache", str(self.name), str(self.spec_hash)) / "stamp.json")) + + def run(self): + from giant.data.setup_cache import sidecar_path + from giant.tools.warm_setup_cache import run_warm_setup_cache + + spec = get_spec() + train = spec.train(str(self.name)) + run_warm_setup_cache( + data=spec.dataset.steps, + config_path=Path(train.config) if train.config else None, + ) + sidecar = Path(sidecar_path(spec.dataset.steps)) + stamp = { + "sidecar": str(sidecar), + "mtime": sidecar.stat().st_mtime if sidecar.exists() else None, + "size": sidecar.stat().st_size if sidecar.exists() else None, + } + out = Path(self.output().path) + out.parent.mkdir(parents=True, exist_ok=True) + out.write_text(json.dumps(stamp, indent=2)) + + +class GeometryOracleTask(b2luigi.Task): + """Build the position -> (material, layer_id) oracle every rollout needs.""" + + spec_hash = b2luigi.Parameter() + + @property + def htcondor_settings(self): + spec = get_spec() + return cpu_settings(spec.condor, request_memory_mb=32768, request_cpus=4, walltime_s=21600) + + def requires(self): + yield DatasetTask(path=get_spec().dataset.steps) + + def output(self): + return b2luigi.LocalTarget( + str(_result_dir("geometry", f"spec_hash={self.spec_hash}") / "oracle.pkl"), + ) + + def run(self): + from giant.tools.geometry_oracle import run_build_geometry_oracle + + spec = get_spec() + g = spec.geometry + out = Path(self.output().path) + out.parent.mkdir(parents=True, exist_ok=True) + run_build_geometry_oracle( + data=Path(spec.dataset.steps), + out=out, + method=g.method, + k=g.k, + subsample=g.subsample, + escape_factor=g.escape_factor, + seed=g.seed, + depth_axis=g.depth_axis, + n_bins=g.n_bins, + ) + + +# --------------------------------------------------------------------------- +# training +# --------------------------------------------------------------------------- + + +def _train_cfg(spec: WorkflowSpec, name: str, epochs: int) -> dict: + """The merged config one training job runs, resolved exactly as `giant train` does.""" + from giant import config as gconfig + + train = spec.train(name) + flags = {**train.overrides, "epochs": epochs} + overrides = gconfig.overrides_from_flags(flags) + cfg = gconfig.merge_cli_overrides( + gconfig.DEFAULT_CONFIG, + Path(train.config) if train.config else None, + overrides, + ) + gconfig.validate_config(cfg, resume=True) + return cfg + + +class TrainEpochTask(b2luigi.Task): + """Epochs up to ``milestone`` of one training, resuming the previous job. + + Target is ``last.pt``. ``best.pt`` is written by the loop *only when that + epoch improved*, and ``best_val_loss`` travels inside the checkpoint, so + the global best comparison stays correct across jobs: "``best.pt`` exists + in milestone dir *k*" means exactly "one of that job's epochs was the best + so far". + """ + + name = b2luigi.Parameter() + spec_hash = b2luigi.Parameter() + milestone = b2luigi.IntParameter() + + @property + def htcondor_settings(self): + spec = get_spec() + train = spec.train(str(self.name)) + return gpu_settings( + spec.condor, + request_gpus=train.request_gpus, + gpu_type=train.gpu_type, + gpu_memory_mb=train.gpu_memory_mb, + request_memory_mb=train.request_memory_mb, + request_cpus=train.request_cpus, + walltime_s=train.walltime_s, + ) + + @property + def _dir(self) -> Path: + return _task_dir("train_epoch", str(self.name), str(self.spec_hash)) / f"epochs={int(self.milestone)}" + + def _previous_milestone(self) -> int | None: + spec = get_spec() + milestones = epoch_milestones(spec.train(str(self.name))) + index = milestones.index(int(self.milestone)) + return milestones[index - 1] if index > 0 else None + + def requires(self): + previous = self._previous_milestone() + if previous is None: + yield WarmCacheTask(name=self.name, spec_hash=get_spec().warm_cache_hash(str(self.name))) + else: + yield TrainEpochTask(name=self.name, spec_hash=self.spec_hash, milestone=previous) + + def output(self): + return b2luigi.LocalTarget(str(self._dir / "last.pt")) + + def run(self): + import torch + + from giant import config as gconfig + from giant.pipeline import run_train_job + + spec = get_spec() + train = spec.train(str(self.name)) + cfg = _train_cfg(spec, str(self.name), int(self.milestone)) + + previous = self._previous_milestone() + resume = None + if previous is not None: + resume = _task_dir("train_epoch", str(self.name), str(self.spec_hash)) / f"epochs={previous}" / "last.pt" + + device = torch.device(train.device) if train.device else gconfig.auto_device() + out_dir = self._dir + out_dir.mkdir(parents=True, exist_ok=True) + run_train_job( + data=Path(spec.dataset.steps), + cfg=cfg, + out_dir=out_dir, + device=device, + shuffle_buffer=train.shuffle_buffer, + num_workers=train.num_workers, + resume=resume, + cache_setup=True, + ) + + +class TrainTask(b2luigi.Task): + """Publish one training's canonical outputs, hiding the epoch fan-out. + + Everything downstream (``RolloutTask``, humans, ``giant analyze metrics``) + points here and never has to know which milestone directory happened to + hold the best checkpoint. + """ + + name = b2luigi.Parameter() + spec_hash = b2luigi.Parameter() + batch_system = "local" + + @property + def _milestones(self) -> list[int]: + return epoch_milestones(get_spec().train(str(self.name))) + + def requires(self): + yield TrainEpochTask(name=self.name, spec_hash=self.spec_hash, milestone=self._milestones[-1]) + + @property + def _dir(self) -> Path: + return _task_dir("train", str(self.name), str(self.spec_hash)) + + def output(self): + d = self._dir + return { + "best.pt": b2luigi.LocalTarget(str(d / "best.pt")), + "last.pt": b2luigi.LocalTarget(str(d / "last.pt")), + "metrics.csv": b2luigi.LocalTarget(str(d / "metrics.csv")), + } + + def run(self): + epoch_base = _task_dir("train_epoch", str(self.name), str(self.spec_hash)) + milestone_dirs = [epoch_base / f"epochs={m}" for m in self._milestones] + + best_dirs = [d for d in milestone_dirs if (d / "best.pt").exists()] + if not best_dirs: + raise FileNotFoundError( + f"no best.pt in any milestone directory under {epoch_base} — " + "did every epoch job run with a validation split?" + ) + out = self._dir + out.mkdir(parents=True, exist_ok=True) + shutil.copy2(best_dirs[-1] / "best.pt", out / "best.pt") + shutil.copy2(milestone_dirs[-1] / "last.pt", out / "last.pt") + for extra in ("config.toml", "run_meta.json"): + src = milestone_dirs[-1] / extra + if src.exists(): + shutil.copy2(src, out / extra) + + # One metrics.csv for the whole run: the first job's header, then + # every job's rows in epoch order, so `giant analyze metrics` sees a + # single continuous training curve. + lines: list[str] = [] + header: str | None = None + for d in milestone_dirs: + csv = d / "metrics.csv" + if not csv.exists(): + continue + rows = csv.read_text().splitlines() + if not rows: + continue + if header is None: + header = rows[0] + lines.extend(rows[1:]) + (out / "metrics.csv").write_text("\n".join([header or ""] + lines) + "\n") + + +# --------------------------------------------------------------------------- +# rollout +# --------------------------------------------------------------------------- + + +class RolloutTask(b2luigi.Task): + """Roll one trained checkpoint forward into full showers. + + ``giant rollout`` has no library-level entry point, so this shells out to + the CLI — with an explicit ``--out``, which puts the YAML sidecar at the + deterministic ``rollout.yaml`` next to the parquet (see + ``giant/cli.py:_write_prediction_ref``). + """ + + name = b2luigi.Parameter() + spec_hash = b2luigi.Parameter() + + @property + def htcondor_settings(self): + spec = get_spec() + ro = spec.rollout(str(self.name)) + return gpu_settings( + spec.condor, + request_gpus=ro.request_gpus, + gpu_type=ro.gpu_type, + gpu_memory_mb=ro.gpu_memory_mb, + request_memory_mb=ro.request_memory_mb, + request_cpus=ro.request_cpus, + walltime_s=ro.walltime_s, + ) + + @property + def _dir(self) -> Path: + return _task_dir("rollout", str(self.name), str(self.spec_hash)) + + def requires(self): + spec = get_spec() + ro = spec.rollout(str(self.name)) + yield TrainTask(name=ro.train, spec_hash=spec.train_hash(ro.train)) + yield GeometryOracleTask(spec_hash=spec.geometry_hash()) + yield DatasetTask(path=spec.dataset.reference) + + def output(self): + d = self._dir + return { + "rollout.parquet": b2luigi.LocalTarget(str(d / "rollout.parquet")), + "rollout.yaml": b2luigi.LocalTarget(str(d / "rollout.yaml")), + } + + def run(self): + spec = get_spec() + ro = spec.rollout(str(self.name)) + out = self._dir / "rollout.parquet" + out.parent.mkdir(parents=True, exist_ok=True) + + checkpoint = _task_dir("train", ro.train, spec.train_hash(ro.train)) / "best.pt" + oracle = _result_dir("geometry", f"spec_hash={spec.geometry_hash()}") / "oracle.pkl" + + cmd = [ + *_giant_cmd(), + "rollout", + spec.dataset.reference, + "--checkpoint", + str(checkpoint), + "--geometry", + str(oracle), + "--out", + str(out), + "--energy-cutoff", + str(ro.energy_cutoff), + "--max-steps", + str(ro.max_steps), + "--steps", + str(ro.steps), + "--batch-size", + str(ro.batch_size), + "--weights", + ro.weights, + ] + for flag, value in ( + ("--n-events", ro.n_events), + ("--max-tracks-per-event", ro.max_tracks_per_event), + ("--escape-threshold", ro.escape_threshold), + ("--seed", ro.seed), + ("--device", ro.device), + ): + if value is not None: + cmd += [flag, str(value)] + subprocess.run(cmd, check=True) + + +# --------------------------------------------------------------------------- +# analysis +# --------------------------------------------------------------------------- + + +class AnalysisPrepTask(b2luigi.Task): + """Resolve the shared bin edges/group sets once, for every compute job. + + Cheap and streaming, so it runs locally: everything after it needs + ``shared.json``/``run_meta.json`` to already exist. + """ + + name = b2luigi.Parameter() + spec_hash = b2luigi.Parameter() + batch_system = "local" + + def requires(self): + spec = get_spec() + for rollout_name in spec.analysis(str(self.name)).rollouts: + yield RolloutTask(name=rollout_name, spec_hash=spec.rollout_hash(rollout_name)) + + @property + def _dir(self) -> Path: + return analysis_dir(get_spec(), str(self.name)) + + def output(self): + d = self._dir + return { + "shared.json": b2luigi.LocalTarget(str(d / "shared.json")), + "run_meta.json": b2luigi.LocalTarget(str(d / "run_meta.json")), + } + + def run(self): + from giant.analysis import prep + + spec = get_spec() + an = spec.analysis(str(self.name)) + yamls = [_task_dir("rollout", r, spec.rollout_hash(r)) / "rollout.yaml" for r in an.rollouts] + prep( + yamls, + run_dir=self._dir, + n_chunks=an.chunks, + labels=list(an.rollouts), + n_energy_bins=an.energy_bins, + n_marginal_bins=an.bins, + top_k_pdg=an.top_pdg, + ) + + +class AnalysisComputeTask(b2luigi.Task): + """One (plot, chunk) streaming reduction — the replaced ``jobs.txt`` row. + + The output path is the on-disk contract ``compute-one``/``merge_one`` + already share (``reduced_partial/__.json``), declared + explicitly rather than through b2luigi's own output naming so that + contract is untouched. + """ + + name = b2luigi.Parameter() + spec_hash = b2luigi.Parameter() + plot_id = b2luigi.Parameter() + chunk = b2luigi.IntParameter() + + @property + def htcondor_settings(self): + # A property, so it is evaluated at submit time — i.e. after prep has + # written run_meta.json, whose row counts size the walltime request. + from giant.analysis import RunMeta + from giant.analysis.runtime_estimate import estimate_runtime_s + + spec = get_spec() + an = spec.analysis(str(self.name)) + walltime = None + meta_path = analysis_dir(spec, str(self.name)) / "run_meta.json" + if meta_path.exists(): + from giant.analysis.catalog import get_spec as get_plot_spec + + meta = RunMeta.load(meta_path) + chunkable = get_plot_spec(str(self.plot_id)).chunkable + n_rows = meta.rows_per_chunk[int(self.chunk)] if chunkable and meta.rows_per_chunk else meta.total_rows + walltime = estimate_runtime_s(str(self.plot_id), n_rows) + return cpu_settings( + spec.condor, + request_memory_mb=an.request_memory_mb, + request_cpus=an.request_cpus, + walltime_s=walltime, + ) + + def requires(self): + yield AnalysisPrepTask(name=self.name, spec_hash=self.spec_hash) + + def output(self): + run_dir = analysis_dir(get_spec(), str(self.name)) + return b2luigi.LocalTarget(str(run_dir / "reduced_partial" / f"{self.plot_id}__{int(self.chunk)}.json")) + + def run(self): + from giant.analysis import compute_one + + compute_one(str(self.plot_id), analysis_dir(get_spec(), str(self.name)), chunk_index=int(self.chunk)) + + +class AnalysisRenderTask(b2luigi.Task): + """Merge every plot's chunk partials, then render the PDFs + gallery. + + Always local — this is the only step that imports plotstyle/LaTeX, which + the compute worker images don't have. + """ + + name = b2luigi.Parameter() + spec_hash = b2luigi.Parameter() + batch_system = "local" + + def requires(self): + spec = get_spec() + for plot_id, chunk in analysis_jobs(spec, str(self.name)): + yield AnalysisComputeTask( + name=self.name, + spec_hash=self.spec_hash, + plot_id=plot_id, + chunk=chunk, + ) + + def output(self): + run_dir = analysis_dir(get_spec(), str(self.name)) + return b2luigi.LocalTarget(str(run_dir / "plots" / "metadata.yaml")) + + def run(self): + # render_run joins every plot's chunk partials (merge_all) before + # rendering, so this one call is the whole merge+render step. + from giant.analysis.render import render_run + + spec = get_spec() + render_run(analysis_dir(spec, str(self.name)), run_gallery=spec.analysis(str(self.name)).gallery) + + +class WorkflowTask(b2luigi.WrapperTask): + """The whole pipeline: every analysis in the spec, rendered.""" + + workflow_name = b2luigi.Parameter() + + def requires(self): + spec = get_spec() + if not spec.analyses: + # A spec with no [[analysis]] still has work to do — fall back to + # the deepest tasks it does define. + for ro in spec.rollouts: + yield RolloutTask(name=ro.name, spec_hash=spec.rollout_hash(ro.name)) + if not spec.rollouts: + for tr in spec.trains: + yield TrainTask(name=tr.name, spec_hash=spec.train_hash(tr.name)) + return + for an in spec.analyses: + yield AnalysisRenderTask(name=an.name, spec_hash=spec.analysis_hash(an.name)) diff --git a/pyproject.toml b/pyproject.toml index fb9455d..391505f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -108,3 +108,12 @@ explicit = true name = "larsbogner" url = "https://git.larsbogner.de/api/packages/lars/pypi/simple/" explicit = true + +# luigi builds task constructors from class-level Parameter descriptors, so a +# static checker sees no keyword parameters at all on `Task(name=..., ...)`. +# The workflow code is written against that API; nothing else in the repo is. +[[tool.ty.overrides]] +include = ["giant/workflow/**", "tests/test_workflow_tasks.py"] + +[tool.ty.overrides.rules] +unknown-argument = "ignore" diff --git a/tests/test_workflow_spec.py b/tests/test_workflow_spec.py new file mode 100644 index 0000000..914f687 --- /dev/null +++ b/tests/test_workflow_spec.py @@ -0,0 +1,148 @@ +"""Workflow spec parsing, validation, and spec hashes (gitea #83).""" + +import pytest + +from giant.workflow.spec import ( + WorkflowSpecError, + epoch_milestones, + load_spec, + parse_spec, + spec_hash, +) + +MINIMAL = { + "workflow": {"name": "wf", "result_dir": "/tmp/wf"}, + "condor": {"accounting_group": "cms", "repo_dir": "/work/lbogner/giant"}, + "dataset": {"steps": "/data/train", "reference": "/data/holdout"}, + "train": [{"name": "a", "epochs": 3}], + "rollout": [{"name": "a", "train": "a"}], + "analysis": [{"name": "cmp", "rollouts": ["a"], "chunks": 4}], +} + + +def _spec(**patch): + raw = {k: (v.copy() if isinstance(v, dict) else list(v)) for k, v in MINIMAL.items()} + raw.update(patch) + return parse_spec(raw) + + +def test_parses_minimal_spec(): + spec = _spec() + assert spec.name == "wf" + assert spec.log_dir == "/tmp/wf/logs" # derived from result_dir + assert spec.train("a").epochs == 3 + assert spec.rollout("a").train == "a" + assert spec.analysis("cmp").rollouts == ("a",) + # defaults come from the dataclasses, not the file + assert spec.geometry.method == "slab" + assert spec.condor.docker_image_gpu == "mschnepf/slc7-condocker" + + +def test_example_config_is_valid(): + spec = load_spec("configs/workflow_example.toml") + assert {t.name for t in spec.trains} == {"baseline", "router-balanced"} + assert spec.analysis("baseline-vs-router").rollouts == ("baseline", "router-balanced") + + +@pytest.mark.parametrize( + "patch, message", + [ + ({"train": [{"name": "a"}, {"name": "a"}]}, "unique"), + ({"rollout": [{"name": "r", "train": "nope"}]}, "names no"), + ({"analysis": [{"name": "c", "rollouts": ["nope"]}]}, "not defined"), + ({"analysis": [{"name": "c", "rollouts": []}]}, "at least one"), + ({"analysis": [{"name": "c", "rollouts": ["a"], "chunks": 0}]}, "chunks must be"), + ({"train": [{"name": "a", "epochs": 0}]}, "epochs must be"), + ({"train": [{"name": "a", "epchs": 3}]}, "unknown key"), + ({"geometry": {"methd": "slab"}}, "unknown key"), + ], +) +def test_validation_errors(patch, message): + with pytest.raises(WorkflowSpecError, match=message): + _spec(**patch) + + +def test_unknown_top_level_table_rejected(): + with pytest.raises(WorkflowSpecError, match="unknown top-level"): + _spec(nonsense={}) + + +def test_missing_required_table_rejected(): + raw = {k: v for k, v in MINIMAL.items() if k != "dataset"} + with pytest.raises(WorkflowSpecError, match=r"missing required \[dataset\]"): + parse_spec(raw) + + +def test_unknown_lookup_names_are_explicit(): + spec = _spec() + with pytest.raises(WorkflowSpecError, match="no \\[\\[train\\]\\] named 'zzz'"): + spec.train("zzz") + + +def test_hash_is_stable_and_order_independent(): + a = _spec() + b = parse_spec( + { + "dataset": MINIMAL["dataset"], + "condor": MINIMAL["condor"], + "workflow": MINIMAL["workflow"], + "train": MINIMAL["train"], + "rollout": MINIMAL["rollout"], + "analysis": MINIMAL["analysis"], + } + ) + assert a.train_hash("a") == b.train_hash("a") + assert a.analysis_hash("cmp") == b.analysis_hash("cmp") + assert len(a.train_hash("a")) == 8 + + +def test_hash_changes_with_own_settings(): + base = _spec() + changed = _spec(train=[{"name": "a", "epochs": 4}]) + assert base.train_hash("a") != changed.train_hash("a") + + +def test_hash_propagates_from_parents(): + """A dataset change must move every downstream task's directory.""" + base = _spec() + changed = _spec(dataset={"steps": "/data/other", "reference": "/data/holdout"}) + assert base.train_hash("a") != changed.train_hash("a") + assert base.rollout_hash("a") != changed.rollout_hash("a") + assert base.analysis_hash("cmp") != changed.analysis_hash("cmp") + + # ... and so must a change to a training the analysis transitively uses. + retrained = _spec(train=[{"name": "a", "epochs": 9}]) + assert retrained.analysis_hash("cmp") != base.analysis_hash("cmp") + # while an unrelated knob on the analysis leaves the training alone + rebinned = _spec(analysis=[{"name": "cmp", "rollouts": ["a"], "chunks": 4, "bins": 99}]) + assert rebinned.train_hash("a") == base.train_hash("a") + assert rebinned.analysis_hash("cmp") != base.analysis_hash("cmp") + + +def test_warm_cache_hash_ignores_epochs(): + """Epoch count doesn't change the setup cache, so it must not re-warm it.""" + base = _spec() + longer = _spec(train=[{"name": "a", "epochs": 50}]) + assert base.warm_cache_hash("a") == longer.warm_cache_hash("a") + other_cfg = _spec(train=[{"name": "a", "epochs": 3, "config": "configs/router.toml"}]) + assert base.warm_cache_hash("a") != other_cfg.warm_cache_hash("a") + + +def test_spec_hash_expands_dataclasses(): + spec = _spec() + assert spec_hash(spec.dataset) == spec_hash(spec.dataset) + assert spec_hash(spec.dataset) != spec_hash(spec.geometry) + + +@pytest.mark.parametrize( + "epochs, per_job, expected", + [ + (3, 1, [1, 2, 3]), + (10, 3, [3, 6, 9, 10]), + (9, 3, [3, 6, 9]), + (1, 5, [1]), + ], +) +def test_epoch_milestones(epochs, per_job, expected): + spec = _spec(train=[{"name": "a", "epochs": epochs, "epochs_per_job": per_job}]) + assert epoch_milestones(spec.train("a")) == expected diff --git a/tests/test_workflow_tasks.py b/tests/test_workflow_tasks.py new file mode 100644 index 0000000..ce09922 --- /dev/null +++ b/tests/test_workflow_tasks.py @@ -0,0 +1,172 @@ +"""Workflow task graph: dependencies, output paths, condor settings (gitea #83).""" + +import pytest + +from giant.analysis.catalog import catalog_ids, get_spec as get_plot_spec +from giant.workflow import tasks +from giant.workflow.spec import parse_spec + +CONDOR = { + "accounting_group": "cms", + "repo_dir": "/work/lbogner/giant", + "env_script": "/work/lbogner/giant/condor_env.sh", +} + +RAW = { + "workflow": {"name": "wf", "result_dir": "/results/wf"}, + "condor": CONDOR, + "dataset": {"steps": "/data/train", "reference": "/data/holdout"}, + "train": [ + {"name": "base", "epochs": 3, "gpu_memory_mb": 20000}, + {"name": "router", "epochs": 2}, + ], + "rollout": [ + {"name": "base", "train": "base"}, + {"name": "router", "train": "router"}, + ], + "analysis": [{"name": "cmp", "rollouts": ["base", "router"], "chunks": 4}], +} + + +@pytest.fixture +def spec(): + s = parse_spec(RAW) + tasks.set_spec(s) + return s + + +def _requires(task): + return list(task.requires() or []) + + +def test_epoch_chain_is_linear_and_rooted_at_warm_cache(spec): + h = spec.train_hash("base") + third = tasks.TrainEpochTask(name="base", spec_hash=h, milestone=3) + second = _requires(third) + assert [type(t) for t in second] == [tasks.TrainEpochTask] + assert second[0].milestone == 2 + first = _requires(second[0])[0] + assert first.milestone == 1 + root = _requires(first) + assert [type(t) for t in root] == [tasks.WarmCacheTask] + # the warm cache is keyed by its own hash, not the training's + assert root[0].spec_hash == spec.warm_cache_hash("base") + + +def test_epoch_task_outputs_last_pt_per_milestone(spec): + h = spec.train_hash("base") + path = tasks.TrainEpochTask(name="base", spec_hash=h, milestone=2).output().path + assert path == f"/results/wf/train_epoch/name=base/spec_hash={h}/epochs=2/last.pt" + + +def test_train_task_requires_final_epoch_and_publishes_canonical_outputs(spec): + h = spec.train_hash("base") + train = tasks.TrainTask(name="base", spec_hash=h) + (dep,) = _requires(train) + assert isinstance(dep, tasks.TrainEpochTask) and dep.milestone == 3 + out = train.output() + assert set(out) == {"best.pt", "last.pt", "metrics.csv"} + assert out["best.pt"].path == f"/results/wf/train/name=base/spec_hash={h}/best.pt" + # local: it only copies files around, no reason to queue a job for it + assert train.batch_system == "local" + + +def test_rollout_requires_training_geometry_and_reference(spec): + ro = tasks.RolloutTask(name="base", spec_hash=spec.rollout_hash("base")) + deps = _requires(ro) + assert [type(d) for d in deps] == [tasks.TrainTask, tasks.GeometryOracleTask, tasks.DatasetTask] + assert deps[0].name == "base" + assert deps[2].path == "/data/holdout" + out = ro.output() + assert out["rollout.yaml"].path.endswith("rollout.yaml") + # the sidecar sits next to the parquet — the deterministic path + # `giant rollout --out` now produces + assert out["rollout.yaml"].path[: -len(".yaml")] == out["rollout.parquet"].path[: -len(".parquet")] + + +def test_analysis_prep_requires_every_named_rollout(spec): + prep = tasks.AnalysisPrepTask(name="cmp", spec_hash=spec.analysis_hash("cmp")) + deps = _requires(prep) + assert [d.name for d in deps] == ["base", "router"] + assert all(isinstance(d, tasks.RolloutTask) for d in deps) + assert prep.batch_system == "local" + + +def test_compute_job_enumeration_collapses_non_chunkable_specs(spec): + jobs = tasks.analysis_jobs(spec, "cmp") + non_chunkable = [i for i in catalog_ids() if not get_plot_spec(i).chunkable] + expected = (len(catalog_ids()) - len(non_chunkable)) * 4 + len(non_chunkable) + assert len(jobs) == expected + assert non_chunkable, "expected some chunkable=False specs in the catalog" + for spec_id in non_chunkable: + assert [c for i, c in jobs if i == spec_id] == [0] + + +def test_compute_output_matches_the_on_disk_contract(spec): + h = spec.analysis_hash("cmp") + task = tasks.AnalysisComputeTask(name="cmp", spec_hash=h, plot_id="event_mean_length", chunk=2) + assert task.output().path == ( + f"/results/wf/analysis/name=cmp/spec_hash={h}/reduced_partial/event_mean_length__2.json" + ) + (dep,) = _requires(task) + assert isinstance(dep, tasks.AnalysisPrepTask) + + +def test_render_requires_every_compute_job_and_runs_locally(spec): + render = tasks.AnalysisRenderTask(name="cmp", spec_hash=spec.analysis_hash("cmp")) + deps = _requires(render) + assert len(deps) == len(tasks.analysis_jobs(spec, "cmp")) + assert render.batch_system == "local" # the only step importing plotstyle/LaTeX + assert render.output().path.endswith("/plots/metadata.yaml") + + +def test_workflow_task_wraps_every_analysis(spec): + deps = _requires(tasks.WorkflowTask(workflow_name="wf")) + assert [(type(d), d.name) for d in deps] == [(tasks.AnalysisRenderTask, "cmp")] + + +def test_workflow_without_analysis_falls_back_to_rollouts(): + raw = {k: v for k, v in RAW.items() if k != "analysis"} + tasks.set_spec(parse_spec(raw)) + deps = _requires(tasks.WorkflowTask(workflow_name="wf")) + assert [type(d) for d in deps] == [tasks.RolloutTask, tasks.RolloutTask] + + +def test_gpu_settings_carry_remote_ceph_and_pins(spec): + settings = tasks.TrainEpochTask(name="base", spec_hash=spec.train_hash("base"), milestone=1).htcondor_settings + assert settings["+RemoteJob"] == "True" + assert settings["RequestGPUs"] == 1 + assert "TARGET.ProvidesEtpCeph =?= True" in settings["requirements"] + assert "TARGET.GPUs_GlobalMemoryMb >= 20000" in settings["requirements"] + assert settings["accounting_group"] == "cms" + assert settings["docker_image"] == "mschnepf/slc7-condocker" + + +def test_cpu_settings_used_for_analysis_compute(spec): + task = tasks.AnalysisComputeTask( + name="cmp", spec_hash=spec.analysis_hash("cmp"), plot_id="event_mean_length", chunk=0 + ) + settings = task.htcondor_settings + assert settings["docker_image"] == "cverstege/alma9-gridjob" + assert "RequestGPUs" not in settings + # no run_meta.json yet (prep hasn't run), so no walltime is claimed + assert "+RequestWalltime" not in settings + + +def test_cpu_settings_local_files_use_provides_etp_resources(): + raw = {**RAW, "condor": {**CONDOR, "remote": False}} + spec = parse_spec(raw) + tasks.set_spec(spec) + settings = tasks.GeometryOracleTask(spec_hash=spec.geometry_hash()).htcondor_settings + assert settings["requirements"] == "TARGET.ProvidesETPResources" + assert "+RemoteJob" not in settings + + +def test_missing_dataset_fails_immediately(spec): + with pytest.raises(FileNotFoundError, match="/ceph"): + tasks.DatasetTask(path="/data/train").complete() + + +def test_dataset_that_exists_is_complete(tmp_path, spec): + (tmp_path / "steps.parquet").write_text("") + assert tasks.DatasetTask(path=str(tmp_path / "steps.parquet")).complete()