File size: 2,348 Bytes
5de8603 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 | #!/usr/bin/env python3
"""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()
|