"""Spec preflight probe: rank-1 and rank-32 Qwen2.5-3B LoRA clients, fixed-batch logprob match, one IS fwd/bwd + AdamW step."""
import json, time, sys, logging
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(levelname)s %(message)s")
for n in ("httpx","httpcore"): logging.getLogger(n).setLevel(logging.WARNING)
import tinker
from tinker import types
from common import *

t0 = time.time()
sc = service(); rest = sc.create_rest_client()
out = {"model": MODEL, "renderer": RENDERER, "arms": {}}
clients = {}
def create(rank):
    for attempt in range(8):
        try:
            return sc.create_lora_training_client(MODEL, rank=rank, user_metadata={"rx_run": "run_40338c73ff", "phase": "probe"})
        except Exception as e:
            print("create retry", rank, repr(e)[:300], flush=True); time.sleep(min(60 * 2 ** attempt, 240))
    raise SystemExit("create failed")
for r in (1, 32):
    clients[r] = create(r)
    print("created", r, time.time() - t0, flush=True)
tok = clients[1].get_tokenizer()
datums, spec, batch_sha = fixed_batch(tok)
out["fixed_batch_sha256"] = batch_sha
json.dump(spec, open("ev/fixed_batch.json", "w"))
for r, tc in clients.items():
    a = {}
    a["get_info_at_creation"] = info_dict(tc.get_info())
    f = tc.forward(datums, "cross_entropy").result()
    vals, fp = fingerprint(f); a["initial_output_sha256"] = fp; a["n_logprobs"] = len(vals)
    json.dump(vals, open(f"ev/probe_initial_logprobs_rank{r}.json", "w"))
    a["adapter_error"] = "attempt 1: save_weights_external -> 404 Not Found on shared Spindle (ev/probe_attempt1.log); not retried"
    print("init done", r, time.time() - t0, json.dumps({k: v for k, v in a.items() if k != "get_info_at_creation"})[:600], flush=True)
    # one importance-sampling fwd/bwd on one GSM8K example (reference solution tokens, advantage 1) + AdamW lr=1e-3
    d = datums[0]; tgt = d.loss_fn_inputs["target_tokens"].tolist(); w = d.loss_fn_inputs["weights"].tolist()
    lp = f.loss_fn_outputs[0]["logprobs"].tolist()
    isd = types.Datum(model_input=d.model_input, loss_fn_inputs={"target_tokens": tgt, "logprobs": lp, "advantages": [float(x) for x in w]})
    print("fwd_bwd start", r, time.time()-t0, flush=True); fb = tc.forward_backward([isd], "importance_sampling").result(timeout=900); print("fwd_bwd done", r, time.time()-t0, flush=True)
    st = tc.optim_step(types.AdamParams(learning_rate=1e-3, beta1=0.9, beta2=0.95, eps=1e-8)).result(timeout=600)
    a["fwd_bwd_metrics"] = fb.metrics; a["optim_step_metrics"] = st.metrics
    a["get_info_after_step"] = info_dict(tc.get_info())
    f2 = tc.forward(datums, "cross_entropy").result(); v2, fp2 = fingerprint(f2)
    a["post_step_output_sha256"] = fp2
    a["post_step_max_abs_logprob_change"] = max(abs(x - y) for x, y in zip(vals, v2))
    out["arms"][f"rank{r}"] = a
    print("step done", r, time.time() - t0, flush=True)
out["initial_outputs_match"] = out["arms"]["rank1"]["initial_output_sha256"] == out["arms"]["rank32"]["initial_output_sha256"]
v1 = json.load(open("ev/probe_initial_logprobs_rank1.json")); v32 = json.load(open("ev/probe_initial_logprobs_rank32.json"))
out["initial_logprob_max_abs_diff"] = max(abs(x - y) for x, y in zip(v1, v32))
out["seconds"] = time.time() - t0
json.dump(out, open("ev/probe_results.json", "w"), indent=1, default=str)
print(json.dumps(out, indent=1, default=str)[:4000])