code/job_entry.py

"""Self-contained job entry: downloads repro code from the HF bucket, installs
deps, runs the full PhotoAgent toy-scale reproduction (Claims 1,2,3,4,5), and
writes all outputs back to the bucket under /outputs.

Run inside an hf Job (python:3.12 image) WITHOUT -v mount:
  hf jobs run python:3.12 python job_entry.py --flavor a10g-small --secrets HF_TOKEN
"""
import os, sys, subprocess, json, time

BUCKET = "hf://buckets/Yashp2003/photoagent-repro-artifacts"
CODE_DIR = "/tmp/repro/code"
OUT_DIR = "/tmp/repro/outputs"


def run(cmd):
    print(">>>", cmd)
    r = subprocess.run(cmd, shell=True)
    return r.returncode


def main():
    t0 = time.time()
    os.makedirs(CODE_DIR, exist_ok=True)
    os.makedirs(OUT_DIR, exist_ok=True)

    # 1) pull code + inputs from bucket
    run(f"hf buckets sync {BUCKET}/code {CODE_DIR}")
    # 2) install deps
    run("pip install -q diffusers transformers accelerate qwen-vl-utils "
        "tensorboard safetensors Pillow numpy huggingface_hub torchvision")
    # 3) run full reproduction
    os.chdir(CODE_DIR)
    rc = run(f"python3 run_gpu_job.py")
    # 4) push outputs back to bucket
    run(f"hf buckets sync {OUT_DIR} {BUCKET}/outputs")
    # 5) also print summary to stdout (immutable in job logs)
    try:
        with open(os.path.join(OUT_DIR, "results.json")) as f:
            print("=== RESULTS ===")
            print(f.read())
    except Exception as e:
        print("no results.json:", e)
    print(f"JOB DONE in {time.time()-t0:.0f}s, rc={rc}")


if __name__ == "__main__":
    main()