# Preflight probe: each arm's exact optimizer config, one forward_backward + optim_step on one example, shared Spindle.
import os, json, time, traceback, tinker, datasets
from tinker import types
from tinker_cookbook import renderers, tokenizer_utils
from tinker_cookbook.supervised.data import conversation_to_datum
MODEL = "Qwen/Qwen3.5-2B"
hdr = {"Modal-Key": os.environ["MODAL_PROXY_TOKEN_ID"], "Modal-Secret": os.environ["MODAL_PROXY_TOKEN_SECRET"]}
sc = tinker.ServiceClient(base_url=os.environ["TINKER_BASE_URL"], api_key=os.environ["TINKER_API_KEY"], default_headers_overrides=hdr)
tok = tokenizer_utils.get_tokenizer(MODEL)
rend = renderers.get_renderer("qwen3_5_disable_thinking", tok)
row = datasets.load_dataset("HuggingFaceH4/no_robots", split="train")[0]
d = conversation_to_datum(row["messages"], rend, 2048)
log = []
def rec(**k):
    k["t"] = time.time(); log.append(k); print(json.dumps(k)[:800], flush=True)
def err(e):
    c = e.__cause__
    return {"type": type(e).__name__, "msg": str(e)[:600], "cause": (type(c).__name__ + ": " + str(c)[:600]) if c else None}
for arm, kw, params in [
    ("dimuon", {"optimizer": types.DimuonOptimizerConfig()}, types.DimuonParams(learning_rate=1e-3, beta1=0.9, grad_clip_norm=0.0)),
    ("adam", {}, types.AdamParams(learning_rate=1e-4, beta1=0.9, beta2=0.95, eps=1e-8)),
]:
    for seedkw in ({"seed": 0}, {}):
        try:
            tc = sc.create_lora_training_client(base_model=MODEL, rank=32, **seedkw, **kw)
            rec(arm=arm, stage="create", seed_kw=seedkw, ok=True)
        except Exception as e:
            rec(arm=arm, stage="create", seed_kw=seedkw, ok=False, error=err(e)); continue
        try:
            fb = tc.forward_backward([d], loss_fn="cross_entropy").result()
            rec(arm=arm, stage="forward_backward", ok=True, metrics=fb.metrics)
            r = tc.optim_step(params).result()
            rec(arm=arm, stage="optim_step", params_type=type(params).__name__, ok=True, resp=str(r)[:300])
        except Exception as e:
            rec(arm=arm, stage="fb/optim_step", params_type=type(params).__name__, ok=False, error=err(e))
        break
json.dump(log, open("probe_shared_log.json", "w"), indent=1)