File size: 7,787 Bytes
84d148b 688fc29 84d148b 688fc29 84d148b 688fc29 84d148b 688fc29 84d148b | 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 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 | """Self-contained remote run: pulls code and data from the Hub, scores, evaluates.
Nothing is downloaded to a workstation. The Guardian tarballs total ~68 GB, and rlbench's test
tarball alone is 45.8 GB because the archives carry `global.avi` videos alongside the frames, so
only the tarballs a given stage actually needs are fetched, on the machine that will read them.
Stages, smallest first, so a failure costs minutes rather than an hour of GPU:
--stage smoke bridge + ur5, --limit episodes per split, zero-shot only (~1.1 GB)
--stage full adds rlbench train, all splits, zero-shot + LoRA (~9.9 GB)
Usage on Kaggle (T4, 4-bit):
python kaggle_run.py --stage smoke --limit 24
python kaggle_run.py --stage full
"""
from __future__ import annotations
import argparse
import json
import os
import subprocess
import sys
import tarfile
import time
from pathlib import Path
CODE_REPO = "Chucks90/judgetron"
# domain -> list of HF dataset repos to pool. prepare_splits does its own task-disjoint split, so
# a source's upstream train/val/test division is not reused; pooling only widens the task pool.
#
# ur5 pools all three of its splits deliberately. It is the OOD domain, so it carries the headline
# numbers, and the bootstrap resamples whole TASKS: ur5 train alone has just 7 taskvars, which after
# the cal_ood/test_ood split leaves 5 task clusters in test_ood and CIs too wide to read. Pooling
# gives 34 taskvars over 570 episodes (train/test taskvars are disjoint upstream) for 0.3 GB total.
UR5_ALL = [("ur5", "paulpacaud/ur5fail_train_dataset"),
("ur5", "paulpacaud/ur5fail_val_dataset"),
("ur5", "paulpacaud/ur5fail_test_dataset")]
SOURCES = {
"smoke": {
"in_domain": [("bridge", "paulpacaud/bdv2fail_train_dataset")],
"ood": UR5_ALL,
},
"full": {
# rlbench train is 8.78 GB; its val/test tarballs (11.9 / 45.8 GB) are deliberately not used
# -- they are mostly .avi video, and one source split already gives 70 taskvars.
"in_domain": [("rlbench", "paulpacaud/rlbenchfail_train_dataset"),
("bridge", "paulpacaud/bdv2fail_train_dataset")],
"ood": UR5_ALL,
},
}
def sh(cmd: list[str], **kw) -> None:
print(f"$ {' '.join(cmd)}", flush=True)
subprocess.run(cmd, check=True, **kw)
def fetch_source(domain: str, repo: str, root: Path) -> Path:
"""Download one source's metadata + image tarball and extract it. Returns the metadata path."""
from huggingface_hub import hf_hub_download
out = root / repo.split("/")[-1] # one dir per repo: pooled repos must not overwrite
out.mkdir(parents=True, exist_ok=True)
meta = hf_hub_download(repo, "metadata_execution.jsonl", repo_type="dataset",
local_dir=out)
marker = out / ".extracted"
if not marker.exists():
t0 = time.time()
tar = hf_hub_download(repo, "records.tar.gz", repo_type="dataset", local_dir=out)
print(f"[{domain}] downloaded records.tar.gz in {time.time()-t0:.0f}s", flush=True)
t0 = time.time()
with tarfile.open(tar) as tf:
# Only the frames are needed; the archives also carry .avi videos, which is most of
# their size. Filtering here keeps peak disk well under Kaggle's working-dir limit.
members = [m for m in tf.getmembers()
if m.isfile() and m.name.lower().endswith((".png", ".jpg", ".jpeg"))]
tf.extractall(out, members=members)
print(f"[{domain}] extracted {len(members)} frames in {time.time()-t0:.0f}s", flush=True)
os.remove(tar) # reclaim the archive immediately
marker.write_text("ok")
return Path(meta)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--stage", choices=["smoke", "full"], default="smoke")
ap.add_argument("--limit", type=int, default=None, help="episodes per split (smoke)")
ap.add_argument("--data", default="/kaggle/tmp/data")
ap.add_argument("--out", default="/kaggle/working")
ap.add_argument("--code", default=None, help="local judgecal checkout; otherwise pulled from HF")
ap.add_argument("--model-id", default="Qwen/Qwen2.5-VL-3B-Instruct")
ap.add_argument("--batch-size", type=int, default=2)
ap.add_argument("--train-steps", type=int, default=1500)
# Effective batch is batch-size x grad-accum, so 2x4 matches the documented 1x8 while halving
# the number of forward/backward passes -- the thing that actually costs wall-clock on a T4.
ap.add_argument("--train-batch-size", type=int, default=2)
ap.add_argument("--grad-accum", type=int, default=4)
ap.add_argument("--zeroshot", action="store_true", help="score zero-shot too (skip if already done)")
ap.add_argument("--skip-lora", action="store_true")
a = ap.parse_args()
out = Path(a.out); out.mkdir(parents=True, exist_ok=True)
data = Path(a.data); data.mkdir(parents=True, exist_ok=True)
# ---- code
if a.code:
code = Path(a.code)
else:
from huggingface_hub import snapshot_download
code = Path(snapshot_download(CODE_REPO, repo_type="model",
local_dir=str(out / "judgetron_code")))
sys.path.insert(0, str(code))
print(f"[code] {code}", flush=True)
# ---- data
spec = SOURCES[a.stage]
src_args = []
for domain, repo in spec["in_domain"]:
m = fetch_source(domain, repo, data)
src_args += ["--source", f"{domain}={m}"]
ood_args = []
for domain, repo in spec["ood"]:
m = fetch_source(domain, repo, data)
ood_args += ["--ood", f"{domain}={m}"]
# ---- splits
manifests = out / "manifests"
sh([sys.executable, str(code / "scripts" / "prepare_splits.py"), *src_args,
*ood_args, "--out", str(manifests)])
# ---- zero-shot
scores = out / "results" / "scores.csv"
scores.parent.mkdir(parents=True, exist_ok=True)
cmd = [sys.executable, str(code / "scripts" / "score.py"),
"--manifests", str(manifests), "--scores", str(scores),
"--model-id", a.model_id, "--load-in-4bit",
"--batch-size", str(a.batch_size), "--name", "qwen3b-zeroshot"]
if a.limit:
cmd += ["--limit", str(a.limit)]
if a.zeroshot or a.skip_lora:
sh(cmd)
else:
print("[skip] zero-shot scoring (--zeroshot not set); scored in its own kernel", flush=True)
# ---- LoRA
if not a.skip_lora:
adapter = out / "adapters" / "qwen3b-lora"
sh([sys.executable, str(code / "scripts" / "train.py"),
"--manifests", str(manifests), "--out", str(adapter),
"--model-id", a.model_id, "--load-in-4bit", "--steps", str(a.train_steps),
"--batch-size", str(a.train_batch_size), "--grad-accum", str(a.grad_accum)])
cmd = [sys.executable, str(code / "scripts" / "score.py"),
"--manifests", str(manifests), "--scores", str(scores),
"--model-id", a.model_id, "--adapter", str(adapter), "--load-in-4bit",
"--batch-size", str(a.batch_size), "--name", "qwen3b-lora"]
if a.limit:
cmd += ["--limit", str(a.limit)]
sh(cmd)
# ---- report. With the grid split across kernels each writes only its own rows, so the report
# is regenerated from the merged scores.csv afterwards (evaluate.py is CPU-only).
if scores.exists():
sh([sys.executable, str(code / "scripts" / "evaluate.py"),
"--scores", str(scores), "--out", str(out / "results")])
print(json.dumps({"stage": a.stage, "scores": str(scores),
"report": str(out / "results" / "report.md")}, indent=2), flush=True)
if __name__ == "__main__":
main()
|