ReNunney/scripts/run_orchestration.py

"""
run_orchestration.py

Small CLI for the minimal SQLite-backed orchestration layer.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path
import sys

REPO_ROOT = Path(__file__).resolve().parents[1]
SRC_DIR = REPO_ROOT / "src"
if str(SRC_DIR) not in sys.path:
    sys.path.insert(0, str(SRC_DIR))

from renunney.orchestration import (
    collate_track1_figure1,
    expand_track1_figure1_manifest,
    initialize_registry,
    list_jobs,
    run_one_job,
    run_worker_loop,
    submit_job_manifest,
    submit_track1_figure1_jobs,
)


def save_payload(payload: dict, path: str | Path) -> None:
    out = Path(path)
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Minimal orchestration CLI for cost-of-substitution sweeps.")
    subparsers = parser.add_subparsers(dest="command", required=True)

    init_db = subparsers.add_parser("init-db")
    init_db.add_argument("--db", required=True)

    submit = subparsers.add_parser("submit")
    submit.add_argument("--db", required=True)
    submit.add_argument("--manifest", required=True)

    submit_fig1 = subparsers.add_parser("submit-figure1")
    submit_fig1.add_argument("--db", required=True)
    submit_fig1.add_argument("--config", required=True)
    submit_fig1.add_argument("--job-prefix", required=True)
    submit_fig1.add_argument("--priority", type=int, default=0)
    submit_fig1.add_argument("--created-by", default="codex")
    submit_fig1.add_argument("--result-prefix", default="results/track1/figure1")

    list_cmd = subparsers.add_parser("list")
    list_cmd.add_argument("--db", required=True)

    run_one = subparsers.add_parser("run-one")
    run_one.add_argument("--db", required=True)
    run_one.add_argument("--result-root", required=True)
    run_one.add_argument("--worker-backend", default="python-track1")
    run_one.add_argument("--worker-host", default=None)
    run_one.add_argument("--scratch-root", default=None)

    run_loop = subparsers.add_parser("run-loop")
    run_loop.add_argument("--db", required=True)
    run_loop.add_argument("--result-root", required=True)
    run_loop.add_argument("--worker-backend", default="python-track1")
    run_loop.add_argument("--worker-host", default=None)
    run_loop.add_argument("--scratch-root", default=None)
    run_loop.add_argument("--max-jobs", type=int, default=None)

    collate = subparsers.add_parser("collate-figure1")
    collate.add_argument("--db", required=True)
    collate.add_argument("--output", default=None)

    return parser


def main() -> int:
    parser = build_parser()
    args = parser.parse_args()

    if args.command == "init-db":
        initialize_registry(args.db)
        print(json.dumps({"db": str(Path(args.db)), "status": "initialized"}, indent=2, sort_keys=True))
        return 0

    if args.command == "submit":
        manifest = json.loads(Path(args.manifest).read_text(encoding="utf-8"))
        job_id = submit_job_manifest(args.db, manifest)
        print(json.dumps({"db": str(Path(args.db)), "job_id": job_id, "status": "submitted"}, indent=2, sort_keys=True))
        return 0

    if args.command == "submit-figure1":
        config = json.loads(Path(args.config).read_text(encoding="utf-8"))
        manifests = expand_track1_figure1_manifest(
            base_job_id_prefix=args.job_prefix,
            config=config,
            priority=args.priority,
            created_by=args.created_by,
            result_prefix=args.result_prefix,
        )
        job_ids = submit_track1_figure1_jobs(
            db_path=args.db,
            base_job_id_prefix=args.job_prefix,
            config=config,
            priority=args.priority,
            created_by=args.created_by,
            result_prefix=args.result_prefix,
        )
        print(
            json.dumps(
                {
                    "db": str(Path(args.db)),
                    "job_count": len(job_ids),
                    "job_ids": job_ids,
                    "status": "submitted",
                    "preview": manifests,
                },
                indent=2,
                sort_keys=True,
            )
        )
        return 0

    if args.command == "list":
        print(json.dumps(list_jobs(args.db), indent=2, sort_keys=True))
        return 0

    if args.command == "run-one":
        result = run_one_job(
            db_path=args.db,
            result_root=args.result_root,
            worker_backend=args.worker_backend,
            worker_host=args.worker_host,
            scratch_root=args.scratch_root,
            cwd=REPO_ROOT,
        )
        print(json.dumps(result, indent=2, sort_keys=True))
        return 0 if result is None or result["status"] == "succeeded" else 1

    if args.command == "run-loop":
        payload = run_worker_loop(
            db_path=args.db,
            result_root=args.result_root,
            worker_backend=args.worker_backend,
            worker_host=args.worker_host,
            scratch_root=args.scratch_root,
            cwd=REPO_ROOT,
            max_jobs=args.max_jobs,
        )
        print(json.dumps(payload, indent=2, sort_keys=True))
        return 0 if payload["failed_jobs"] == 0 else 1

    if args.command == "collate-figure1":
        payload = collate_track1_figure1(args.db)
        if args.output:
            save_payload(payload, args.output)
        print(json.dumps(payload, indent=2, sort_keys=True))
        return 0

    raise ValueError(args.command)


if __name__ == "__main__":
    raise SystemExit(main())