Merge pull request #309 from bokiko/fix/bench-fetch-robust

bench: survive transient hub errors — retry with backoff, per-task isolation, atomic writes (#304)

c/coli

@@ -661,6 +661,16 @@ def cmd_bench(a):
         print(f"  {C.dim}downloading missing datasets: {', '.join(missing)}{C.r}")
         subprocess.call([py, os.path.join(TOOLS,"fetch_benchmarks.py"),
                          "--out", a.data, "--tasks", ",".join(missing), "--limit", str(max(a.limit,200))])
+        # il fetch riprova da solo (#304), ma se l'hub resta giu' il bench gira sui task
+        # disponibili invece di passare a eval file inesistenti.
+        # EN: the fetch retries on its own (#304), but if the hub stays down the bench
+        # runs on the available tasks instead of handing eval nonexistent files.
+        still=[t for t in tasks.split(",") if not os.path.exists(os.path.join(a.data,f"{t}.jsonl"))]
+        if still:
+            tasks=",".join(t for t in tasks.split(",") if t not in still)
+            print(f"  {C.yel}skipping (download failed, rerun later): {', '.join(still)}{C.r}")
+            if not tasks:
+                print(f"  {C.yel}no datasets available — nothing to bench{C.r}"); sys.exit(1)
     cmd=[py, os.path.join(TOOLS,"eval_glm.py"), "--glm", GLM, "--snap",a.model,
          "--tasks", tasks, "--limit", str(a.limit), "--data", a.data]
     if a.ram: cmd+=["--ram",str(a.ram)]

c/tools/fetch_benchmarks.py

@@ -8,7 +8,7 @@ USO:
 Poi:
   python3 tools/eval_glm.py --snap /home/vincenzo/glm52_i4 --data ./bench --tasks mmlu --limit 40 --ram 15
 """
-import os, json, argparse, random
+import os, sys, json, time, argparse, random
 
 def f_hellaswag(d):
     ctx = (d["activity_label"] + ": " + d["ctx_a"] + " " + d["ctx_b"].capitalize()).strip()
@@ -38,19 +38,44 @@ TASKS = {  # task: (path, config, split, formatter)
     "openbookqa":    ("allenai/openbookqa", "main", "validation", f_openbookqa),
 }
 
+def load_retry(path, cfg, split, tries=5):
+    """L'hub restituisce 5xx/timeout transitori (#304): riprova con backoff invece di
+    morire al primo HEAD fallito. hf_hub riprende i download parziali dalla cache,
+    quindi il retry riparte da dove si era fermato, non da zero.
+    EN: the hub throws transient 5xx/timeouts (#304): retry with backoff instead of
+    dying on the first failed HEAD. hf_hub resumes partial downloads from its cache,
+    so a retry continues where it stopped rather than starting over."""
+    from datasets import load_dataset
+    for k in range(tries):
+        try:
+            return load_dataset(path, cfg, split=split)
+        except (KeyboardInterrupt, SystemExit):
+            raise
+        except Exception as e:
+            if k == tries - 1: raise
+            wait = 2 ** (k + 1)
+            print(f"  {path}: {type(e).__name__}: {e} — retry {k+1}/{tries-1} in {wait}s")
+            time.sleep(wait)
+
 def main():
     ap = argparse.ArgumentParser()
     ap.add_argument("--out", default="./bench")
     ap.add_argument("--tasks", default="hellaswag,arc_challenge,mmlu")
     ap.add_argument("--limit", type=int, default=300)
     ap.add_argument("--seed", type=int, default=1234)
+    ap.add_argument("--tries", type=int, default=5)
     a = ap.parse_args()
-    from datasets import load_dataset
     os.makedirs(a.out, exist_ok=True)
+    failed = []
     for t in [x.strip() for x in a.tasks.split(",") if x.strip()]:
-        if t not in TASKS: print("unknown task:", t); continue
+        if t not in TASKS: print("unknown task:", t); failed.append(t); continue
         path, cfg, split, fn = TASKS[t]
-        ds = load_dataset(path, cfg, split=split)
+        try:
+            ds = load_retry(path, cfg, split, tries=max(a.tries, 1))
+        except Exception as e:
+            # un task fallito non deve uccidere gli altri / one failed task must not kill the rest
+            print(f"{t}: FAILED after {a.tries} tries ({type(e).__name__}: {e}) — skipping")
+            failed.append(t); continue
         idx = list(range(len(ds))); random.Random(a.seed).shuffle(idx)
         rows, n = [], 0
         for i in idx:
@@ -61,9 +86,18 @@ def main():
             except Exception: continue
             if n >= a.limit: break
         outp = os.path.join(a.out, t + ".jsonl")
-        with open(outp, "w") as f:
+        # scrittura atomica: coli controlla solo l'ESISTENZA del file, quindi un jsonl
+        # troncato da un run interrotto bloccherebbe il re-download per sempre.
+        # EN: atomic write: coli only checks the file EXISTS, so a truncated jsonl from
+        # an interrupted run would block re-download forever.
+        tmp = outp + ".part"
+        with open(tmp, "w") as f:
             for r in rows: f.write(json.dumps(r) + "\n")
+        os.replace(tmp, outp)
         print(f"{t}: {len(rows)} -> {outp}")
+    if failed:
+        print(f"incomplete: {', '.join(failed)} — rerun when the hub recovers (cached progress is kept)")
+        sys.exit(1)
 
 if __name__ == "__main__":
     main()