1"""Coordinator / worker protocol over TCP for real multi-process and multi-node benchmarks (Phase 8).23Stages (mandate §11.1): A workers fit local and split MLEs; B workers upload corrected local estimates (+ optional local Hessian);4C coordinator averages and broadcasts theta^(0); D workers evaluate score (+ Hessian) at theta^(0); E upload; F coordinator solves5the Newton system; G repeat D-F q times. Every message is length-prefixed; bytes, message counts and wall-clock per stage are logged.6Data never move: each worker simulates its own block from a seed (as if the data were already local).7"""8from __future__ import annotations9import json, pickle, socket, struct, time, platform, os10import numpy as np11from .likelihoods import make_model, gather_blocks12from .algorithms import local_split_fits, worker_corrected_estimates, pilot_estimator, newton_step, quantize_relative1314HDR = struct.Struct("!Q")151617def send(sock, obj):18 payload = pickle.dumps(obj, protocol=5)19 sock.sendall(HDR.pack(len(payload)) + payload)20 return len(payload) + HDR.size212223def recv(sock):24 n = _recv_exact(sock, HDR.size)25 (ln,) = HDR.unpack(n)26 payload = _recv_exact(sock, ln)27 return pickle.loads(payload), ln + HDR.size282930def _recv_exact(sock, n):31 buf = bytearray()32 while len(buf) < n:33 chunk = sock.recv(min(1 << 20, n - len(buf)))34 if not chunk:35 raise ConnectionError("socket closed")36 buf.extend(chunk)37 return bytes(buf)383940def theta_true(p, model):41 th = np.array([0.3] + [0.5 * (-1) ** j / (1 + 0.15 * j) for j in range(p - 1)])42 return th * (0.5 if model == "poisson" else 1.0) if p > 1 else np.array([1.0])434445# ----------------------------------------------------------------------------- worker46def worker_main(host, port, worker_id):47 s = socket.create_connection((host, port), timeout=600)48 s.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)49 send(s, {"type": "hello", "id": worker_id, "host": socket.gethostname(), "machine": platform.machine(), "python": platform.python_version(),50 "numpy": np.__version__, "cpu": os.cpu_count()})51 model = None; db = None52 while True:53 msg, _ = recv(s)54 t = msg["type"]55 if t == "fit":56 t0 = time.perf_counter()57 model = make_model(msg["model"], msg["p"]); p = model.p58 rng = np.random.default_rng(msg["seed"])59 th0 = theta_true(p, msg["model"])60 data = model.simulate(rng, msg["m"], th0) if msg["model"] != "exponential" else {"x": rng.exponential(1.0, size=msg["m"])}61 db = {k: v[None] for k, v in data.items()}62 t_gen = time.perf_counter() - t063 t1 = time.perf_counter(); fits = local_split_fits(model, db, msg["splits"]); t_fit = time.perf_counter() - t164 J, ok = worker_corrected_estimates(fits, msg["splits"])65 out = {"type": "estimate", "id": worker_id, "J": J[0], "ok": bool(ok[0]), "t_gen": t_gen, "t_fit": t_fit,66 "n_fits": int(sum(msg["splits"])), "iters": {int(sk): int(fits[sk][1].size) for sk in fits}}67 if msg.get("send_local_hessian"):68 out["H_local"] = fits[1][2][0, 0]69 send(s, out)70 elif t == "derivatives":71 t0 = time.perf_counter()72 th = np.asarray(msg["theta"]); g, H = model.score_hessian(th[None], db)73 out = {"type": "derivatives", "id": worker_id, "g": g[0], "t_der": time.perf_counter() - t0}74 if msg.get("want_hessian", True):75 H0 = H[0]76 if msg.get("bits"):77 H0 = quantize_relative(H0, msg["bits"])78 out["H"] = H0[np.triu_indices(len(th))] # upper triangle only79 send(s, out)80 elif t == "bye":81 s.close(); return828384# ----------------------------------------------------------------------------- coordinator85class Coordinator:86 def __init__(self, port, K, model, p, N, splits=(1, 2), q=1, hessian_mode="exact", seed=1, bits=None, log=None):87 self.port, self.K, self.model_name, self.p, self.N = port, K, model, p, N88 self.splits, self.q, self.hessian_mode, self.seed, self.bits = tuple(splits), q, hessian_mode, seed, bits89 self.m = N // K; self.log = log or (lambda *a: None)90 self.stats = {"bytes_up": 0, "bytes_down": 0, "messages_up": 0, "messages_down": 0, "stages": {}}9192 def _bcast(self, msg):93 t0 = time.perf_counter()94 for s in self.socks:95 self.stats["bytes_down"] += send(s, msg); self.stats["messages_down"] += 196 return time.perf_counter() - t09798 def _gather(self):99 out = []; lat = []100 t0 = time.perf_counter()101 for s in self.socks: # sequential drain; latency = time until each worker's message is fully received102 o, n = recv(s); self.stats["bytes_up"] += n; self.stats["messages_up"] += 1103 out.append(o); lat.append(time.perf_counter() - t0)104 return out, lat105106 def run(self):107 srv = socket.socket(); srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1); srv.bind(("0.0.0.0", self.port)); srv.listen(self.K + 8)108 self.socks = []; hellos = []109 t_start = time.perf_counter()110 while len(self.socks) < self.K:111 c, _ = srv.accept(); c.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)112 h, n = recv(c); self.stats["bytes_up"] += n; hellos.append(h); self.socks.append(c)113 self.socks.sort(key=lambda s: 0); t_connect = time.perf_counter() - t_start114 hosts = sorted({h["host"] for h in hellos})115 self.log(f"{self.K} workers connected from {len(hosts)} hosts in {t_connect:.2f}s")116 p = self.p117 # Stage A/B118 tA = time.perf_counter()119 for k, s in enumerate(self.socks): # per-worker seeds: K small individual messages120 self.stats["bytes_down"] += send(s, {"type": "fit", "model": self.model_name, "p": p, "m": self.m, "splits": list(self.splits),121 "seed": self.seed * 1000003 + k, "send_local_hessian": self.hessian_mode == "stale"})122 self.stats["messages_down"] += 1123 ests, latA = self._gather()124 tB = time.perf_counter()125 J = np.stack([e["J"] for e in ests]); ok = np.array([e["ok"] for e in ests])126 t_fit = np.array([e["t_fit"] for e in ests]); t_gen = np.array([e["t_gen"] for e in ests])127 theta = pilot_estimator(J)128 H_stale = np.mean([e["H_local"] for e in ests], axis=0) if self.hessian_mode == "stale" else None129 self.stats["stages"]["A_B_local_fit_and_upload"] = {"wall": tB - tA, "worker_fit_s_mean": float(t_fit.mean()), "worker_fit_s_max": float(t_fit.max()),130 "worker_gen_s_mean": float(t_gen.mean()), "latency_p50": float(np.percentile(latA, 50)),131 "latency_p90": float(np.percentile(latA, 90)), "latency_p99": float(np.percentile(latA, 99)), "valid_workers": int(ok.sum())}132 path = [theta.copy()]; diags = []133 for step in range(self.q):134 tC = time.perf_counter()135 want_H = self.hessian_mode in ("exact", "quantized")136 t_b = self._bcast({"type": "derivatives", "theta": theta, "want_hessian": want_H or (self.hessian_mode == "single"), "bits": self.bits if self.hessian_mode == "quantized" else None})137 ders, lat = self._gather()138 tE = time.perf_counter()139 g = np.mean([d["g"] for d in ders], axis=0)140 if want_H:141 Hs = np.zeros((p, p)); iu = np.triu_indices(p)142 Hu = np.mean([d["H"] for d in ders], axis=0); Hs[iu] = Hu; Hs = Hs + Hs.T - np.diag(np.diag(Hs)); H = Hs143 elif self.hessian_mode == "single":144 Hs = np.zeros((p, p)); iu = np.triu_indices(p); Hs[iu] = ders[0]["H"]; H = Hs + Hs.T - np.diag(np.diag(Hs))145 else:146 H = H_stale147 tS = time.perf_counter(); theta, d = newton_step(theta, g, H); t_solve = time.perf_counter() - tS148 path.append(theta.copy()); d["grad_norm"] = float(np.linalg.norm(g)); diags.append(d)149 t_der = np.array([x["t_der"] for x in ders])150 self.stats["stages"][f"D_F_round{step+1}"] = {"wall": time.perf_counter() - tC, "broadcast_s": t_b, "worker_der_s_mean": float(t_der.mean()), "worker_der_s_max": float(t_der.max()),151 "gather_s": tE - tC - t_b, "solve_s": t_solve, "latency_p50": float(np.percentile(lat, 50)), "latency_p90": float(np.percentile(lat, 90)),152 "latency_p99": float(np.percentile(lat, 99)), "grad_norm": d["grad_norm"], "cond": d["cond"]}153 t_total = time.perf_counter() - tA154 self._bcast({"type": "bye"})155 for s in self.socks:156 s.close()157 srv.close()158 res = {"K": self.K, "N": self.N, "m": self.m, "p": p, "model": self.model_name, "splits": list(self.splits), "q": self.q, "hessian_mode": self.hessian_mode,159 "bits": self.bits, "hosts": hosts, "n_hosts": len(hosts), "workers": hellos, "t_connect": t_connect, "t_total_after_connect": t_total,160 "theta_path": [t.tolist() for t in path], "newton": diags, "stats": self.stats, "coordinator": socket.gethostname(),161 "critical_path_estimate": float(t_fit.max()) + sum(v["wall"] for k, v in self.stats["stages"].items() if k.startswith("D_F"))}162 return res163164165166def coordinator_run(port, K, model, p, N, splits, q, hessian_mode, seed, bits=None, log=print):167 return Coordinator(port, K, model, p, N, splits, q, hessian_mode, seed, bits, log).run()168