"""Coordinator / worker protocol over TCP for real multi-process and multi-node benchmarks (Phase 8). Stages (mandate ยง11.1): A workers fit local and split MLEs; B workers upload corrected local estimates (+ optional local Hessian); C coordinator averages and broadcasts theta^(0); D workers evaluate score (+ Hessian) at theta^(0); E upload; F coordinator solves the Newton system; G repeat D-F q times. Every message is length-prefixed; bytes, message counts and wall-clock per stage are logged. Data never move: each worker simulates its own block from a seed (as if the data were already local). """ from __future__ import annotations import json, pickle, socket, struct, time, platform, os import numpy as np from .likelihoods import make_model, gather_blocks from .algorithms import local_split_fits, worker_corrected_estimates, pilot_estimator, newton_step, quantize_relative HDR = struct.Struct("!Q") def send(sock, obj): payload = pickle.dumps(obj, protocol=5) sock.sendall(HDR.pack(len(payload)) + payload) return len(payload) + HDR.size def recv(sock): n = _recv_exact(sock, HDR.size) (ln,) = HDR.unpack(n) payload = _recv_exact(sock, ln) return pickle.loads(payload), ln + HDR.size def _recv_exact(sock, n): buf = bytearray() while len(buf) < n: chunk = sock.recv(min(1 << 20, n - len(buf))) if not chunk: raise ConnectionError("socket closed") buf.extend(chunk) return bytes(buf) def theta_true(p, model): th = np.array([0.3] + [0.5 * (-1) ** j / (1 + 0.15 * j) for j in range(p - 1)]) return th * (0.5 if model == "poisson" else 1.0) if p > 1 else np.array([1.0]) # ----------------------------------------------------------------------------- worker def worker_main(host, port, worker_id): s = socket.create_connection((host, port), timeout=600) s.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) send(s, {"type": "hello", "id": worker_id, "host": socket.gethostname(), "machine": platform.machine(), "python": platform.python_version(), "numpy": np.__version__, "cpu": os.cpu_count()}) model = None; db = None while True: msg, _ = recv(s) t = msg["type"] if t == "fit": t0 = time.perf_counter() model = make_model(msg["model"], msg["p"]); p = model.p rng = np.random.default_rng(msg["seed"]) th0 = theta_true(p, msg["model"]) data = model.simulate(rng, msg["m"], th0) if msg["model"] != "exponential" else {"x": rng.exponential(1.0, size=msg["m"])} db = {k: v[None] for k, v in data.items()} t_gen = time.perf_counter() - t0 t1 = time.perf_counter(); fits = local_split_fits(model, db, msg["splits"]); t_fit = time.perf_counter() - t1 J, ok = worker_corrected_estimates(fits, msg["splits"]) out = {"type": "estimate", "id": worker_id, "J": J[0], "ok": bool(ok[0]), "t_gen": t_gen, "t_fit": t_fit, "n_fits": int(sum(msg["splits"])), "iters": {int(sk): int(fits[sk][1].size) for sk in fits}} if msg.get("send_local_hessian"): out["H_local"] = fits[1][2][0, 0] send(s, out) elif t == "derivatives": t0 = time.perf_counter() th = np.asarray(msg["theta"]); g, H = model.score_hessian(th[None], db) out = {"type": "derivatives", "id": worker_id, "g": g[0], "t_der": time.perf_counter() - t0} if msg.get("want_hessian", True): H0 = H[0] if msg.get("bits"): H0 = quantize_relative(H0, msg["bits"]) out["H"] = H0[np.triu_indices(len(th))] # upper triangle only send(s, out) elif t == "bye": s.close(); return # ----------------------------------------------------------------------------- coordinator class Coordinator: def __init__(self, port, K, model, p, N, splits=(1, 2), q=1, hessian_mode="exact", seed=1, bits=None, log=None): self.port, self.K, self.model_name, self.p, self.N = port, K, model, p, N self.splits, self.q, self.hessian_mode, self.seed, self.bits = tuple(splits), q, hessian_mode, seed, bits self.m = N // K; self.log = log or (lambda *a: None) self.stats = {"bytes_up": 0, "bytes_down": 0, "messages_up": 0, "messages_down": 0, "stages": {}} def _bcast(self, msg): t0 = time.perf_counter() for s in self.socks: self.stats["bytes_down"] += send(s, msg); self.stats["messages_down"] += 1 return time.perf_counter() - t0 def _gather(self): out = []; lat = [] t0 = time.perf_counter() for s in self.socks: # sequential drain; latency = time until each worker's message is fully received o, n = recv(s); self.stats["bytes_up"] += n; self.stats["messages_up"] += 1 out.append(o); lat.append(time.perf_counter() - t0) return out, lat def run(self): 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) self.socks = []; hellos = [] t_start = time.perf_counter() while len(self.socks) < self.K: c, _ = srv.accept(); c.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) h, n = recv(c); self.stats["bytes_up"] += n; hellos.append(h); self.socks.append(c) self.socks.sort(key=lambda s: 0); t_connect = time.perf_counter() - t_start hosts = sorted({h["host"] for h in hellos}) self.log(f"{self.K} workers connected from {len(hosts)} hosts in {t_connect:.2f}s") p = self.p # Stage A/B tA = time.perf_counter() for k, s in enumerate(self.socks): # per-worker seeds: K small individual messages self.stats["bytes_down"] += send(s, {"type": "fit", "model": self.model_name, "p": p, "m": self.m, "splits": list(self.splits), "seed": self.seed * 1000003 + k, "send_local_hessian": self.hessian_mode == "stale"}) self.stats["messages_down"] += 1 ests, latA = self._gather() tB = time.perf_counter() J = np.stack([e["J"] for e in ests]); ok = np.array([e["ok"] for e in ests]) t_fit = np.array([e["t_fit"] for e in ests]); t_gen = np.array([e["t_gen"] for e in ests]) theta = pilot_estimator(J) H_stale = np.mean([e["H_local"] for e in ests], axis=0) if self.hessian_mode == "stale" else None 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()), "worker_gen_s_mean": float(t_gen.mean()), "latency_p50": float(np.percentile(latA, 50)), "latency_p90": float(np.percentile(latA, 90)), "latency_p99": float(np.percentile(latA, 99)), "valid_workers": int(ok.sum())} path = [theta.copy()]; diags = [] for step in range(self.q): tC = time.perf_counter() want_H = self.hessian_mode in ("exact", "quantized") 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}) ders, lat = self._gather() tE = time.perf_counter() g = np.mean([d["g"] for d in ders], axis=0) if want_H: Hs = np.zeros((p, p)); iu = np.triu_indices(p) Hu = np.mean([d["H"] for d in ders], axis=0); Hs[iu] = Hu; Hs = Hs + Hs.T - np.diag(np.diag(Hs)); H = Hs elif self.hessian_mode == "single": Hs = np.zeros((p, p)); iu = np.triu_indices(p); Hs[iu] = ders[0]["H"]; H = Hs + Hs.T - np.diag(np.diag(Hs)) else: H = H_stale tS = time.perf_counter(); theta, d = newton_step(theta, g, H); t_solve = time.perf_counter() - tS path.append(theta.copy()); d["grad_norm"] = float(np.linalg.norm(g)); diags.append(d) t_der = np.array([x["t_der"] for x in ders]) 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()), "gather_s": tE - tC - t_b, "solve_s": t_solve, "latency_p50": float(np.percentile(lat, 50)), "latency_p90": float(np.percentile(lat, 90)), "latency_p99": float(np.percentile(lat, 99)), "grad_norm": d["grad_norm"], "cond": d["cond"]} t_total = time.perf_counter() - tA self._bcast({"type": "bye"}) for s in self.socks: s.close() srv.close() 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, "bits": self.bits, "hosts": hosts, "n_hosts": len(hosts), "workers": hellos, "t_connect": t_connect, "t_total_after_connect": t_total, "theta_path": [t.tolist() for t in path], "newton": diags, "stats": self.stats, "coordinator": socket.gethostname(), "critical_path_estimate": float(t_fit.max()) + sum(v["wall"] for k, v in self.stats["stages"].items() if k.startswith("D_F"))} return res def coordinator_run(port, K, model, p, N, splits, q, hessian_mode, seed, bits=None, log=print): return Coordinator(port, K, model, p, N, splits, q, hessian_mode, seed, bits, log).run()