scripts/preflight.py

#!/usr/bin/env python3
"""Verify the whole setup end-to-end before spending the training budget.

    python3 scripts/preflight.py

Runs nine checks, cheapest first, stopping at the first hard failure. Checks 1-5
are free and offline. Checks 6-9 hit the Tinker API and cost a few cents in
total: one training step on 2 examples, and one short generation.

What it proves, in order:
    1  .env present and TINKER_API_KEY loaded
    2  tinker + tinker_cookbook importable at the expected versions
    3  datasets present and well-formed
    4  reward function still calibrated (catches a broken refactor)
    5  no example exceeds max_length (silent right-truncation would eat claims)
    6  API key authenticates and the target model is actually served
    7  renderer masks the prompt and trains only on the assistant tokens
    8  forward_backward + optim_step complete and return a finite loss
    9  the model generates, and the generation is scorable

Exit code 0 means you are clear to run training/train_sft.py.
"""

import json
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from claim_drafter.config import DEFAULTS, load_env

OK, BAD, WARN = "  [PASS]", "  [FAIL]", "  [WARN]"
_state = {"failed": 0, "warned": 0}


def check(name):
    print("\n%s" % name)


def ok(msg):
    print("%s %s" % (OK, msg))


def warn(msg):
    _state["warned"] += 1
    print("%s %s" % (WARN, msg))


def fail(msg, fatal=True):
    _state["failed"] += 1
    print("%s %s" % (BAD, msg))
    if fatal:
        print("\nPreflight stopped. Fix the above and re-run.")
        sys.exit(1)


AUTH_URL = ("https://tinker.thinkingmachines.dev/services/tinker-prod"
            "/api/v1/auth/token")


def _probe_auth_directly(api_key, base_url=None):
    """Ask the auth endpoint whether the key is good, using stdlib urllib.

    Returns "valid", "rejected", or a short description of the transport failure.
    Deliberately does NOT use the tinker SDK: the point is to take a different
    network path so we can tell a bad key apart from a blocked connection.
    """
    import urllib.error
    import urllib.request

    url = (base_url.rstrip("/") + "/api/v1/auth/token") if base_url else AUTH_URL
    request = urllib.request.Request(
        url, data=b"{}", method="POST",
        headers={"X-API-Key": api_key, "Content-Type": "application/json",
                 "User-Agent": "claim-drafter-preflight/1.0"},
    )
    try:
        with urllib.request.urlopen(request, timeout=45) as response:
            return "valid" if response.status == 200 else "http %d" % response.status
    except urllib.error.HTTPError as e:
        return "rejected" if e.code in (401, 403) else "http %d" % e.code
    except Exception as e:
        return "%s: %s" % (type(e).__name__, str(e)[:80])


