← back to Exo Preflight

preflight.py

180 lines

#!/usr/bin/env python3
"""
exo-preflight — cluster-health gate for the model-routing policy.

Contrarian's fix #2 (DTD verdict A, 2026-08-13): before a pipeline routes a
PRODUCTION-FEEDING job to the exo ring at "$0 local", it must verify the ring
can actually hold the load. A degraded cluster (a node offline, the requested
model not loaded) must NOT silently proceed — the caller queues or escalates to
a paid API and LOGS it.

This script is that gate. It reads exo's /state, judges HEALTHY / DEGRADED /
UNREACHABLE, prints a one-line verdict, and sets its exit code so a shell/py
caller can branch. It touches nothing — pure read.

Exit codes:
  0  HEALTHY      -> safe to route the job to the ring ($0 local)
  2  DEGRADED     -> caller must queue or escalate-to-paid + log the fallback
  3  UNREACHABLE  -> ring not answering; treat as DEGRADED

Usage:
  preflight.py [--node HOST] [--model MODELID] [--expect-nodes N] [--json] [-q]

Examples:
  # Is the ring healthy enough for any bulk job right now?
  preflight.py

  # Is Llama-70B specifically loaded + all its shards ready?
  preflight.py --model llama-3.1-70b

  # In a pipeline:
  if python3 preflight.py --model llama-3.1-70b -q; then
      run_on_ring            # $0 local
  else
      escalate_to_paid_and_log   # per the routing policy
  fi
"""
import argparse, json, os, sys, urllib.request

# The two nodes most likely to be up (this box's loopback first, then the M1 Max
# that is usually online). We try them in order until one answers.
DEFAULT_NODES = ["127.0.0.1", "192.168.1.133", "192.168.1.54"]
PORT = 52415
# Normal live-ring size. The 4th M1 Max is often offline BY DESIGN, so the
# baseline "healthy" count is 3. Override with --expect-nodes or EXO_EXPECT_NODES.
DEFAULT_EXPECT = int(os.environ.get("EXO_EXPECT_NODES", "3"))


def fetch_state(node, timeout=8):
    url = f"http://{node}:{PORT}/state"
    with urllib.request.urlopen(url, timeout=timeout) as r:
        return json.load(r)


def get_state(nodes):
    """Try each node until one returns /state. Returns (state, node) or (None, None)."""
    for n in nodes:
        try:
            return fetch_state(n), n
        except Exception:
            continue
    return None, None


def runner_ready(runners, rid):
    """A runner is ready only if its entry is the RunnerReady variant."""
    v = runners.get(rid)
    return isinstance(v, dict) and "RunnerReady" in v


def model_readiness(state, want):
    """
    Find instances whose modelId matches `want` (exact, or case-insensitive
    substring) and report whether ALL their shard-runners are RunnerReady.
    Returns (matched_model_id_or_None, all_shards_ready_bool, ready_count, total_count).
    """
    runners = state.get("runners", {}) or {}
    instances = state.get("instances", {}) or {}
    w = want.lower()
    matched_id = None
    total = ready = 0
    for inst in instances.values():
        ri = inst.get("MlxRingInstance") or {}
        sa = ri.get("shardAssignments") or {}
        mid = sa.get("modelId", "")
        if mid and (mid == want or w in mid.lower()):
            matched_id = mid
            for rid in (sa.get("runnerToShard") or {}).keys():
                total += 1
                if runner_ready(runners, rid):
                    ready += 1
    return matched_id, (total > 0 and ready == total), ready, total


def main():
    ap = argparse.ArgumentParser(description="exo cluster-health preflight gate")
    ap.add_argument("--node", help="specific node host to query (else auto-try)")
    ap.add_argument("--model", help="require this model loaded + all shards ready")
    ap.add_argument("--expect-nodes", type=int, default=DEFAULT_EXPECT,
                    help=f"min live ring nodes for HEALTHY (default {DEFAULT_EXPECT})")
    ap.add_argument("--json", action="store_true", help="emit machine-readable JSON")
    ap.add_argument("-q", "--quiet", action="store_true", help="suppress human line")
    args = ap.parse_args()

    nodes = [args.node] if args.node else DEFAULT_NODES
    state, served_by = get_state(nodes)

    result = {"verdict": None, "reason": None, "served_by": served_by,
              "live_nodes": None, "expect_nodes": args.expect_nodes,
              "model": args.model, "model_ready": None, "rdma_enabled": None,
              "os_mismatch": None}

    if state is None:
        result["verdict"] = "UNREACHABLE"
        result["reason"] = f"no exo node answered /state on {nodes}"
        emit(result, args)
        return 3

    live = state.get("topology", {}).get("nodes", []) or []
    result["live_nodes"] = len(live)

    ids = state.get("nodeIdentities", {}) or {}
    rdma = state.get("nodeRdmaCtl", {}) or {}
    result["rdma_enabled"] = any(v.get("enabled") for v in rdma.values())
    # Informational: do the live nodes share one macOS build? (RDMA needs it; routing doesn't)
    builds = {ids.get(n, {}).get("osBuildVersion") for n in live if ids.get(n)}
    result["os_mismatch"] = len([b for b in builds if b]) > 1

    reasons = []
    degraded = False

    if len(live) < args.expect_nodes:
        degraded = True
        reasons.append(f"only {len(live)}/{args.expect_nodes} ring nodes live")

    if args.model:
        mid, ready, rc, tc = model_readiness(state, args.model)
        result["model_ready"] = ready
        if mid is None:
            degraded = True
            reasons.append(f"model '{args.model}' not loaded on the ring")
        elif not ready:
            degraded = True
            reasons.append(f"model '{mid}' shards not all ready ({rc}/{tc})")

    if degraded:
        result["verdict"] = "DEGRADED"
        result["reason"] = "; ".join(reasons)
        emit(result, args)
        return 2

    result["verdict"] = "HEALTHY"
    ok = f"{len(live)} nodes live"
    if args.model:
        ok += f"; model '{args.model}' ready"
    result["reason"] = ok
    emit(result, args)
    return 0


def emit(result, args):
    if args.json:
        print(json.dumps(result))
        return
    if args.quiet:
        return
    icon = {"HEALTHY": "✅", "DEGRADED": "⚠️", "UNREACHABLE": "✖"}.get(result["verdict"], "?")
    line = f"{icon} exo {result['verdict']}: {result['reason']}"
    extras = []
    if result.get("os_mismatch"):
        extras.append("OS builds differ across nodes (fine for routing, blocks RDMA)")
    if result.get("rdma_enabled") is False:
        extras.append("RDMA off")
    if extras:
        line += "  [" + "; ".join(extras) + "]"
    print(line, file=(sys.stderr if result["verdict"] != "HEALTHY" else sys.stdout))


if __name__ == "__main__":
    sys.exit(main())