1"""Coordinator-side connection to one expert worker."""2from __future__ import annotations3import os, socket, time4import numpy as np5from transport.protocol import Framer678class WorkerClient:9 def __init__(self, host: str, port: int, token: bytes, model_type: str, timeout: float = 30.0, name: str = None):10 self.host, self.port, self.name = host, port, name or host11 self.sock = socket.create_connection((host, port), timeout=timeout)12 self.sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)13 self.f = Framer(self.sock, token)14 self.f.send("hello", {"proto": "TX1", "model_type": model_type})15 kind, h, _ = self.f.recv()16 if kind != "hello_ack":17 raise ConnectionError(f"{host}: {h}")18 self.info = h19 self.exec_calls = 0; self.exec_wall = 0.0; self.rows_sent = 02021 def ping(self, n=5):22 rtts = []23 for _ in range(n):24 t = time.perf_counter(); self.f.send("ping", {}); self.f.recv(); rtts.append((time.perf_counter() - t) * 1e3)25 return rtts2627 def load(self, assignments):28 self.f.send("load", {"assignments": {str(l): list(map(int, es)) for l, es in assignments.items()}})29 kind, h, _ = self.f.recv()30 if kind != "load_reply":31 raise RuntimeError(h)32 return h3334 def exec_async(self, layer: int, rows: np.ndarray, pairs, dtype: str = "float16"):35 """Send an exec request (non-blocking on the reply). rows: float16, or uint16 bit patterns when dtype == 'bfloat16'."""36 self.f.send("exec", {"layer": int(layer), "rows": int(rows.shape[0]), "dim": int(rows.shape[1]), "dtype": dtype, "pairs": [[int(r), int(e)] for r, e in pairs]}, rows.tobytes())37 self._t0 = time.perf_counter(); self.rows_sent += rows.shape[0]3839 def exec_wait(self):40 kind, h, payload = self.f.recv()41 if kind != "exec_reply":42 raise RuntimeError(h)43 self.exec_calls += 1; self.exec_wall += time.perf_counter() - self._t044 return np.frombuffer(payload, dtype=np.uint16 if h.get("dtype") == "bfloat16" else np.float16).reshape(h["n"], -1), h4546 def stats(self):47 return {"worker": self.name, "exec_calls": self.exec_calls, "exec_wall_s": self.exec_wall, "bytes_sent": self.f.bytes_sent, "bytes_recv": self.f.bytes_recv}4849 def close(self):50 self.sock.close()51