1#!/usr/bin/env python32"""Expert worker: holds a set of (layer, expert) pairs resident (loaded from its local copy of the checkpoint) and executes them3on activations sent by the coordinator. Exact: the same gather_qmm + swiglu sequence as the local block, per (row, k) slot.45 python -m transport.worker --model <dir> --listen 0.0.0.0:9411 --token-file ~/.titan_token [--allow 192.168.2.80]6"""7from __future__ import annotations8import argparse, json, os, socket, sys, threading, time9import numpy as np10import mlx.core as mx11sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))12from mlx_lm.models.activations import swiglu13from titan.registry import load_spec14from titan.store import ExpertStore15from titan.slots import SlotTable16from transport.protocol import Framer171819class Worker:20 def __init__(self, model_dir, capacity_per_layer=None):21 self.spec = load_spec(model_dir)22 self.store = ExpertStore(self.spec)23 q = self.spec.quantization24 self.gs, self.bits = q["group_size"], q["bits"]25 self.tables = {}26 self.cap = capacity_per_layer or self.spec.num_experts27 self.lock = threading.Lock()28 self.exec_count = 0; self.exec_time = 0.02930 def table(self, l):31 if l not in self.tables:32 self.tables[l] = SlotTable(l, self.cap, self.store)33 return self.tables[l]3435 def load(self, assignments):36 """assignments: {layer: [experts]} -> make them resident."""37 n = 038 for l, es in assignments.items():39 t = self.table(int(l)); t.ensure([int(e) for e in es]); n += len(es)40 mx.eval(*[b for t in self.tables.values() for b in t.buf.values()])41 return n4243 def execute(self, layer, rows_f16: np.ndarray, pairs):44 """rows_f16: [R, D] activations; pairs: list of (row, expert). Returns [len(pairs), D] float16 outputs."""45 t0 = time.perf_counter()46 tab = self.table(layer)47 experts = [e for _, e in pairs]48 slots = tab.ensure(list(dict.fromkeys(experts)))49 slot_of = dict(zip(dict.fromkeys(experts), slots))50 x = mx.array(rows_f16) # [R, D]51 ridx = mx.array([r for r, _ in pairs], dtype=mx.uint32)52 sidx = mx.array([slot_of[e] for e in experts], dtype=mx.uint32)53 xs = x[ridx] # [P, D]54 xe = mx.expand_dims(xs, (-2, -3)) # [P, 1, 1, D]55 b = tab.buf56 def qmm(h, m):57 return mx.gather_qmm(h, b[f"{m}.weight"], b[f"{m}.scales"], b.get(f"{m}.biases"), rhs_indices=sidx.reshape(-1, 1),58 transpose=True, group_size=self.gs, bits=self.bits)59 up = qmm(xe, "up"); gate = qmm(xe, "gate")60 y = qmm(swiglu(gate, up), "down").squeeze(-2).squeeze(-2) # [P, D]61 out = np.array(y.astype(mx.float16))62 self.exec_count += 1; self.exec_time += time.perf_counter() - t063 return out646566def serve(args):67 token = open(os.path.expanduser(args.token_file), "rb").read().strip()68 allow = set(args.allow.split(",")) if args.allow else None69 w = Worker(args.model, args.capacity)70 host, port = args.listen.rsplit(":", 1)71 srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM); srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)72 srv.bind((host, int(port))); srv.listen(4)73 print(f"[worker] {socket.gethostname()} listening on {args.listen}, model {w.spec.model_type}, cap {w.cap}", flush=True)74 while True:75 conn, addr = srv.accept()76 if allow and addr[0] not in allow:77 print(f"[worker] refused {addr}", flush=True); conn.close(); continue78 conn.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)79 conn.settimeout(args.timeout)80 threading.Thread(target=handle, args=(conn, addr, token, w), daemon=True).start()818283def handle(conn, addr, token, w: Worker):84 f = Framer(conn, token)85 try:86 kind, h, _ = f.recv()87 if kind != "hello" or h.get("proto") != "TX1" or h.get("model_type") != w.spec.model_type:88 f.send("error", {"msg": "bad hello"}); conn.close(); return89 f.send("hello_ack", {"host": socket.gethostname(), "num_layers": w.spec.num_layers, "num_experts": w.spec.num_experts,90 "cap": w.cap, "resident": {str(l): t.resident() for l, t in w.tables.items()}})91 print(f"[worker] coordinator {addr} connected", flush=True)92 while True:93 kind, h, payload = f.recv()94 if kind == "ping":95 f.send("pong", {"t": time.time()})96 elif kind == "load":97 n = w.load(h["assignments"])98 f.send("load_reply", {"loaded": n, "resident": {str(l): t.resident() for l, t in w.tables.items()}})99 elif kind == "exec":100 R, D = h["rows"], h["dim"]101 x = np.frombuffer(payload, dtype=np.float16).reshape(R, D)102 out = w.execute(h["layer"], x, [tuple(p) for p in h["pairs"]])103 f.send("exec_reply", {"layer": h["layer"], "n": out.shape[0], "exec_ms": w.exec_time * 1e3}, out.tobytes())104 else:105 f.send("error", {"msg": f"unexpected {kind}"})106 except (ConnectionError, socket.timeout, OSError) as e:107 print(f"[worker] {addr} closed: {e}", flush=True)108 finally:109 conn.close()110111112if __name__ == "__main__":113 ap = argparse.ArgumentParser()114 ap.add_argument("--model", required=True)115 ap.add_argument("--listen", default="0.0.0.0:9411")116 ap.add_argument("--token-file", default="~/.titan_token")117 ap.add_argument("--allow", default=None, help="comma-separated coordinator IPs allowed to connect")118 ap.add_argument("--capacity", type=int, default=None)119 ap.add_argument("--timeout", type=float, default=600.0)120 serve(ap.parse_args())121