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: np.ndarray, pairs, dtype: str = "float16"):44 """rows: [R, D] activations (float16, or uint16 bit patterns of bfloat16); pairs: list of (row, expert).45 Returns [len(pairs), D] outputs in the same dtype (bfloat16 returned as uint16 bit patterns)."""46 t0 = time.perf_counter()47 tab = self.table(layer)48 experts = [e for _, e in pairs]49 slots = tab.ensure(list(dict.fromkeys(experts)))50 slot_of = dict(zip(dict.fromkeys(experts), slots))51 x = mx.array(rows).view(mx.bfloat16) if dtype == "bfloat16" else mx.array(rows) # [R, D]52 ridx = mx.array([r for r, _ in pairs], dtype=mx.uint32)53 sidx = mx.array([slot_of[e] for e in experts], dtype=mx.uint32)54 xs = x[ridx] # [P, D]55 xe = mx.expand_dims(xs, (-2, -3)) # [P, 1, 1, D]56 b = tab.buf57 def qmm(h, m):58 return mx.gather_qmm(h, b[f"{m}.weight"], b[f"{m}.scales"], b.get(f"{m}.biases"), rhs_indices=sidx.reshape(-1, 1),59 transpose=True, group_size=self.gs, bits=self.bits)60 up = qmm(xe, "up"); gate = qmm(xe, "gate")61 y = qmm(swiglu(gate, up), "down").squeeze(-2).squeeze(-2) # [P, D]62 out = np.array(y.view(mx.uint16)) if dtype == "bfloat16" else np.array(y.astype(mx.float16))63 self.exec_count += 1; self.exec_time += time.perf_counter() - t064 return out656667def serve(args):68 token = open(os.path.expanduser(args.token_file), "rb").read().strip()69 allow = set(args.allow.split(",")) if args.allow else None70 w = Worker(args.model, args.capacity)71 host, port = args.listen.rsplit(":", 1)72 srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM); srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)73 srv.bind((host, int(port))); srv.listen(4)74 print(f"[worker] {socket.gethostname()} listening on {args.listen}, model {w.spec.model_type}, cap {w.cap}", flush=True)75 while True:76 conn, addr = srv.accept()77 if allow and addr[0] not in allow:78 print(f"[worker] refused {addr}", flush=True); conn.close(); continue79 conn.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)80 conn.settimeout(args.timeout)81 threading.Thread(target=handle, args=(conn, addr, token, w), daemon=True).start()828384def handle(conn, addr, token, w: Worker):85 f = Framer(conn, token)86 try:87 kind, h, _ = f.recv()88 if kind != "hello" or h.get("proto") != "TX1" or h.get("model_type") != w.spec.model_type:89 f.send("error", {"msg": "bad hello"}); conn.close(); return90 f.send("hello_ack", {"host": socket.gethostname(), "num_layers": w.spec.num_layers, "num_experts": w.spec.num_experts,91 "cap": w.cap, "resident": {str(l): t.resident() for l, t in w.tables.items()}})92 print(f"[worker] coordinator {addr} connected", flush=True)93 while True:94 kind, h, payload = f.recv()95 if kind == "ping":96 f.send("pong", {"t": time.time()})97 elif kind == "load":98 n = w.load(h["assignments"])99 f.send("load_reply", {"loaded": n, "resident": {str(l): t.resident() for l, t in w.tables.items()}})100 elif kind == "exec":101 R, D = h["rows"], h["dim"]; dt = h.get("dtype", "float16")102 x = np.frombuffer(payload, dtype=np.uint16 if dt == "bfloat16" else np.float16).reshape(R, D)103 out = w.execute(h["layer"], x, [tuple(p) for p in h["pairs"]], dt)104 f.send("exec_reply", {"layer": h["layer"], "n": out.shape[0], "dtype": dt, "exec_ms": w.exec_time * 1e3}, out.tobytes())105 else:106 f.send("error", {"msg": f"unexpected {kind}"})107 except (ConnectionError, socket.timeout, OSError) as e:108 print(f"[worker] {addr} closed: {e}", flush=True)109 finally:110 conn.close()111112113if __name__ == "__main__":114 ap = argparse.ArgumentParser()115 ap.add_argument("--model", required=True)116 ap.add_argument("--listen", default="0.0.0.0:9411")117 ap.add_argument("--token-file", default="~/.titan_token")118 ap.add_argument("--allow", default=None, help="comma-separated coordinator IPs allowed to connect")119 ap.add_argument("--capacity", type=int, default=None)120 ap.add_argument("--timeout", type=float, default=600.0)121 serve(ap.parse_args())122