import argparse import asyncio import dataclasses import hashlib import json import time import numpy as np import tinker from research_exchange.restoration import output_fingerprint, validate from run_preflight import preflight from rx_gs1 import MODEL, ROOT, guard_grader, serialized, service, write_json def digest(value): return hashlib.sha256(json.dumps(value, separators=(",", ":"), sort_keys=True, allow_nan=False).encode()).hexdigest() async def create(sc): started = time.monotonic() for delay in (60, 120, 240, None): try: return await sc.create_lora_training_client_async(base_model=MODEL, rank=32, optimizer=tinker.AdamOptimizerConfig()) except Exception as error: message = str(error) if not any(x in message.lower() for x in ("capacity exhausted", "429", "503")): raise with (ROOT / "capacity-retries.log").open("a") as f: f.write(json.dumps(dict(elapsed=time.monotonic()-started, error=message, next_delay=delay)) + "\n") if delay is None or time.monotonic() - started + delay > 600: raise await asyncio.sleep(delay) def restore_datum(record): tokens = [t for chunk in record["model_input"]["chunks"] for t in chunk["tokens"]] return tinker.Datum(model_input=tinker.ModelInput.from_ints(tokens), loss_fn_inputs={k: tinker.TensorData(**v) for k, v in record["loss_fn_inputs"].items()}) async def run_a(): started = time.monotonic() sc = service() status = "errored" timings = {} try: tc = await create(sc) timings["init_create_seconds"] = time.monotonic() - started t = time.monotonic() if (ROOT / "init-save-response.json").exists(): checkpoint = json.loads((ROOT / "init-save-response.json").read_text())["path"] future = await tc.load_state_async(checkpoint) loaded = await future.result_async() write_json(ROOT / "init-resume-load-response.json", serialized(loaded)) saved = None else: future = await tc.save_state_async("rx-gs1-init-run_5207121d26") saved = await future.result_async() write_json(ROOT / "init-save-response.json", serialized(saved)) timings["init_save_seconds"] = time.monotonic() - t info = await tc.get_info_async() write_json(ROOT / "init-info-response.json", serialized(info)) print("Init state:", (ROOT / "init-save-response.json").read_text(), flush=True) t = time.monotonic() await preflight(tc) timings["correctness_seconds"] = time.monotonic() - t status = "success" finally: await sc.close(status) timings["phase_a_total_seconds"] = time.monotonic() - started write_json(ROOT / "timings-a.json", timings) print("Phase A session closed", flush=True) async def run_b(): checkpoint = json.loads((ROOT / "init-save-response.json").read_text())["path"] fixture = json.loads((ROOT / "fixed-batch.json").read_text()) initial_data = [restore_datum(fixture["stock_datums"][0])] token_fixture = [dict(tokens=d.model_input.to_ints(), targets=d.loss_fn_inputs["target_tokens"].data, masks=fixture["stock_masks"][i]) for i, d in enumerate(initial_data)] write_json(ROOT / "initial-batch-tokens.json", token_fixture) batch_hash = digest(token_fixture) provenance = dict(version=1, arms=[], initialization_protocol_override=dict( reason="Authorized after Miles rejected per-model seed=0; seed-rejection.log", per_model_seed_kwarg="omitted", data_shuffle_seed=0, protocol="one never-updated init save_state checkpoint; both arms explicitly load_state that checkpoint")) probes = dict(arms=[]) output_arrays = [] phase_start = time.monotonic() for arm in ("grpo_gs8", "gs1_kernel"): started = time.monotonic() sc = service() status = "errored" evidence = dict(arm=arm, checkpoint_requested=checkpoint) timing = {} try: tc = await create(sc) timing["create_seconds"] = time.monotonic()-started info = await tc.get_info_async() evidence["creation_info"] = serialized(info) t = time.monotonic() future = await tc.load_state_async(checkpoint) loaded = await future.result_async() evidence["load_response"] = serialized(loaded) timing["load_seconds"] = time.monotonic()-t write_json(ROOT / f"{arm}-server-responses.json", evidence) print(arm, "load:", json.dumps(serialized(loaded)), flush=True) t = time.monotonic() future = await tc.forward_async(initial_data, loss_fn="importance_sampling") forward = await future.result_async() arrays = [o["logprobs"].data for o in forward.loss_fn_outputs] output_arrays.append(np.concatenate([np.array(a, dtype=np.float64) for a in arrays])) evidence["initial_forward"] = dict(metrics=forward.metrics, outputs=[{k: serialized(v) for k, v in o.items()} for o in forward.loss_fn_outputs]) write_json(ROOT / f"{arm}-server-responses.json", evidence) output_hash = output_fingerprint(output_arrays[-1].tolist()) timing["initial_forward_seconds"] = time.monotonic()-t observation = dict(source="declared", actual=None, actual_config=None, requested="adamw", requested_config=dict(lr=4e-5, beta1=0.9, beta2=0.95, eps=1e-8), evidence_ref=f"{arm}-server-responses.json", before_next_update=True, reason="get_info/load do not inspect the concrete active optimizer; SDK GetInfoResponse supplies an adamw default even when omitted by server. Raw response fields serialized with exclude_unset=True.") provenance["arms"].append(dict(name=arm, restore_count=1, initialization=dict(base_checkpoint=MODEL, checkpoint_uri=checkpoint, loaded_checkpoint_uri=loaded.path, load_succeeded=True, checkpoint_immutable=True, checkpoint_scope="full_initialization", source="server_response", evidence_ref=f"init-save-response.json;{arm}-server-responses.json", before_first_update=True), optimizer_checks=[dict(observation, after_restore=i) for i in (0, 1)], initial_output=dict(available=True, batch_sha256=batch_hash, output_sha256=output_hash, before_first_update=True, evidence_ref=f"initial-batch-tokens.json;{arm}-server-responses.json"))) write_json(ROOT / "initialization_provenance.json", provenance) records = fixture["stock_datums"] if arm == "grpo_gs8" else fixture["enabled_datums"] index = next(i for i, d in enumerate(records) if any(abs(v) > 1e-6 for v in d["loss_fn_inputs"]["advantages"]["data"])) datum = restore_datum(records[index]) t = time.monotonic() future = await tc.forward_backward_async([datum], loss_fn="importance_sampling") fb = await future.result_async() evidence["probe_forward_backward"] = dict(metrics=fb.metrics, outputs=[{k: serialized(v) for k, v in o.items()} for o in fb.loss_fn_outputs]) loss = fb.metrics.get("loss:sum") if loss is None: loss = sum(o["loss:sum"].data[0] for o in fb.loss_fn_outputs) if not np.isfinite(loss): raise ValueError("nonfinite probe loss") params = tinker.AdamParams(learning_rate=4e-5, beta1=0.9, beta2=0.95, eps=1e-8) future = await tc.optim_step_async(params) optim = await future.result_async() evidence["optim_step_response"] = serialized(optim) write_json(ROOT / f"{arm}-server-responses.json", evidence) timing["probe_update_seconds"] = time.monotonic()-t probes["arms"].append(dict(name=arm, datum_index=index, loss=loss, finite=True, forward_backward_succeeded=True, optim_step_succeeded=True, adam_params=serialized(params), evidence_ref=f"{arm}-server-responses.json", timings=timing)) write_json(ROOT / "probe.json", probes) print(arm, "probe loss:", loss, flush=True) status = "success" finally: await sc.close(status) timing["total_seconds"] = time.monotonic()-started write_json(ROOT / f"timings-{arm}.json", timing) print(arm, "session closed; probe client discarded", flush=True) provenance["paired_outputs"] = dict(equal=bool(np.array_equal(*output_arrays)), max_abs_diff=float(np.max(np.abs(output_arrays[0]-output_arrays[1]))), batch_sha256=batch_hash) write_json(ROOT / "initialization_provenance.json", provenance) spec = json.loads((ROOT / "effective-spec.json").read_text()) write_json(ROOT / "restoration-validation.json", validate(spec, provenance)) write_json(ROOT / "timings-b.json", dict(total_seconds=time.monotonic()-phase_start)) if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("phase", choices=["a", "b"]) args = parser.parse_args() guard_grader() asyncio.run(run_a() if args.phase == "a" else run_b())