Spaces:
Running
Running
| """Run one of our scripts on an HF Job GPU. | |
| Ships the code folder and mounts the private dataset repo at /data. The HF token is read | |
| from the HF_TOKEN env var and passed as an encrypted job secret (never written to disk). | |
| Usage: | |
| python launch_job.py --name smoke --flavor a10g-small --timeout 1h \ | |
| --pip "transformers accelerate peft" -- student.py --split pool --limit 20 ... | |
| Everything after "--" is the script and its arguments; --data/--out are added for you. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import os | |
| import shlex | |
| import sys | |
| import time | |
| from huggingface_hub import HfApi, Volume, fetch_job_logs, inspect_job, run_job, sync_job_volume | |
| HERE = os.path.dirname(os.path.abspath(__file__)) | |
| DATA_REPO = "HideIron/selfheal-sql-spider" | |
| RUNS_REPO = "HideIron/selfheal-sql-runs" | |
| IMAGE = "pytorch/pytorch:2.6.0-cuda12.4-cudnn9-devel" | |
| def main() -> None: | |
| sys.stdout.reconfigure(encoding="utf-8", errors="replace") | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--name", required=True) | |
| ap.add_argument("--flavor", default="a10g-small") | |
| ap.add_argument("--timeout", default="1h") | |
| ap.add_argument("--image", default=IMAGE) | |
| ap.add_argument("--pip", default="transformers accelerate peft huggingface_hub jinja2>=3.1.6") | |
| ap.add_argument("--detach", action="store_true") | |
| ap.add_argument("rest", nargs=argparse.REMAINDER) | |
| a = ap.parse_args() | |
| rest = a.rest[1:] if a.rest and a.rest[0] == "--" else a.rest | |
| script, script_args = rest[0], rest[1:] | |
| token = os.environ["HF_TOKEN"] | |
| api = HfApi(token=token) | |
| api.create_repo(RUNS_REPO, repo_type="dataset", private=True, exist_ok=True) | |
| code = sync_job_volume(HERE, "/code", token=token) | |
| data = Volume(type="dataset", source=DATA_REPO, mount_path="/data") | |
| inner = ( | |
| "export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True && " | |
| "pip install -q -U " + " ".join(shlex.quote(x) for x in a.pip.split()) + " && " | |
| f"python /code/{script} --data /data --out /tmp/out " | |
| f"--push-repo {RUNS_REPO} --push-path runs/{a.name} " + " ".join(shlex.quote(x) for x in script_args) | |
| ) | |
| job = run_job( | |
| image=a.image, command=["bash", "-lc", inner], flavor=a.flavor, timeout=a.timeout, | |
| name=a.name, volumes=[code, data], secrets={"HF_TOKEN": token}, | |
| ) | |
| print("job:", job.url, flush=True) | |
| if a.detach: | |
| print("job_id:", job.id) | |
| return | |
| while True: | |
| st = inspect_job(job_id=job.id).status | |
| if st.stage in ("COMPLETED", "ERROR", "CANCELED", "DELETED"): | |
| break | |
| time.sleep(10) | |
| print("stage:", st.stage, "|", st.message) | |
| lines = [ln for ln in fetch_job_logs(job_id=job.id) if "[A" not in ln and "it/s]" not in ln and "B/s]" not in ln] | |
| for line in lines[-80:]: | |
| print(line) | |
| sys.exit(0 if st.stage == "COMPLETED" else 1) | |
| if __name__ == "__main__": | |
| main() | |