shadow: extract render_shadow_yaml host builders

radiken · claude · radiken · commit 08cf069070ff · 2026-07-16T16:09:16.000+01:00
Pull the peer env and the peer / bootstrap / publisher host construction out of
render_shadow_yaml into small pure functions so each is unit-testable. Output is
unchanged; adds tests for the extracted builders.

Co-Authored-By: Claude Opus 4.8 (1M context) &lt;noreply@anthropic.com&gt;
diff --git a/src/deployments/shadow/builders.py b/src/deployments/shadow/builders.py
@@ -42,36 +42,18 @@
 )
 
 
-def render_shadow_yaml(
+def _peer_env(
     *,
     num_nodes: int,
-    sim_stop_time_s: int,
-    publisher_start_s: int,
-    connect_to: int = 2,
-    muxer: Literal["yamux", "mplex", "quic"] = "yamux",
-    discovery: Literal["static", "kad-dht"] = "static",
-    start_sleep: int = 60,
-    metrics_interval_s: int = 15,
-    seed: int = 1,
-    model_unblocked_syscall_latency: bool = False,
-    strace_logging_mode: str = "off",
-    lsquic_tick_floor_us: int = 0,
-    start_jitter_ms: int = 0,
-    requester_app_path: str = _REQUESTER_APP_PATH,
+    connect_to: int,
+    muxer: str,
+    discovery: str,
+    start_sleep: int,
+    metrics_interval_s: int,
+    lsquic_tick_floor_us: int,
 ) -> dict:
-    """Build the shadow.yaml dict: N peer hosts running `./main` + a publisher host
-    running the pod-api-requester in batch mode against the peers' `/publish`
-    endpoints. The traffic shape (message count, size, pacing) lives in the
-    requester's own config (see `render_publisher_config`), mounted at
-    `{_CONFIG_MOUNT}/{_PUBLISHER_CONFIG}`.
-
-    discovery selects mesh formation: "static" dials CONNECTTO peers by pod-N
-    hostname; "kad-dht" adds a `bootstrap-0` anchor host that peers discover
-    through (Shadow resolves it by hostname, so no k8s Service is needed)."""
-    if connect_to >= num_nodes:
-        raise ValueError(f"connect_to ({connect_to}) must be smaller than num_nodes ({num_nodes}).")
-
-    peer_env = {
+    """Environment for a relay (pod-*) node."""
+    env = {
         "PEERS": str(num_nodes),
         "CONNECTTO": str(connect_to),
         "SHADOWENV": "true",  # env.nim requires the literal string "true"
@@ -82,14 +64,18 @@ def render_shadow_yaml(
     }
     if lsquic_tick_floor_us > 0:
         # Needs the tick-floor test-node image (stock images ignore the env var).
-        peer_env["LSQUIC_TICK_FLOOR_US"] = str(lsquic_tick_floor_us)
+        env["LSQUIC_TICK_FLOOR_US"] = str(lsquic_tick_floor_us)
     if discovery == "kad-dht":
-        peer_env["NODE_ROLE"] = "RoleNormal"
-        peer_env["SERVICE"] = "bootstrap-0"
-    # start_jitter_ms staggers per-pod process start so peers don't wake and dial at
-    # one simulated instant (lockstep wakes force simultaneous-dial collisions that
-    # never occur on real hosts).
-    hosts = {
+        env["NODE_ROLE"] = "RoleNormal"
+        env["SERVICE"] = "bootstrap-0"
+    return env
+
+
+def _peer_hosts(num_nodes: int, peer_env: dict, start_jitter_ms: int) -> dict:
+    """The N relay hosts (pod-0..pod-(N-1)). start_jitter_ms staggers per-pod process
+    start so peers don't wake and dial at one simulated instant (lockstep wakes force
+    simultaneous-dial collisions that never occur on real hosts)."""
+    return {
         f"pod-{i}": {
             "network_node_id": 0,
             "processes": [
@@ -104,35 +90,47 @@ def render_shadow_yaml(
         }
         for i in range(num_nodes)
     }
-    if discovery == "kad-dht":
-        hosts["bootstrap-0"] = {
-            "network_node_id": 0,
-            "processes": [
-                {
-                    "path": "./main",
-                    "start_time": "5s",
-                    "expected_final_state": "running",
-                    "environment": {
-                        "PEERS": str(num_nodes),
-                        "SHADOWENV": "true",
-                        "MUXER": muxer,
-                        "DISCOVERY": "kad-dht",
-                        "NODE_ROLE": "RoleBootstrap",
-                        # the single anchor must accept every node's bootstrap dial,
-                        # so lift its cap above the network size (default is 250).
-                        "MAXCONNECTIONS": str(num_nodes + 100),
-                        "STARTSLEEP": str(start_sleep),
-                        "METRICS_INTERVAL_S": str(metrics_interval_s),
-                        **(
-                            {"LSQUIC_TICK_FLOOR_US": str(lsquic_tick_floor_us)}
-                            if lsquic_tick_floor_us > 0
-                            else {}
-                        ),
-                    },
-                }
-            ],
-        }
-    hosts["publisher"] = {
+
+
+def _bootstrap_host(
+    *,
+    num_nodes: int,
+    muxer: str,
+    start_sleep: int,
+    metrics_interval_s: int,
+    lsquic_tick_floor_us: int,
+) -> dict:
+    """The kad-dht anchor (bootstrap-0); peers discover through it by hostname."""
+    env = {
+        "PEERS": str(num_nodes),
+        "SHADOWENV": "true",
+        "MUXER": muxer,
+        "DISCOVERY": "kad-dht",
+        "NODE_ROLE": "RoleBootstrap",
+        # the single anchor must accept every node's bootstrap dial, so lift its cap
+        # above the network size (default is 250).
+        "MAXCONNECTIONS": str(num_nodes + 100),
+        "STARTSLEEP": str(start_sleep),
+        "METRICS_INTERVAL_S": str(metrics_interval_s),
+    }
+    if lsquic_tick_floor_us > 0:
+        env["LSQUIC_TICK_FLOOR_US"] = str(lsquic_tick_floor_us)
+    return {
+        "network_node_id": 0,
+        "processes": [
+            {
+                "path": "./main",
+                "start_time": "5s",
+                "expected_final_state": "running",
+                "environment": env,
+            }
+        ],
+    }
+
+
+def _publisher_host(publisher_start_s: int, requester_app_path: str) -> dict:
+    """The publisher host: runs the pod-api-requester in batch mode against the peers."""
+    return {
         "network_node_id": 0,
         "processes": [
             {
@@ -146,6 +144,56 @@ def render_shadow_yaml(
             }
         ],
     }
+
+
+def render_shadow_yaml(
+    *,
+    num_nodes: int,
+    sim_stop_time_s: int,
+    publisher_start_s: int,
+    connect_to: int = 2,
+    muxer: Literal["yamux", "mplex", "quic"] = "yamux",
+    discovery: Literal["static", "kad-dht"] = "static",
+    start_sleep: int = 60,
+    metrics_interval_s: int = 15,
+    seed: int = 1,
+    model_unblocked_syscall_latency: bool = False,
+    strace_logging_mode: str = "off",
+    lsquic_tick_floor_us: int = 0,
+    start_jitter_ms: int = 0,
+    requester_app_path: str = _REQUESTER_APP_PATH,
+) -> dict:
+    """Build the shadow.yaml dict: N peer hosts running `./main` + a publisher host
+    running the pod-api-requester in batch mode against the peers' `/publish`
+    endpoints. The traffic shape (message count, size, pacing) lives in the
+    requester's own config (see `render_publisher_config`), mounted at
+    `{_CONFIG_MOUNT}/{_PUBLISHER_CONFIG}`.
+
+    discovery selects mesh formation: "static" dials CONNECTTO peers by pod-N
+    hostname; "kad-dht" adds a `bootstrap-0` anchor host that peers discover
+    through (Shadow resolves it by hostname, so no k8s Service is needed)."""
+    if connect_to >= num_nodes:
+        raise ValueError(f"connect_to ({connect_to}) must be smaller than num_nodes ({num_nodes}).")
+
+    peer_env = _peer_env(
+        num_nodes=num_nodes,
+        connect_to=connect_to,
+        muxer=muxer,
+        discovery=discovery,
+        start_sleep=start_sleep,
+        metrics_interval_s=metrics_interval_s,
+        lsquic_tick_floor_us=lsquic_tick_floor_us,
+    )
+    hosts = _peer_hosts(num_nodes, peer_env, start_jitter_ms)
+    if discovery == "kad-dht":
+        hosts["bootstrap-0"] = _bootstrap_host(
+            num_nodes=num_nodes,
+            muxer=muxer,
+            start_sleep=start_sleep,
+            metrics_interval_s=metrics_interval_s,
+            lsquic_tick_floor_us=lsquic_tick_floor_us,
+        )
+    hosts["publisher"] = _publisher_host(publisher_start_s, requester_app_path)
     # Always render the seed so the run's shadow.yaml records it (Shadow defaults to 1).
     config = {
         "general": {
diff --git a/src/deployments/shadow/tests/test_builders.py b/src/deployments/shadow/tests/test_builders.py
@@ -4,6 +4,10 @@
 from ruamel.yaml import YAML
 
 from src.deployments.shadow.builders import (
+    _bootstrap_host,
+    _peer_env,
+    _peer_hosts,
+    _publisher_host,
     build_configmap,
     build_pvc,
     build_shadow_job,
@@ -69,6 +73,77 @@ def test_request_and_action_cross_references_resolve(self):
         assert cfg["actions"][0]["targets"] == [cfg["targets"][0]["name"]]
 
 
+# --------------------------------------------------------------------------- #
+# host builders (the pieces render_shadow_yaml assembles)
+# --------------------------------------------------------------------------- #
+class TestHostBuilders:
+    def _env(self, **over):
+        base = dict(
+            num_nodes=5,
+            connect_to=2,
+            muxer="yamux",
+            discovery="static",
+            start_sleep=60,
+            metrics_interval_s=15,
+            lsquic_tick_floor_us=0,
+        )
+        base.update(over)
+        return _peer_env(**base)
+
+    def test_peer_env_static_omits_kad_and_tickfloor(self):
+        env = self._env()
+        assert env["MUXER"] == "yamux" and env["DISCOVERY"] == "static"
+        assert "NODE_ROLE" not in env and "SERVICE" not in env
+        assert "LSQUIC_TICK_FLOOR_US" not in env
+
+    def test_peer_env_kad_adds_role_and_service(self):
+        env = self._env(discovery="kad-dht")
+        assert env["NODE_ROLE"] == "RoleNormal"
+        assert env["SERVICE"] == "bootstrap-0"
+
+    def test_peer_env_tickfloor_only_when_set(self):
+        assert (
+            _peer_env(
+                num_nodes=5,
+                connect_to=2,
+                muxer="quic",
+                discovery="static",
+                start_sleep=60,
+                metrics_interval_s=15,
+                lsquic_tick_floor_us=10000,
+            )["LSQUIC_TICK_FLOOR_US"]
+            == "10000"
+        )
+
+    def test_peer_hosts_stagger_start_time_by_jitter(self):
+        hosts = _peer_hosts(3, {"MUXER": "yamux"}, start_jitter_ms=50)
+        assert [hosts[f"pod-{i}"]["processes"][0]["start_time"] for i in range(3)] == [
+            "5000ms",
+            "5050ms",
+            "5100ms",
+        ]
+        assert all(hosts[h]["processes"][0]["expected_final_state"] == "running" for h in hosts)
+
+    def test_bootstrap_host_lifts_connection_cap(self):
+        h = _bootstrap_host(
+            num_nodes=1000,
+            muxer="yamux",
+            start_sleep=60,
+            metrics_interval_s=15,
+            lsquic_tick_floor_us=0,
+        )
+        env = h["processes"][0]["environment"]
+        assert env["NODE_ROLE"] == "RoleBootstrap"
+        assert env["MAXCONNECTIONS"] == "1100"  # num_nodes + 100
+
+    def test_publisher_host_runs_requester_batch(self):
+        h = _publisher_host(90, "/app/api_requester.py")
+        proc = h["processes"][0]
+        assert proc["path"] == "/usr/bin/python3"
+        assert "--mode batch" in proc["args"]
+        assert proc["start_time"] == "90s"
+
+
 # --------------------------------------------------------------------------- #
 # render_shadow_yaml  (peer hosts + publisher host)
 # --------------------------------------------------------------------------- #
