"""Generic job wrapper: pull code from the Hub repo, run one run_all.py stage.
The stage is selected with the JOB_STAGE env variable."""

import os
import subprocess
import sys

REPO = "eac123/ptr-recurrent-depth"
STAGE = os.environ["JOB_STAGE"]

from huggingface_hub import snapshot_download

snapshot_download(repo_id=REPO, repo_type="model",
                  allow_patterns=["*.py", "ptr/*.py"], local_dir=".")
sys.path.insert(0, os.getcwd())
subprocess.run([sys.executable, "run_all.py", STAGE, "--hub-repo", REPO],
               check=True)
print(f"STAGE {STAGE} COMPLETE")