import argparse import asyncio import copy import hashlib import json import math import shutil import time from pathlib import Path import numpy as np import tinker from research_exchange.research_decisions import assess, paired_interval from research_exchange.restoration import output_fingerprint from tinker_cookbook import renderers from tinker_cookbook.completers import TinkerTokenCompleter from tinker_cookbook.rl.data_processing import assemble_training_data from tinker_cookbook.rl.rollouts import do_group_rollout from eval import canonical, digest, freeze, summarize from pilot import pilot_advantages from remote_stages import restore_datum from rx_gs1 import MODEL, ROOT, advantages, dataset, datums, guard_grader, service, write_json ARMS = ("grpo_gs8", "gs1_kernel") INIT = "tinker://cf264e49a4dd83e5b827d28f49b9c43a/weights/rx-gs1-init-run_5207121d26" EVALUATOR = "5ce61cac464d6d7739964d336937b9a4de9d9b6a48b9905f59a0d978824a1d70" OUT = ROOT / "full" def atomic(path, value): staging = path.with_suffix(path.suffix + ".new") write_json(staging, value) staging.replace(path) def sha(path): return hashlib.sha256(path.read_bytes()).hexdigest() def lost(error): return any(s in str(error).lower() for s in ("model_lost", "model lost", "checkpoint recovery required")) def capacity(error): return any(s in str(error).lower() for s in ("capacity exhausted", "429", "503")) async def rpc(method, *args, **kwargs): async with asyncio.timeout(300): result = await method(*args, **kwargs) if hasattr(result, "result_async"): result = await result.result_async() return result async def retry(method, *args, **kwargs): start = time.monotonic() for delay in (60, 120, 240, None): try: return await rpc(method, *args, **kwargs) except Exception as error: if not capacity(error) or delay is None or time.monotonic() - start + delay > 600: raise with (OUT / "capacity-retries.jsonl").open("a") as f: f.write(canonical(dict(operation=method.__name__, error=str(error), backoff=delay)) + "\n") print("capacity backoff", method.__name__, delay, flush=True) await asyncio.sleep(delay) def observation(event): return dict(after_restore=event, source="declared", actual=None, actual_config=None, requested="adamw", requested_config=dict(lr=4e-5, beta1=0.9, beta2=0.95, eps=1e-8), before_next_update=True, reason="Shared Spindle accepts adam_params; no concrete active optimizer/config response.") async def timed_rollouts(sampler, builders): start = time.monotonic() async def feature(builder): prompt, _ = await builder.env_thunk().initial_observation() logprobs = await sampler.compute_logprobs_async(prompt) values = [p for p in logprobs if p is not None] if not values or not np.isfinite(values).all(): raise ValueError("missing/nonfinite prompt logprobs") return [math.log(prompt.length), float(np.mean(values))] async with asyncio.timeout(300): features = await asyncio.gather(*(feature(b) for b in builders)) feature_seconds = time.monotonic() - start start = time.monotonic() policy = TinkerTokenCompleter(sampler, max_tokens=512, temperature=1.0) async with asyncio.timeout(300): groups = await asyncio.gather(*(do_group_rollout(b, policy) for b in builders)) return groups, [f for f, g in zip(features, groups) for _ in g.trajectories_G], dict( logprob_feature_seconds=feature_seconds, rollout_seconds=time.monotonic() - start) class Arm: def __init__(self, name, resume): self.name = name self.directory = OUT / name self.directory.mkdir(parents=True, exist_ok=True) self.transient = Path("/tmp/rx-gs1") / name self.transient.mkdir(parents=True, exist_ok=True) self.state_path = self.directory / "state.json" if self.state_path.exists() and not resume: raise ValueError("existing run state; explicitly pass --resume") self.state = json.loads(self.state_path.read_text()) if self.state_path.exists() else dict( arm=name, completed_steps=0, checkpoint=INIT, checkpoint_step=0, checkpoints=[], recoveries=[], extra_replay_updates=0) self.sc = service() self.tc = None self.history = [] self.provenance = dict(name=name, restore_count=0, optimizer_checks=[observation(0)], restores=[]) if resume and (self.directory / "initialization.json").exists(): self.provenance = json.loads((self.directory / "initialization.json").read_text()) def save(self): atomic(self.state_path, self.state) atomic(self.directory / "initialization.json", self.provenance) def step_path(self, n): return self.directory / f"step-{n:03d}.json" def input_path(self, n): return self.directory / f"step-{n:03d}-inputs.json" async def load(self, optimizer=False): self.tc = await retry(self.sc.create_lora_training_client_async, base_model=MODEL, rank=32, optimizer=tinker.AdamOptimizerConfig()) path = self.state["checkpoint"] if optimizer else INIT method = self.tc.load_state_with_optimizer_async if optimizer else self.tc.load_state_async response = await retry(method, path) if response.path != path: raise ValueError("server loaded unexpected checkpoint") self.provenance["restore_count"] += 1 self.provenance["optimizer_checks"].append(observation(self.provenance["restore_count"])) self.provenance["restores"].append(dict(requested=path, with_optimizer=optimizer, response=response, checkpoint_step=self.state["checkpoint_step"])) self.save() async def initialize(self, resume): self.ds, _ = await dataset(8 if self.name == "grpo_gs8" else 64, 8 if self.name == "grpo_gs8" else 1) if resume: for n in range(1, 41): if self.step_path(n).exists(): self.state["completed_steps"] = n else: break await self.load(optimizer=True) await self.replay() self.rebuild_history() return await self.load() fixture = json.loads((ROOT / "fixed-batch.json").read_text()) data = [restore_datum(fixture["stock_datums"][0])] output = await retry(self.tc.forward_async, data, loss_fn="importance_sampling") arrays = [float(x) for row in output.loss_fn_outputs for x in row["logprobs"].data] write_json(self.directory / "initial-forward.json", output) batch = json.loads((ROOT / "initial-batch-tokens.json").read_text()) self.provenance.update(initialization=dict(base_checkpoint=MODEL, checkpoint_uri=INIT, loaded_checkpoint_uri=self.provenance["restores"][-1]["response"].path, load_succeeded=True, checkpoint_immutable=True, checkpoint_scope="full_initialization", source="server_response", evidence_ref=f"full/{self.name}/initialization.json", before_first_update=True), initial_output=dict(available=True, batch_sha256=digest(batch), output_sha256=output_fingerprint(arrays), before_first_update=True, evidence_ref=f"full/{self.name}/initial-forward.json")) self.save() def rebuild_history(self): self.history = [] for n in range(1, self.state["completed_steps"] + 1): step = json.loads(self.step_path(n).read_text()) self.history = (self.history + list(zip(step["features"], step["rewards"])))[-256:] async def update(self, data): start = time.monotonic() fb = await rpc(self.tc.forward_backward_async, data, loss_fn="importance_sampling") if not math.isfinite(fb.metrics["loss:sum"]): raise ValueError("nonfinite training loss") fb_seconds = time.monotonic() - start start = time.monotonic() optim = await rpc(self.tc.optim_step_async, tinker.AdamParams( learning_rate=4e-5, beta1=0.9, beta2=0.95, eps=1e-8)) if optim.metrics.get("update_successful:mean") != 1: raise ValueError("optimizer did not report successful update") return fb, optim, dict(forward_backward_seconds=fb_seconds, optim_seconds=time.monotonic() - start) async def replay(self): for n in range(self.state["checkpoint_step"] + 1, self.state["completed_steps"] + 1): draft = json.loads(self.input_path(n).read_text()) await self.update([restore_datum(d) for d in draft["datums"]]) self.state["extra_replay_updates"] += 1 self.save() print(self.name, "replayed committed update", n, "without resampling", flush=True) async def recover(self, error): self.state["recoveries"].append(dict(error=str(error), completed_steps=self.state["completed_steps"], checkpoint=self.state["checkpoint"], checkpoint_step=self.state["checkpoint_step"])) self.save() await self.sc.close("errored") self.sc = service() await self.load(optimizer=True) await self.replay() async def checkpoint(self, n): name = f"rx-full-{self.name}-step{n}-recovery{len(self.state['recoveries'])}" response = await retry(self.tc.save_state_async, name) self.state.update(checkpoint=response.path, checkpoint_step=n) self.state["checkpoints"].append(dict(step=n, response=response)) self.save() async def step(self, n): if self.state["completed_steps"] >= n: return start = time.monotonic() for attempt in range(4): phase = "save" try: if self.input_path(n).exists(): draft = json.loads(self.input_path(n).read_text()) else: t = time.monotonic() name = f"rx-full-{self.name}-step{n}-attempt{attempt}" sampler = await retry(self.tc.save_weights_and_get_sampling_client_async, name=name) save_seconds = time.monotonic() - t phase = "rollout" groups, features, timings = await timed_rollouts(sampler, list(self.ds.get_batch(n - 1))) phase = "assembly" rewards = [r for g in groups for r in g.get_total_rewards()] if len(rewards) != 64: raise ValueError("not exactly 64 rollouts") adv, baseline, neighbors = (pilot_advantages(groups, features, self.history, bandwidth=1.0) if self.name == "gs1_kernel" else advantages(groups, self.name, features, self.history)) data, _ = datums(groups, adv) masked, _ = assemble_training_data(groups, adv, include_provenance=False) flat = np.concatenate([a.numpy() for a in adv]) draft = dict(step=n, features=features, rewards=rewards, reward_mean=float(np.mean(rewards)), executed_group_size=[len(g.trajectories_G) for g in groups], b_i=baseline.tolist(), advantages=flat.tolist(), kernel_effective_neighbors=neighbors.tolist(), neff_p10=float(np.percentile(neighbors, 10)) if self.name == "gs1_kernel" else None, frac_nonzero_advantage=float(np.mean(np.abs(flat) > 1e-6)), datums=data, masks=[d.loss_fn_inputs["mask"].data for d in masked], sampler_save_name=name, bandwidth=1.0 if self.name == "gs1_kernel" else None, updates_since_sampling=0, timings=dict(timings, save_seconds=save_seconds)) atomic(self.input_path(n), draft) draft = json.loads(self.input_path(n).read_text()) phase = "update" fb, optim, timings = await self.update([restore_datum(d) for d in draft["datums"]]) outputs = self.directory / f"step-{n:03d}-outputs.json" write_json(outputs, fb) record = {k: v for k, v in draft.items() if k not in ("datums", "masks")} record["timings"].update(timings, step_seconds=time.monotonic() - start) record.update(metrics=fb.metrics, optim_response=optim, inputs_sha256=sha(self.input_path(n)), outputs_sha256=sha(outputs), attempts=attempt + 1) atomic(self.step_path(n), record) shutil.copy2(self.step_path(n), self.transient / self.step_path(n).name) self.state["completed_steps"] = n self.history = (self.history + list(zip(draft["features"], draft["rewards"])))[-256:] self.save() print(self.name, "step", n, "reward", record["reward_mean"], "seconds", record["timings"]["step_seconds"], flush=True) if n % 10 == 0: phase = "checkpoint" await self.checkpoint(n) return except Exception as error: if attempt == 3 or not (lost(error) or capacity(error)): raise if phase == "rollout": raise RuntimeError("Partial rollout failure: refusing to resample silently or exceed equal rollout budget") from error if capacity(error): delay = (60, 120, 240)[attempt] print(self.name, "capacity backoff", delay, flush=True) await asyncio.sleep(delay) await self.recover(error) if self.state["completed_steps"] >= n: await self.checkpoint(n) return async def final(self): if self.state.get("final"): return self.state["final"] if self.state["completed_steps"] != 40: raise ValueError("final save requires 40 steps") state = await retry(self.tc.save_state_async, f"rx-full-{self.name}-final") sampling = await retry(self.tc.save_weights_for_sampler_async, name=f"rx-full-{self.name}-final") self.state["final"] = dict(state=state, sampling=sampling) self.save() return self.state["final"] async def score(arm, frozen, builders): start = time.monotonic() final = json.loads(arm.state_path.read_text())["final"] checkpoint, sample_path = final["state"]["path"], final["sampling"]["path"] sampler = await retry(arm.sc.create_sampling_client_async, model_path=sample_path) policy = TinkerTokenCompleter(sampler, max_tokens=512, temperature=0) path = arm.directory / "eval.jsonl" rows = [json.loads(line) for line in path.read_text().splitlines()] if path.exists() else [] if rows: summarize(rows) for row in rows: if row["evaluator_id"] != EVALUATOR or row["checkpoint"] != checkpoint or row["arm"] != arm.name: raise ValueError("existing evaluation identity mismatch") completed = {r["item_id"] for r in rows} if path.exists(): shutil.copy2(path, arm.transient / "eval.jsonl") with path.open("a") as out, (arm.transient / "eval.jsonl").open("a") as mirror: async def item(i, builder): identity = f"gsm8k-test-{i:04d}" if identity in completed: return env = builder.env_thunk() prompt, stop = await env.initial_observation() action = await policy(prompt, stop) message, _ = env.renderer.parse_response(action.tokens) prediction = renderers.get_text_content(message) row = dict(item_id=identity, arm=arm.name, seed=0, evaluator_id=frozen["evaluator_id"], metric="accuracy", prediction=prediction, reference=env.get_reference_answer(), value=int(env.check_answer(prediction)), checkpoint=checkpoint, sampler_checkpoint=sample_path) line = canonical(row) + "\n" out.write(line); out.flush() mirror.write(line); mirror.flush() rows.append(row) if len(rows) % 200 == 0: print(arm.name, "evaluated", len(rows), flush=True) async with asyncio.timeout(600): await asyncio.gather(*(item(i, b) for i, b in enumerate(builders))) if {r["item_id"] for r in rows} != set(frozen["evaluators"][EVALUATOR]["items"]["ids"]): raise ValueError("incomplete evaluation") summary = dict(summarize(rows), arm=arm.name, metric="accuracy", seed=0, evaluator_id=EVALUATOR, checkpoint=checkpoint, predictions_ref=f"full/{arm.name}/eval.jsonl", seconds=time.monotonic() - start) write_json(arm.directory / "evaluation.json", summary) return rows, summary def decision(rows, summaries, frozen): summaries = [dict(s, value=s.get("value", s.get("accuracy"))) for s in summaries] observations = [{key: row[key] for key in ("item_id", "arm", "seed", "metric", "value", "evaluator_id")} for arm_rows in rows for row in arm_rows] path = ROOT / "scored-observations.jsonl" path.write_text("".join(canonical(r) + "\n" for r in sorted(observations, key=lambda r: (r["arm"], r["item_id"])))) rule = copy.deepcopy(json.loads((ROOT / "spec.json").read_text())["decision_rule"]) for criterion in rule["criteria"]: criterion["weights"] = {dict(baseline="grpo_gs8", treatment="gs1_kernel")[a]: w for a, w in criterion["weights"].items()} evaluation = dict(version=1, evaluators=frozen["evaluators"], scores=summaries, comparisons=[dict(metric="accuracy", arms=list(ARMS))]) result = assess(rule, dict(version=1, observations=observations), evaluation) if result["issues"]: raise ValueError(f"Decision evidence invalid: {result['issues']}") indexed = [{r["item_id"]: r["value"] for r in arm_rows} for arm_rows in rows] effects = [indexed[1][i] - indexed[0][i] for i in frozen["evaluators"][EVALUATOR]["items"]["ids"]] write_json(ROOT / "paired-hoeffding.json", paired_interval(effects, -1, 1, 0.05)) write_json(ROOT / "decision-full.json", result) write_json(ROOT / "evaluation-full.json", evaluation) write_json(ROOT / "decision-rule-execution.json", rule) write_json(ROOT / "decision-evidence.json", dict(version=1, observations_file=dict(name=path.name, sha256=sha(path)))) print("DECISION", canonical(result), flush=True) return result def provenance(arms): atomic(ROOT / "initialization_provenance_full.json", dict(version=1, arms=[a.provenance for a in arms])) async def main(resume, wall_seconds): start = time.monotonic() guard_grader() OUT.mkdir(exist_ok=True) arms = [] status = "errored" try: async with asyncio.timeout(wall_seconds): frozen, builders = await freeze() if frozen["evaluator_id"] != EVALUATOR: raise ValueError("frozen evaluator identity changed") calibration = json.loads((ROOT / "activation-pilot-gate-bw1.0.json").read_text()) if not calibration["gate"]["passed"]: raise ValueError("activation gate unresolved") arms = [Arm(name, resume) for name in ARMS] async with asyncio.TaskGroup() as group: for arm in arms: group.create_task(arm.initialize(resume)) provenance(arms) if not resume and len({a.provenance["initial_output"]["output_sha256"] for a in arms}) != 1: raise ValueError("initial arm outputs differ") for n in range(1, 41): async with asyncio.TaskGroup() as group: for arm in arms: group.create_task(arm.step(n)) provenance(arms) atomic(OUT / "progress.json", dict(step=n, elapsed_seconds=time.monotonic() - start, arms={a.name: a.state["completed_steps"] for a in arms})) await asyncio.gather(*(arm.final() for arm in arms)) scored = await asyncio.gather(*(score(arm, frozen, builders) for arm in arms)) result = decision([r for r, s in scored], [s for r, s in scored], frozen) atomic(OUT / "result.json", dict(decision=result, scores=[s for r, s in scored], checkpoints={a.name: a.state["final"] for a in arms}, wall_seconds=time.monotonic() - start)) status = "success" except BaseException as error: atomic(OUT / "failure.json", dict(error=repr(error), elapsed_seconds=time.monotonic() - start)) raise finally: await asyncio.gather(*(a.sc.close(status) for a in arms), return_exceptions=True) atomic(OUT / "timings.json", dict(wall_seconds=time.monotonic() - start, status=status, wall_budget_seconds=wall_seconds, clients_released=True)) print("Full run ended", status, "seconds", time.monotonic() - start, flush=True) if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--resume", action="store_true") parser.add_argument("--wall-seconds", type=float, default=4500) args = parser.parse_args() asyncio.run(main(args.resume, args.wall_seconds))