1"""Transport layer for E2/E6: TCP framing with explicit byte accounting, a remote worker that receives caches/texts and2runs the receiver model, and a client. Loopback and LAN are both supported (labelled by the caller)."""3from __future__ import annotations4import socket, struct, json, time, sys, traceback56MAGIC = b"WP30"78def send_frame(sock: socket.socket, header: dict, payload: bytes) -> tuple[int, float]:9 h = json.dumps(header).encode("utf-8")10 t0 = time.perf_counter()11 sock.sendall(MAGIC + struct.pack("!II", len(h), len(payload)) + h + payload)12 return len(h) + len(payload) + 12, (time.perf_counter() - t0) * 10001314def _recvall(sock, n):15 buf = bytearray()16 while len(buf) < n:17 chunk = sock.recv(min(1 << 20, n - len(buf)))18 if not chunk:19 raise ConnectionError("socket closed")20 buf += chunk21 return bytes(buf)2223def recv_frame(sock: socket.socket) -> tuple[dict, bytes, float]:24 t0 = time.perf_counter()25 hdr = _recvall(sock, 12)26 assert hdr[:4] == MAGIC, "bad magic"27 hl, pl = struct.unpack("!II", hdr[4:])28 h = json.loads(_recvall(sock, hl).decode("utf-8"))29 p = _recvall(sock, pl)30 return h, p, (time.perf_counter() - t0) * 10003132class Client:33 def __init__(self, host: str, port: int, timeout: float = 600.0):34 self.sock = socket.create_connection((host, port), timeout=timeout)35 self.sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)36 def call(self, header: dict, payload: bytes = b"") -> tuple[dict, bytes, dict]:37 t0 = time.perf_counter()38 nbytes, send_ms = send_frame(self.sock, header, payload)39 h, p, recv_ms = recv_frame(self.sock)40 return h, p, {"transport_bytes": nbytes, "send_ms": send_ms, "roundtrip_ms": (time.perf_counter() - t0) * 1000, "reply_bytes": len(p) + 12}41 def close(self):42 self.sock.close()4344def serve(port: int, model_repo: str) -> int:45 """Remote receiver worker. Commands: ping | run_cache (payload = serialized cache) | run_text | prefill_time."""46 from .backends.mlx_backend import MLXModel, memory_snapshot, reset_peak47 from .cache.serialize import deserialize_cache, codec_decompress, IncompatibleCache48 from .cache.sparse import cache_tensor_bytes49 import mlx.core as mx50 m = MLXModel(model_repo)51 srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)52 srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)53 srv.bind(("0.0.0.0", port)); srv.listen(4)54 print(json.dumps({"worker": "ready", "model": model_repo, "port": port, "revision": m.info.revision}), flush=True)55 while True:56 conn, addr = srv.accept()57 conn.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)58 try:59 while True:60 h, p, recv_ms = recv_frame(conn)61 t_recv_done = time.perf_counter()62 cmd = h.get("cmd")63 reply = {"cmd": cmd, "recv_ms": recv_ms, "payload_bytes": len(p)}64 try:65 if cmd == "ping":66 reply["model_revision"] = m.info.revision; reply["config_hash"] = m.info.config_hash67 elif cmd == "shutdown":68 send_frame(conn, {"ok": True}, b""); conn.close(); srv.close(); return 069 elif cmd == "run_text":70 ids = h["token_ids"]71 reset_peak()72 g = m.generate_text(ids, max_tokens=h.get("max_tokens", 48))73 reply.update({"text": g["text"], "prefill_ms": g["prefill_ms"], "ttft_ms": g["ttft_ms"], "decode_ms": g["decode_ms"], "output_tokens": g["output_tokens"], "peak_bytes": int(mx.get_peak_memory())})74 elif cmd == "run_cache":75 codec = h.get("codec", "none")76 t0 = time.perf_counter()77 raw, dec_ms = codec_decompress(p, codec)78 try:79 cache, man = deserialize_cache(raw, m.info, h.get("prefix_hash"))80 except IncompatibleCache as e:81 reply.update({"status": "incompatible", "error": str(e)}); send_frame(conn, reply, b""); continue82 load_ms = (time.perf_counter() - t0) * 100083 reset_peak()84 g = m.generate_from_cache(cache, h["query_ids"], max_tokens=h.get("max_tokens", 48))85 reply.update({"status": "success", "text": g["text"], "decompress_ms": dec_ms, "load_ms": load_ms, "ttft_ms": g["ttft_ms"], "decode_ms": g["decode_ms"],86 "output_tokens": g["output_tokens"], "cache_tensor_bytes": cache_tensor_bytes(cache) , "peak_bytes": int(mx.get_peak_memory())})87 elif cmd == "prefill_cache":88 # worker builds a cache from token ids and returns it serialized (for pull-style transfer)89 from .cache.serialize import build_manifest, serialize_cache, codec_compress90 ids = h["token_ids"]91 c, t = m.prefill(ids)92 man = build_manifest(m.info, ids, c, h.get("document_id", "?"), "B1", {})93 payload = serialize_cache(c, man)94 payload, cms = codec_compress(payload, h.get("codec", "none"))95 reply.update({"prefill_ms": t["prefill_ms"], "serialize_ms": cms, "prefix_hash": man["prefix_token_hash"]})96 send_frame(conn, reply, payload); continue97 else:98 reply["error"] = "unknown cmd"99 except Exception as e:100 reply.update({"status": "implementation_error", "error": f"{type(e).__name__}: {e}", "trace": traceback.format_exc()[-800:]})101 reply["server_ms"] = (time.perf_counter() - t_recv_done) * 1000102 send_frame(conn, reply, b"")103 except (ConnectionError, AssertionError, OSError):104 pass105 finally:106 try: conn.close()107 except Exception: pass108