Add the giant/workflow b2luigi task graph (gitea #83)
One workflow TOML now parameterises a whole experiment and `giant workflow run <spec.toml>` 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 <noreply@anthropic.com>
This commit is contained in:
@@ -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
|
||||
@@ -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.",
|
||||
|
||||
@@ -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 <spec.toml>`` 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",
|
||||
]
|
||||
@@ -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)})",
|
||||
}
|
||||
@@ -0,0 +1,103 @@
|
||||
#!/usr/bin/env python
|
||||
"""Entry point b2luigi re-executes on every worker.
|
||||
|
||||
Locally this is what ``giant workflow run <spec.toml>`` 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 `<executable> [<basename of this file>] --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()
|
||||
@@ -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 = <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
|
||||
@@ -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 ``<result_dir>/<kind>/name=<name>/spec_hash=
|
||||
<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 = <job k-1>/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 <same file>` (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:
|
||||
"""``<result_dir>/<kind>/name=<name>/spec_hash=<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 (``<data>.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/<id>__<chunk>.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))
|
||||
@@ -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"
|
||||
|
||||
@@ -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
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user