code/bundle/run_job.py

# /// script
# requires-python = ">=3.11"
# dependencies = [
#   "torch",
#   "torchvision",
#   "numpy",
#   "scipy",
#   "scikit-learn",
#   "pillow",
#   "timm",
#   "open-clip-torch",
#   "datasets",
#   "huggingface-hub",
# ]
# ///
"""HF Jobs entrypoint: expects code at /data/code (bucket mount) or alongside this file."""

from __future__ import annotations

import os
import shutil
import subprocess
import sys
from pathlib import Path


def resolve_code_root() -> Path:
    for candidate in (
        Path("/mnt/repro/code"),
        Path("/data/code"),
        Path(__file__).resolve().parent,
    ):
        if (candidate / "evaluate.py").exists():
            return candidate
    raise FileNotFoundError("evaluate.py not found under /mnt/repro/code or script directory")


def main() -> None:
    root = resolve_code_root()
    sys.path.insert(0, str(root))
    os.chdir(root)
    (root / "outputs").mkdir(exist_ok=True)

    suite = os.environ.get("SUITE", "core")
    n_binding = os.environ.get("N_BINDING", "500")
    n_jer = os.environ.get("N_JER", "40")
    n_geom = os.environ.get("N_GEOM", "500")
    out = Path(os.environ.get("OUT", str(root / "outputs" / f"results_{suite}.json")))

    skip_geom = os.environ.get("SKIP_GEOM", "0") == "1"
    skip_jer = os.environ.get("SKIP_JER", "0") == "1"
    cmd = [
        sys.executable,
        str(root / "evaluate.py"),
        "--suite",
        suite,
        "--n-binding",
        n_binding,
        "--n-samediff",
        n_binding,
        "--n-jer",
        n_jer,
        "--n-geom",
        n_geom,
        "--out",
        str(out),
    ]
    if skip_geom:
        cmd.append("--skip-geom")
    if skip_jer:
        cmd.append("--skip-jer")
    print("Running:", " ".join(cmd), flush=True)
    subprocess.check_call(cmd)
    subprocess.check_call([sys.executable, str(root / "analyze.py"), "--results", str(out)])

    # Mirror outputs into bucket mount if present
    for bucket_out in (Path("/mnt/repro/outputs"), Path("/data/outputs")):
        if bucket_out.parent.exists():
            bucket_out.mkdir(parents=True, exist_ok=True)
            for p in (root / "outputs").glob("*"):
                if p.is_file():
                    shutil.copy2(p, bucket_out / p.name)
                    print(f"[upload] copied {p.name} -> {bucket_out / p.name}", flush=True)
            break


if __name__ == "__main__":
    main()