1"""Two-process exactness test: a worker on localhost holds half of every layer's experts; the coordinator executes those2remotely and the rest locally. Logits must equal mlx-lm's bit-for-bit (same kernels, same bytes, same aggregation order)."""3import os, secrets, socket, subprocess, sys, time, json4import mlx.core as mx5import pytest6from pathlib import Path7from mlx_lm.utils import load_model, load_tokenizer8from titan.runtime import TitanModel9from transport.client import WorkerClient1011PROMPTS = ["The capital of Canada is", "def fibonacci(n):\n \"\"\"Return the n-th Fibonacci number.\"\"\"\n"]1213@pytest.fixture(scope="module")14def worker(model_dir, tmp_path_factory):15 tok = secrets.token_hex(32).encode()16 tf = tmp_path_factory.mktemp("tok") / "token"; tf.write_bytes(tok)17 port = 9400 + os.getpid() % 10018 p = subprocess.Popen([sys.executable, "-m", "transport.worker", "--model", model_dir, "--listen", f"127.0.0.1:{port}", "--token-file", str(tf), "--allow", "127.0.0.1"],19 cwd=str(Path(__file__).resolve().parents[1]), stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)20 for _ in range(100):21 try:22 s = socket.create_connection(("127.0.0.1", port), timeout=0.5); s.close(); break23 except OSError:24 time.sleep(0.2)25 yield {"port": port, "token": tok, "proc": p}26 p.terminate()2728def test_remote_half_experts_bit_identical(model_dir, worker):29 refm, _ = load_model(Path(model_dir)); tok = load_tokenizer(Path(model_dir))30 tm = TitanModel(model_dir, capacity=max(8, tm_cap(model_dir)))31 E = tm.spec.num_experts; L = tm.spec.num_layers32 placement = {(l, e): ("w0" if e % 2 == 1 else None) for l in range(L) for e in range(E)} # odd experts remote33 w = WorkerClient("127.0.0.1", worker["port"], worker["token"], tm.spec.model_type, name="w0")34 rtt = w.ping(5)35 w.load({l: [e for e in range(E) if e % 2 == 1] for l in range(L)})36 rx = tm.attach_workers({"w0": w}, placement)37 rec = {"rtt_ms_localhost": rtt, "prompts": []}38 for p in PROMPTS:39 ids = mx.array(tok.encode(p))[None]40 ref = refm(ids); mx.eval(ref)41 tm.reset_stats()42 out = tm(ids); mx.eval(out)43 d = float(mx.max(mx.abs(out.astype(mx.float32) - ref.astype(mx.float32))))44 rec["prompts"].append({"prompt": p[:30], "max_abs_dlogit": d, "argmax_equal": bool(mx.array_equal(mx.argmax(out, -1), mx.argmax(ref, -1))), "remote": rx.stats()})45 assert d <= 1e-3, (p, d)46 # greedy decode token-identical47 from mlx_lm import generate48 a = generate(refm, tok, prompt=PROMPTS[0], max_tokens=24, verbose=False)49 b = generate(tm.model, tok, prompt=PROMPTS[0], max_tokens=24, verbose=False)50 rec["greedy_equal"] = a == b; rec["greedy_ref"] = a; rec["greedy_titan"] = b51 assert a == b52 os.makedirs("benchmarks/raw/tests", exist_ok=True)53 json.dump(rec, open(f"benchmarks/raw/tests/distributed_localhost_{tm.spec.model_type}_{time.strftime('%Y%m%dT%H%M%S')}.json", "w"), indent=2)54 w.close()5556def tm_cap(model_dir):57 from titan.registry import load_spec58 return load_spec(model_dir).num_experts // 259