File size: 5,951 Bytes
0298809 371ff91 0298809 371ff91 0298809 371ff91 0298809 75a9ca0 0298809 371ff91 0298809 371ff91 0298809 | 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 | """Per-checkpoint eval worker: watch the adapter repo, eval each new checkpoint.
Loop: poll ADAPTER_REPO for checkpoint-N dirs without eval results, hot-load each
as a LoRA into the already-running vLLM (runtime LoRA updates — no server restart),
run the held-out eval on a fixed task subset, and upload per-checkpoint results
plus a rolling eval/summary.json (the live learning curve).
Base-model cells are computed once and reused for every checkpoint.
Env: ADAPTER_REPO (required), BASE_MODEL, N_TASKS (default 100), POLL_S (600).
Stops when eval/final_results.json appears (the training job's own final eval).
"""
import json
import os
import subprocess
import time
import urllib.request
from pathlib import Path
from huggingface_hub import HfApi, snapshot_download
REPO = os.environ["ADAPTER_REPO"]
BASE_MODEL = os.environ.get("BASE_MODEL", "LiquidAI/LFM2.5-2.6B")
POLL_S = int(os.environ.get("POLL_S", "600"))
VLLM = "http://127.0.0.1:8001"
api = HfApi()
def repo_files():
return set(api.list_repo_files(REPO))
def pending_checkpoints(fs):
ckpts = sorted(
{f.split("/")[0] for f in fs if f.startswith("checkpoint-") and f.endswith("adapter_config.json")},
key=lambda s: int(s.rsplit("-", 1)[1]),
)
done = {f[len("eval/"):-len("_results.json")] for f in fs if f.startswith("eval/") and f.endswith("_results.json")}
return [c for c in ckpts if c not in done]
def run_eval(model_name, out_path):
env = dict(os.environ, ONLY_MODEL=model_name, OUT=out_path)
subprocess.run(["python3", "/app/eval_adapter.py"], env=env, check=True)
def upload(local, remote):
api.upload_file(path_or_fileobj=str(local), path_in_repo=remote, repo_id=REPO)
def lora_rpc(route, name, path=None):
payload = {"lora_name": name}
if path is not None:
payload["lora_path"] = str(path)
req = urllib.request.Request(
f"{VLLM}/v1/{route}",
data=json.dumps(payload).encode(),
headers={"Content-Type": "application/json"},
)
urllib.request.urlopen(req, timeout=900)
def agg(cells_path):
"""pass@1 plus efficiency means over solved tasks (the article's second axis)."""
cells = json.loads(Path(cells_path).read_text())["cells"]
scored = [c["correct"] for c in cells.values() if not c["failed"]]
solved = [c["metrics"] for c in cells.values() if c["correct"] and c.get("metrics")]
def mean(key):
vals = [m[key] for m in solved if key in m]
return round(sum(vals) / len(vals), 1) if vals else None
return {
"pass_at_1": (sum(scored) / len(scored)) if scored else None,
"scorable": len(scored),
"solved": len(solved),
"tool_calls_solved_mean": mean("tool_calls"),
"turns_solved_mean": mean("turns"),
"completion_tokens_solved_mean": mean("completion_tokens"),
"prompt_tokens_solved_mean": mean("prompt_tokens_total"),
}
def main():
fs = repo_files()
# 1. Base-model cells, once. Every checkpoint compares against these.
# Recomputed if the cached file predates per-cell metrics.
if "eval/base_cells.json" in fs:
local = api.hf_hub_download(REPO, "eval/base_cells.json")
Path("/tmp/base_cells.json").write_text(Path(local).read_text())
if "eval/base_cells.json" not in fs or any(
c.get("metrics") is None
for c in json.loads(Path("/tmp/base_cells.json").read_text())["cells"].values()
if c["correct"]
):
print("[watch] evaluating base model first", flush=True)
run_eval(BASE_MODEL, "/tmp/base_cells.json")
upload("/tmp/base_cells.json", "eval/base_cells.json")
summary = {}
if "eval/summary.json" in fs:
summary = json.loads(Path(api.hf_hub_download(REPO, "eval/summary.json")).read_text())
loaded = []
while True:
fs = repo_files()
if "eval/final_results.json" in fs:
print("[watch] training job's final eval landed; stopping", flush=True)
return
for ckpt in pending_checkpoints(fs):
print(f"[watch] new checkpoint: {ckpt}", flush=True)
local = Path(snapshot_download(REPO, allow_patterns=[f"{ckpt}/*"])) / ckpt
lora_rpc("load_lora_adapter", ckpt, local)
loaded.append(ckpt)
if len(loaded) > 4: # stay under --max-loras
lora_rpc("unload_lora_adapter", loaded.pop(0))
try:
run_eval(ckpt, f"/tmp/{ckpt}_cells.json")
except subprocess.CalledProcessError as exc:
# Transient eval failure (client reconnect flake, sandbox hiccup):
# skip it — the next poll recomputes pending and retries it.
print(f"[watch] {ckpt} eval failed ({exc}); will retry next poll", flush=True)
continue
upload(f"/tmp/{ckpt}_cells.json", f"eval/{ckpt}_results.json")
a, b = agg(f"/tmp/{ckpt}_cells.json"), agg("/tmp/base_cells.json")
entry = dict(a)
entry["base_pass_at_1"] = b["pass_at_1"]
if a["pass_at_1"] is not None and b["pass_at_1"] is not None:
entry["delta"] = a["pass_at_1"] - b["pass_at_1"]
for key in ("tool_calls_solved_mean", "turns_solved_mean",
"completion_tokens_solved_mean", "prompt_tokens_solved_mean"):
if a.get(key) is not None and b.get(key):
entry[key.replace("_mean", "_vs_base")] = round(a[key] / b[key], 3)
summary[ckpt] = entry
Path("/tmp/summary.json").write_text(json.dumps(summary, indent=2))
upload("/tmp/summary.json", "eval/summary.json")
print(f"[watch] {ckpt}: pass@1={a['pass_at_1']} (base {b['pass_at_1']}), "
f"tool_calls={a['tool_calls_solved_mean']} (base {b['tool_calls_solved_mean']})",
flush=True)
time.sleep(POLL_S)
if __name__ == "__main__":
main()
|