| |
| """Submit the fail-closed build, smoke, and resumable production chain.""" |
| from __future__ import annotations |
|
|
| from datetime import datetime |
| import json |
| from pathlib import Path |
| import subprocess |
|
|
|
|
| HERE = Path(__file__).resolve().parent |
| STATUS = HERE / "caption_pipeline_jobs.json" |
|
|
|
|
| def submit(args, dependency=None): |
| command = ["sbatch", "--parsable"] |
| if dependency: |
| command += [f"--dependency=afterok:{dependency}"] |
| command += [str(x) for x in args] |
| result = subprocess.run(command, cwd=HERE, check=True, text=True, capture_output=True) |
| return result.stdout.strip().split(";")[0] |
|
|
|
|
| def main(): |
| if STATUS.exists(): |
| previous = json.loads(STATUS.read_text()) |
| previous_asset = previous.get("jobs", {}).get("assets", "unknown") |
| archive = HERE / f"caption_pipeline_jobs.failed-{previous_asset}.json" |
| assert not archive.exists(), f"Refusing to overwrite pipeline archive: {archive}" |
| STATUS.replace(archive) |
| jobs = {} |
| jobs["assets"] = submit(["build_caption_assets.sbatch"]) |
| jobs["stages"] = submit(["build_caption_stages.sbatch"], jobs["assets"]) |
| jobs["freeze"] = submit(["freeze_caption_plans.sbatch"], jobs["stages"]) |
| jobs["smoke_4gpu"] = submit([ |
| "--nodes=1", "--time=01:00:00", "--job-name=m2-cap-smoke4", |
| "train.sbatch", "plan_caption_smoke_4gpu.json"], jobs["freeze"]) |
| jobs["smoke_32nodes"] = submit([ |
| "--nodes=32", "--time=01:00:00", "--job-name=m2-cap-smoke32", |
| "train.sbatch", "plan_caption_smoke_32nodes.json"], jobs["smoke_4gpu"]) |
| dependency = jobs["smoke_32nodes"] |
| production = [] |
| for index in range(1, 6): |
| job = submit(["--nodes=32", "--time=12:00:00", f"--job-name=m2-cap-prod{index}", |
| "train.sbatch", "plan_caption_production_32nodes.json"], dependency) |
| production.append(job); dependency = job |
| jobs["production_chain"] = production |
| payload = {"status": "SUBMITTED", "created_at": datetime.now().astimezone().isoformat(), |
| "dependency_policy": "afterok", "jobs": jobs, |
| "production_authorized": True, |
| "throughput_gate_seconds_per_update": 5.0} |
| STATUS.write_text(json.dumps(payload, indent=2) + "\n") |
| print(json.dumps(payload, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
