Satyam Srivastava
e71d01648c
- Preserve existing CI and add configurable GPU backend routing - Limit execution to 2 active PRs, 4 GPUs per PR, and 8 GPUs total - Add persistent admission, cancellation recovery, and backend statuses - Document deployment and add infrastructure regression tests
99 lines
4.6 KiB
Python
99 lines
4.6 KiB
Python
"""Operator CLI. Use the isolated entrypoint from an immutable installation."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import threading
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from .dispatcher import (cancellation_signals, execute, make_store, recover,
|
|
run_directory, upload_artifacts)
|
|
from .pipeline import render_pipeline
|
|
from .policy import DEFAULT_REPOSITORY, backend_name, build_request
|
|
|
|
ROOT = Path(__file__).resolve().parents[2]
|
|
|
|
|
|
def load_config(path: Path) -> dict[str, Any]:
|
|
config = json.loads(path.read_text())
|
|
if not isinstance(config, dict):
|
|
raise ValueError("Operator config must be a JSON object")
|
|
for key in ("state_path", "artifacts_dir"):
|
|
value = config.get(key)
|
|
if not isinstance(value, str) or not Path(value).is_absolute():
|
|
raise ValueError(f"{key} must be an absolute local path")
|
|
for key, default in (("max_active_prs", 2), ("max_gpus_per_pr", 4), ("max_gpus", 8),
|
|
("queue_timeout_seconds", 21600)):
|
|
value = config.get(key, default)
|
|
if type(value) is not int or value < 1:
|
|
raise ValueError(f"{key} must be a positive integer")
|
|
config[key] = value
|
|
# The initial rollout deliberately cannot raise the user's agreed ceilings.
|
|
for key, ceiling in (("max_active_prs", 2), ("max_gpus_per_pr", 4), ("max_gpus", 8)):
|
|
if config[key] > ceiling:
|
|
raise ValueError(f"{key} exceeds the reviewed ceiling {ceiling}")
|
|
if config.get("repository", DEFAULT_REPOSITORY) != DEFAULT_REPOSITORY:
|
|
raise ValueError("This worker installation is allowlisted only for hao-ai-lab/FastVideo")
|
|
backend_name(config.get("default_backend"))
|
|
return config
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
subparsers = parser.add_subparsers(dest="command", required=True)
|
|
for name in ("run", "render", "upload", "status", "recover"):
|
|
sub = subparsers.add_parser(name)
|
|
sub.add_argument("--config", type=Path, required=True)
|
|
if name in ("render", "upload"):
|
|
sub.add_argument("--canonical", type=Path, default=ROOT / ".buildkite/pipeline.yml")
|
|
if name == "recover":
|
|
sub.add_argument("--build-id", required=True)
|
|
args = parser.parse_args(argv)
|
|
try:
|
|
config = load_config(args.config)
|
|
if args.command == "status":
|
|
print(json.dumps(make_store(config).snapshot(), indent=2))
|
|
return 0
|
|
if args.command == "recover":
|
|
recover(config, args.build_id)
|
|
return 0
|
|
backend = backend_name(os.environ.get("CI_GPU_BACKEND"), config.get("default_backend", "slurm"))
|
|
if args.command == "upload" and backend == "slurm":
|
|
if config.get("default_backend", "slurm") != "slurm":
|
|
raise ValueError("Slurm overrides are disabled after promoting another backend: legacy statuses are unqualified")
|
|
legacy = config.get("slurm_uploader")
|
|
if (not isinstance(legacy, list) or not legacy or
|
|
not all(isinstance(x, str) and x for x in legacy) or not Path(legacy[0]).is_absolute()):
|
|
raise ValueError("Configure slurm_uploader as the existing trusted uploader argv")
|
|
return subprocess.run(legacy, check=False).returncode
|
|
if args.command in ("render", "upload"):
|
|
import yaml
|
|
canonical = yaml.safe_load(args.canonical.read_text())
|
|
rendered = render_pipeline(canonical, backend, os.environ.get("TEST_SCOPE", ""),
|
|
config.get("queue", "gpu-ci-dispatch"))
|
|
payload = json.dumps(rendered, indent=2) + "\n"
|
|
if args.command == "render":
|
|
print(payload, end="")
|
|
return 0
|
|
return subprocess.run(["buildkite-agent", "pipeline", "upload", "--no-interpolation"],
|
|
input=payload, text=True, check=False).returncode
|
|
request = build_request(os.environ, config, ROOT)
|
|
cancel = threading.Event()
|
|
with cancellation_signals(cancel):
|
|
code = execute(request, config, cancel=cancel)
|
|
if config.get("upload_artifacts", True):
|
|
upload_artifacts(run_directory(config, request["build_id"]))
|
|
return code
|
|
except (OSError, ValueError, RuntimeError, KeyError, subprocess.SubprocessError) as error:
|
|
print(f"GPU CI dispatcher error: {error}", file=sys.stderr)
|
|
return 97
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|