| |
| """Run model configurations concurrently without fail-fast behavior.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import subprocess |
| import sys |
| from concurrent.futures import ThreadPoolExecutor, as_completed |
| from pathlib import Path |
| from typing import Any |
|
|
| from pipeline_common import REPO_ROOT, atomic_write_json, run_id, shell_join, utc_now |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("configs", nargs="+", type=Path) |
| parser.add_argument("--workers", type=int, default=2) |
| parser.add_argument("--stage", action="append", dest="stages") |
| parser.add_argument("--dry-run", action="store_true") |
| parser.add_argument("--force", action="store_true") |
| parser.add_argument("--resume-failed", action="store_true") |
| return parser.parse_args() |
|
|
|
|
| def run_one(config: Path, batch_dir: Path, args: argparse.Namespace) -> dict[str, Any]: |
| name = config.stem |
| stdout_path = batch_dir / f"{name}.stdout.log" |
| stderr_path = batch_dir / f"{name}.stderr.log" |
| argv = [sys.executable, str(REPO_ROOT / "scripts" / "run_model.py"), "--config", str(config.resolve())] |
| for stage in args.stages or []: |
| argv.extend(["--stage", stage]) |
| for enabled, flag in ( |
| (args.dry_run, "--dry-run"), (args.force, "--force"), (args.resume_failed, "--resume-failed") |
| ): |
| if enabled: |
| argv.append(flag) |
| started_at = utc_now() |
| with stdout_path.open("wb") as stdout_handle, stderr_path.open("wb") as stderr_handle: |
| completed = subprocess.run(argv, cwd=REPO_ROOT, stdout=stdout_handle, stderr=stderr_handle, check=False) |
| return { |
| "config": str(config.resolve()), |
| "command": shell_join(argv), |
| "started_at": started_at, |
| "ended_at": utc_now(), |
| "exit_code": completed.returncode, |
| "status": "PASS" if completed.returncode == 0 else "FAIL", |
| "stdout_log": str(stdout_path), |
| "stderr_log": str(stderr_path), |
| } |
|
|
|
|
| def main() -> int: |
| args = parse_args() |
| if args.workers < 1: |
| raise SystemExit("--workers must be at least 1") |
| current_run_id = run_id() |
| batch_dir = REPO_ROOT / "logs" / "batch" / current_run_id |
| batch_dir.mkdir(parents=True, exist_ok=False) |
| configs = [path.resolve() for path in args.configs] |
| results: list[dict[str, Any]] = [] |
| with ThreadPoolExecutor(max_workers=args.workers) as executor: |
| futures = {executor.submit(run_one, config, batch_dir, args): config for config in configs} |
| for future in as_completed(futures): |
| config = futures[future] |
| try: |
| results.append(future.result()) |
| except Exception as error: |
| results.append({ |
| "config": str(config), "command": "", "started_at": utc_now(), |
| "ended_at": utc_now(), "exit_code": 127, "status": "FAIL", |
| "stdout_log": "", "stderr_log": "", "error": str(error), |
| }) |
| results.sort(key=lambda item: item["config"]) |
| summary = { |
| "schema_version": "1.0", |
| "run_id": current_run_id, |
| "worker_count": args.workers, |
| "total": len(results), |
| "passed": sum(item["status"] == "PASS" for item in results), |
| "failed": sum(item["status"] == "FAIL" for item in results), |
| "results": results, |
| } |
| atomic_write_json(batch_dir / "batch_summary.json", summary) |
| print(json.dumps(summary, sort_keys=True)) |
| return 0 if summary["failed"] == 0 else 1 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
|
|
