1"""End-to-end test of the TCP protocol with worker subprocesses on the local host."""2import subprocess, sys, threading, time, pathlib, socket, os3import numpy as np4from pamle.protocol import coordinator_run56ROOT = pathlib.Path(__file__).resolve().parents[1]789def _free_port():10 s = socket.socket(); s.bind(("", 0)); port = s.getsockname()[1]; s.close(); return port111213def test_protocol_local_workers():14 port = _free_port(); K = 3; res = {}15 def co():16 res.update(coordinator_run(port, K, "logistic", 4, 3 * 2048, (1, 2), 2, "exact", seed=5, log=lambda *a: None))17 t = threading.Thread(target=co, daemon=True); t.start(); time.sleep(0.3)18 env = dict(os.environ, PYTHONPATH=str(ROOT / "src"))19 procs = [subprocess.Popen([sys.executable, "-m", "pamle.worker", "--connect", f"127.0.0.1:{port}", "--id", str(k)], env=env) for k in range(K)]20 t.join(timeout=120); assert not t.is_alive()21 for pr in procs:22 assert pr.wait(timeout=30) == 023 assert res["stats"]["messages_up"] == K * (1 + 2) and res["stats"]["messages_down"] == K * (1 + 2 + 1)24 assert res["stats"]["bytes_up"] > 0 and res["stats"]["bytes_down"] > 025 assert res["newton"][1]["grad_norm"] < res["newton"][0]["grad_norm"]26 assert len(res["theta_path"]) == 3 and res["n_hosts"] == 127 assert res["stats"]["stages"]["A_B_local_fit_and_upload"]["valid_workers"] == K28