def main():
    model = os.environ.get("CLAIM_DRAFTER_MODEL", DEFAULTS["model"])
    print("=" * 68)
    print("Claim drafter preflight   model=%s" % model)
    print("=" * 68)

    # ---------------------------------------------------------------- 1
    check("1. Environment")
    if not os.path.exists(".env"):
        fail(".env not found. Run:  cp .env.example .env   and add your key.")
    load_env()
    key = os.environ["TINKER_API_KEY"]
    if len(key) < 8:
        fail("TINKER_API_KEY looks too short to be real.")
    ok("TINKER_API_KEY loaded (%s…%s, %d chars)" % (key[:4], key[-2:], len(key)))

    # ---------------------------------------------------------------- 2
    check("2. Packages")
    try:
        import tinker
        ok("tinker %s" % getattr(tinker, "__version__", "(version attr missing)"))
    except ImportError:
        fail("tinker not installed.  pip install -r requirements.txt")
    try:
        import tinker_cookbook  # noqa: F401
        ok("tinker_cookbook importable")
    except ImportError:
        fail("tinker_cookbook not installed.  pip install -r requirements.txt")

    # ---------------------------------------------------------------- 3
    check("3. Datasets")
    train_path = "data/sft/train.jsonl"
    if not os.path.exists(train_path):
        fail("%s missing. Rebuild with:  make data   (see docs/pipeline.md)" % train_path)
    with open(train_path) as f:
        rows = [json.loads(l) for l in f if l.strip()]
    ok("%s: %d examples" % (train_path, len(rows)))
    roles = [m["role"] for m in rows[0]["messages"]]
    if roles != ["system", "user", "assistant"]:
        fail("unexpected role sequence %s; expected system/user/assistant" % roles)
    ok("role sequence system/user/assistant")
    dpo_path = "data/dpo/comparisons.jsonl"
    if os.path.exists(dpo_path):
        n = sum(1 for _ in open(dpo_path))
        ok("%s: %d preference pairs" % (dpo_path, n))
    else:
        warn("%s missing — stages 2/3 unavailable until you run make dpo" % dpo_path)

    # ---------------------------------------------------------------- 4
    check("4. Reward function calibration")
    from claim_drafter.rewards import claim_reward
    scores = [claim_reward(r["messages"][2]["content"]) for r in rows[:300]]
    mean = sum(scores) / len(scores)
    if mean < 0.95:
        fail("mean reward on real granted claims is %.3f; expected >0.95. "
             "The checker is broken, not the claims." % mean)
    ok("mean reward on 300 real granted claim sets: %.3f" % mean)
    broken = claim_reward("1. A thing.\n\n3. The thing of claim 9.")
    if broken >= mean:
        fail("reward does not punish malformed claims (%.3f vs %.3f)" % (broken, mean))
    ok("malformed claim set scores lower (%.3f)" % broken)

    # ---------------------------------------------------------------- 5
    check("5. Sequence lengths")
    from tinker_cookbook.tokenizer_utils import get_tokenizer
    tokenizer = get_tokenizer(model)
    longest = 0
    for r in rows:
        n = sum(len(tokenizer.encode(m["content"])) for m in r["messages"])
        longest = max(longest, n)
    if longest + 32 > DEFAULTS["max_length"]:
        fail("longest example is %d tokens vs max_length %d. The cookbook "
             "truncates silently FROM THE RIGHT, which would cut off the end of "
             "the claim set." % (longest, DEFAULTS["max_length"]))
    ok("longest example %d tokens, under max_length %d — nothing truncates"
       % (longest, DEFAULTS["max_length"]))

    # ---------------------------------------------------------------- 6
    check("6. API authentication and model availability")
    import tinker
    try:
        service_client = tinker.ServiceClient()
        caps = service_client.get_server_capabilities()
        served = [m.model_name for m in caps.supported_models]
        ok("authenticated; server lists %d models" % len(served))
    except Exception as e:
        # Distinguish "bad key" from "cannot reach Tinker". The SDK reports both
        # as APIConnectionError, which is genuinely confusing. Re-test the auth
        # endpoint with stdlib urllib: it takes a different network path than the
        # SDK's HTTP client, so a 200 here means the key is fine and the problem
        # is transport (corporate proxy, VPN, sandboxed shell).
        verdict = _probe_auth_directly(key)
        if verdict == "valid":
            fail("Your API KEY IS VALID, but the tinker SDK could not connect.\n"
                 "         This is a transport problem, not a credentials problem.\n"
                 "         Do NOT set TINKER_BASE_URL -- the default endpoint is correct.\n"
                 "         The usual cause is pyqwest 0.7.0, which builds an EMPTY CA root\n"
                 "         store and fails every TLS handshake. requirements.txt pins\n"
                 "         pyqwest<0.7.0 and claim_drafter/transport.py repairs an already-\n"
                 "         installed 0.7.0 at runtime. If you still see this, check that\n"
                 "         load_env() ran, and try TINKER_TRANSPORT=httpx.\n"
                 "         Underlying error: %s" % e)
        elif verdict == "rejected":
            fail("TINKER_API_KEY was REJECTED by the server (HTTP 401).\n"
                 "         Check the value in .env against your Tinker dashboard.")
        else:
            fail("Could not reach Tinker at all (%s).\n"
                 "         Underlying error: %s" % (verdict, e))
    if model not in served:
        fail("%s is not in the served model list. Available Qwen models: %s"
             % (model, [m for m in served if "Qwen" in m][:8]))
    ok("%s is served" % model)

    # ---------------------------------------------------------------- 7
    check("7. Renderer and loss masking")
    from tinker_cookbook import renderers
    from tinker_cookbook.renderers import TrainOnWhat
    from tinker_cookbook.supervised.data import conversation_to_datum
    renderer = renderers.get_renderer(DEFAULTS["renderer"], tokenizer)
    ok("renderer %s" % DEFAULTS["renderer"])
    datum = conversation_to_datum(
        rows[0]["messages"], renderer, DEFAULTS["max_length"],
        train_on_what=TrainOnWhat.LAST_ASSISTANT_MESSAGE,
    )
    weights = datum.loss_fn_inputs["weights"]
    weights = getattr(weights, "data", weights)
    weights = list(weights)
    nonzero = sum(1 for w in weights if w > 0)
    if nonzero == 0:
        fail("every loss weight is zero — nothing would be trained on")
    if nonzero == len(weights):
        fail("every loss weight is non-zero — the prompt is not masked")
    ok("%d/%d tokens carry loss (prompt masked, assistant trained)"
       % (nonzero, len(weights)))

    # ---------------------------------------------------------------- 8
    check("8. One training step (costs a few cents)")
    training_client = service_client.create_lora_training_client(
        base_model=model, rank=DEFAULTS["lora_rank"])
    ok("LoRA training client created (rank %d)" % DEFAULTS["lora_rank"])
    batch = [
        conversation_to_datum(r["messages"], renderer, DEFAULTS["max_length"],
                              train_on_what=TrainOnWhat.LAST_ASSISTANT_MESSAGE)
        for r in rows[:2]
    ]
    n_tokens = sum(d.model_input.length for d in batch)
    # Submit both, then await both, so the server pipelines them.
    fwd_bwd_future = training_client.forward_backward(batch, loss_fn="cross_entropy")
    optim_future = training_client.optim_step(
        tinker.AdamParams(learning_rate=1e-4, beta1=0.9, beta2=0.95, eps=1e-8))
    fwd_bwd_result = fwd_bwd_future.result()
    optim_future.result()
    ok("forward_backward + optim_step completed on %d tokens" % n_tokens)
    try:
        logprobs = fwd_bwd_result.loss_fn_outputs[0]["logprobs"]
        vals = list(getattr(logprobs, "data", logprobs))
        if not all(v == v for v in vals[:50]):     # NaN check
            fail("logprobs contain NaN")
        ok("logprobs finite")
    except (KeyError, IndexError, TypeError):
        warn("could not read logprobs from the result; step still succeeded")

    # ---------------------------------------------------------------- 9
    check("9. Sampling")
    # No `name=`: it is deprecated and has no effect -- these checkpoints are
    # always ephemeral. Use save_weights_for_sampler(name=...) for a durable one.
    sampling_client = training_client.save_weights_and_get_sampling_client()
    prompt = renderer.build_generation_prompt(rows[0]["messages"][:2])
    result = sampling_client.sample(
        prompt=prompt,
        sampling_params=tinker.types.SamplingParams(
            max_tokens=256, temperature=0.0, stop=renderer.get_stop_sequences()),
        num_samples=1,
    ).result()
    message, _ = renderer.parse_response(result.sequences[0].tokens)
    text = renderers.get_text_content(message)
    if not text.strip():
        fail("model returned an empty generation")
    ok("generated %d chars" % len(text))
    print("\n  ---- first 300 chars ----")
    for line in text[:300].splitlines():
        print("  | %s" % line)
    print("  -------------------------")
    warn("untrained LoRA output will look poor — this only proves the path works")

    print("\n" + "=" * 68)
    if _state["failed"]:
        print("PREFLIGHT FAILED (%d)" % _state["failed"])
        sys.exit(1)
    print("PREFLIGHT PASSED%s" % (" with %d warning(s)" % _state["warned"]
                                  if _state["warned"] else ""))
    print("Next:  python3 training/train_sft.py --epochs 2")
    print("=" * 68)


if __name__ == "__main__":
    main()