Spaces:
Running
Running
File size: 2,987 Bytes
f636e22 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | """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()
|