lars
a482b04761
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>
104 lines
4.1 KiB
Python
104 lines
4.1 KiB
Python
#!/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()
|