1#!/usr/bin/env python32"""Coordinator-side distributed benchmark (Gate E): compares local bounded residency vs remote execution of a placed subset.3--workers name=host:port[,name=host:port] --placement {none,odd,half_layers,popular:<traces>} --capacity C4Writes JSONL records (same schema as bench_decode + remote stats) and a correctness check vs mlx-lm on the first prompt.5"""6import argparse, json, os, random, socket, subprocess, sys, time, statistics7import mlx.core as mx, numpy as np, yaml8from pathlib import Path9sys.path.insert(0, str(Path(__file__).resolve().parents[2]))10from mlx_lm import stream_generate11from mlx_lm.utils import load_model, load_tokenizer12from titan.runtime import TitanModel13from transport.client import WorkerClient14sys.path.insert(0, str(Path(__file__).resolve().parent))15from bench_decode import run_one, env_record161718def make_placement(kind, spec, workers, traces=None):19 L, E = spec.num_layers, spec.num_experts20 names = list(workers)21 if kind == "none":22 return {}23 if kind == "odd":24 return {(l, e): names[0] for l in range(L) for e in range(E) if e % 2 == 1}25 if kind == "all_remote":26 return {(l, e): names[(l * E + e) % len(names)] for l in range(L) for e in range(E)}27 if kind.startswith("popular:"): # least popular experts go remote (coordinator keeps the hot set), split round-robin over workers28 d = kind.split(":", 1)[1]; idx = json.load(open(os.path.join(d, "index.json"))); cnt = np.zeros((L, E))29 for p in idx["prompts"]:30 if p["split"] != "fit": continue31 z = np.load(os.path.join(d, p["name"] + ".npz")); t = z["topk"]32 for l in range(L):33 cnt[l] += np.bincount(t[:, l].reshape(-1), minlength=E)34 pl = {}35 for l in range(L):36 order = np.argsort(cnt[l]) # ascending popularity37 for i, e in enumerate(order[: E // 2]):38 pl[(l, int(e))] = names[i % len(names)]39 return pl40 raise ValueError(kind)414243def main():44 ap = argparse.ArgumentParser()45 ap.add_argument("--model", required=True)46 ap.add_argument("--workers", default="")47 ap.add_argument("--token-file", default="~/.titan_token")48 ap.add_argument("--placement", default="none")49 ap.add_argument("--capacity", type=int, required=True)50 ap.add_argument("--predictor", default="none"); ap.add_argument("--m", type=int, default=8); ap.add_argument("--traces", default=None)51 ap.add_argument("--workloads", default=str(Path(__file__).resolve().parents[1] / "configs" / "workloads_v2.yaml"))52 ap.add_argument("--families", default="general,code"); ap.add_argument("--split", default="heldout")53 ap.add_argument("--max-tokens", type=int, default=64); ap.add_argument("--repeats", type=int, default=2); ap.add_argument("--warmup", type=int, default=1)54 ap.add_argument("--seed", type=int, default=0); ap.add_argument("--tag", default=""); ap.add_argument("--out", required=True)55 a = ap.parse_args()56 tm = TitanModel(a.model, capacity=a.capacity, predictor=a.predictor, m=a.m, traces_dir=a.traces)57 token = open(os.path.expanduser(a.token_file), "rb").read().strip()58 workers = {}59 for spec in [w for w in a.workers.split(",") if w]:60 name, hp = spec.split("="); host, port = hp.rsplit(":", 1)61 workers[name] = WorkerClient(host, int(port), token, tm.spec.model_type, name=name)62 placement = make_placement(a.placement, tm.spec, workers, a.traces)63 rtts = {n: w.ping(10) for n, w in workers.items()}64 if placement:65 assign = {}66 for (l, e), n in placement.items():67 assign.setdefault(n, {}).setdefault(l, []).append(e)68 for n, asg in assign.items():69 t = time.perf_counter(); r = workers[n].load(asg); print(f"worker {n}: loaded {r['loaded']} experts in {time.perf_counter()-t:.1f}s")70 rx = tm.attach_workers(workers, placement)71 else:72 rx = None73 # correctness vs mlx-lm on one prompt74 refm, _ = load_model(Path(a.model)); tok = tm.tokenizer75 ids = mx.array(tok.encode("The capital of Canada is"))[None]76 ref = refm(ids); mx.eval(ref); out = tm(ids); mx.eval(out)77 dlogit = float(mx.max(mx.abs(out.astype(mx.float32) - ref.astype(mx.float32))))78 del refm79 wl = yaml.safe_load(open(a.workloads)); prompts = [(f, p) for f in a.families.split(",") for p in wl["families"][f][a.split]]80 rng = random.Random(a.seed); order = []81 for rep in range(a.warmup + a.repeats):82 ps = prompts[:]; rng.shuffle(ps); order += [(rep, f, p) for f, p in ps]83 base = {"backend": "titan-distributed", "capacity": a.capacity, "placement": a.placement, "workers": list(workers), "rtt_ms": rtts,84 "n_remote_experts": len(placement), "predictor": a.predictor, "m": a.m, "max_tokens": a.max_tokens, "split": a.split, "tag": a.tag,85 "model": os.path.realpath(a.model), "cmd": " ".join(sys.argv), "env": env_record(), "seed": a.seed, "dlogit_vs_mlx_lm": dlogit}86 os.makedirs(os.path.dirname(a.out), exist_ok=True)87 with open(a.out, "a") as f:88 for rep, fam, p in order:89 tm.reset_stats(); mx.reset_peak_memory()90 if rx: rx.remote_pairs = rx.local_pairs = rx.layers_remote = rx.failures = 0; rx.remote_wait = 0.091 r = run_one(tm.model, tok, p, a.max_tokens)92 cs = tm.cache_stats()93 rec = dict(base, warmup=rep < a.warmup, repeat=rep - a.warmup, family=fam, prompt=p[:60], ts=time.strftime("%Y-%m-%dT%H:%M:%S"), **r,94 cache_hit_rate=cs["hit_rate"], cache_misses=cs["misses"], stall_s=cs["stall_time_s"], remote=rx.stats() if rx else None)95 f.write(json.dumps(rec) + "\n"); f.flush()96 print(f"dist cap={a.capacity} {a.placement:10s} rep={rep} {fam:10s} ttft={r['ttft_ms']:.0f}ms decode={r['gen_tps']:.1f} tok/s hit={cs['hit_rate']:.3f} stall={cs['stall_time_s']:.2f}s"97 + (f" remote_wait={rx.remote_wait:.2f}s remote_pairs={rx.remote_pairs}" if rx else ""))98 for w in workers.values(): w.close()99100101if __name__ == "__main__":102 main()103