Skip to content

Commit 18a4531

Browse files
rasmusfaberclaude
andcommitted
feat(network_sandbox): support per-service network modes
`network_sandbox` applied one `network_mode` to every service, so the mixed shape that matters for network isolation -- a connected agent container next to an isolated one holding reference solutions or hidden tests -- could not be expressed. That blocks a live smoke test of METR/hawk#1248, which makes `network_mode: none` actually isolate a service on Kubernetes. Add an optional `service_network_modes` mapping. Per-service entries win for the services they name; `network_mode` covers the rest and still defaults to "none", so every existing call emits byte-identical compose. Contradictions (unknown service, unknown mode, empty `services`, or a `network_mode` fully shadowed by `service_network_modes`) raise ValueError instead of resolving to a silently-picked winner. `bridge_network_pattern` works per service too: the top-level `networks` block is emitted whenever at least one service joins the shared network, and a "none" service is simply left off it -- `network_mode: none` plus `networks` is rejected by Hawk and by the inspect_k8s_sandbox converter. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
1 parent f15f30f commit 18a4531

3 files changed

Lines changed: 299 additions & 13 deletions

File tree

‎CLAUDE.md‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,7 @@ The library has 9 modules in `inspect_test_utils/`:
6666
| `sometimes_fails_setup` | Randomly fails during setup phase | `sample_count`, `fail_setup_on_epochs`, `failure_rate` |
6767
| `sometimes_fails_scoring` | Randomly fails during scoring phase | `sample_count`, `fail_score_on_epochs`, `failure_rate` |
6868
| `configurable_sandbox` | K8s sandbox with resource configuration; optional crash injector (`crash_after`) for agent-agnostic deployment resume tests | `cpu`, `memory`, `storage`, `gpu`, `gpu_model`, `allow_internet`, `crash_after`, `crash_hard` |
69-
| `network_sandbox` | Docker network mode testing | `network_mode` ("none", "bridge", "bridge_network_pattern"), `services` |
69+
| `network_sandbox` | Docker network mode testing, uniform or per-service | `network_mode` ("none", "bridge", "bridge_network_pattern"), `services`, `service_network_modes` |
7070

7171
## HardcodedModelAPI
7272

@@ -181,8 +181,20 @@ inspect eval inspect_test_utils/network_sandbox \
181181
inspect eval inspect_test_utils/network_sandbox \
182182
--task-arg network_mode=bridge_network_pattern \
183183
--task-arg 'services=["default", "server"]'
184+
185+
# Mixed: a connected agent container next to an isolated one
186+
inspect eval inspect_test_utils/network_sandbox \
187+
--task-arg 'services=["default", "solution"]' \
188+
--task-arg 'service_network_modes={"default": "bridge", "solution": "none"}'
184189
```
185190

191+
`service_network_modes` overrides `network_mode` for the services it names;
192+
`network_mode` covers the rest (and defaults to `none`). Passing a
193+
`service_network_modes` that covers *every* service alongside `network_mode`, or
194+
naming a service that is not in `services`, raises `ValueError`. A service set to
195+
`none` is never put on the shared network — `network_mode: none` plus `networks`
196+
is rejected by Hawk and by the `inspect_k8s_sandbox` converter.
197+
186198
## Test Utilities
187199

188200
For writing tests against Inspect AI evaluations:

‎inspect_test_utils/tasks.py‎

Lines changed: 106 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import asyncio
22
import os
33
import tempfile
4-
from typing import Any, Literal
4+
from typing import Any, Literal, get_args
55

66
import yaml
77
from inspect_ai import Task, task
@@ -18,6 +18,11 @@
1818
use_critic_role,
1919
)
2020

21+
NetworkMode = Literal["none", "bridge", "bridge_network_pattern"]
22+
"""How a ``network_sandbox`` service is attached to the network."""
23+
24+
NETWORK_MODES: tuple[NetworkMode, ...] = get_args(NetworkMode)
25+
2126

2227
@task
2328
def sometimes_fails_setup(
@@ -406,25 +411,116 @@ def say_hello_with_tools(
406411
)
407412

408413

414+
def _resolve_network_modes(
415+
services: list[str],
416+
network_mode: NetworkMode | None,
417+
service_network_modes: dict[str, NetworkMode] | None,
418+
) -> dict[str, NetworkMode]:
419+
"""Resolve the effective network mode of every ``network_sandbox`` service.
420+
421+
Args:
422+
services: The service names, in compose order.
423+
network_mode: The uniform mode, applied to every service that has no
424+
per-service entry. ``None`` falls back to ``"none"`` (today's default).
425+
service_network_modes: Per-service overrides.
426+
427+
Returns:
428+
A mode for each service in ``services``.
429+
430+
Raises:
431+
ValueError: If ``services`` is empty, if a mode is not a valid
432+
``NetworkMode``, if ``service_network_modes`` names a service that is
433+
not in ``services``, or if ``network_mode`` is given while
434+
``service_network_modes`` already covers every service (the uniform
435+
mode could never apply, so the caller is contradicting themselves).
436+
"""
437+
if not services:
438+
raise ValueError("services must not be empty")
439+
440+
overrides = service_network_modes or {}
441+
442+
invalid = {
443+
name: mode for name, mode in overrides.items() if mode not in NETWORK_MODES
444+
}
445+
if network_mode is not None and network_mode not in NETWORK_MODES:
446+
invalid = {"network_mode": network_mode, **invalid}
447+
if invalid:
448+
raise ValueError(
449+
f"invalid network mode(s) {invalid}; must be one of {list(NETWORK_MODES)}"
450+
)
451+
452+
unknown = sorted(name for name in overrides if name not in services)
453+
if unknown:
454+
raise ValueError(
455+
f"service_network_modes names unknown service(s) {unknown}; "
456+
+ f"services are {services}"
457+
)
458+
459+
if network_mode is not None and all(name in overrides for name in services):
460+
raise ValueError(
461+
f"network_mode={network_mode!r} is contradicted by a "
462+
+ "service_network_modes that covers every service: the uniform mode "
463+
+ "could never apply. Pass one or the other."
464+
)
465+
466+
return {name: overrides.get(name, network_mode or "none") for name in services}
467+
468+
409469
@task
410470
def network_sandbox(
411471
sample_count: int = 1,
412-
network_mode: Literal["none", "bridge", "bridge_network_pattern"] | None = None,
472+
network_mode: NetworkMode | None = None,
413473
services: list[str] | None = None,
474+
service_network_modes: dict[str, NetworkMode] | None = None,
414475
) -> Task:
415476
"""Task for testing network configurations in Docker sandbox.
416477
478+
Every service runs an HTTP server on port 8000, so reachability between
479+
services (and the lack of it) is directly testable from inside the sandbox.
480+
481+
Modes:
482+
- "none": ``network_mode: none`` -- no network at all
483+
- "bridge": ``network_mode: bridge``
484+
- "bridge_network_pattern": joins the shared ``networks: ["shared"]``
485+
bridge network (a top-level ``networks`` block is emitted whenever at
486+
least one service uses this mode)
487+
488+
Precedence: ``service_network_modes[service]`` wins for the services it names;
489+
every other service gets ``network_mode``; if that is ``None`` too, the
490+
service gets ``"none"`` (the historical default). Passing ``network_mode``
491+
*and* a ``service_network_modes`` that covers every service is rejected rather
492+
than silently resolved, as is naming a service that is not in ``services``.
493+
494+
Mixed modes are the point: ``services=["default", "server"]`` with
495+
``service_network_modes={"default": "bridge", "server": "none"}`` gives an
496+
agent container with normal connectivity next to an isolated one -- the shape
497+
a platform's network isolation has to get right. A ``"none"`` service is never
498+
put on the shared network, because ``network_mode: none`` plus ``networks`` is
499+
rejected by Hawk and by the ``inspect_k8s_sandbox`` converter.
500+
417501
Args:
418502
sample_count: Number of samples
419-
network_mode:
420-
- None/"none": No network access
421-
- "bridge": Uses network_mode: bridge
422-
- "bridge_network_pattern": Uses shared bridge network pattern
503+
network_mode: Uniform mode for services without a per-service entry
504+
(default: "none")
423505
services: List of service names (default: ["default"])
506+
service_network_modes: Per-service modes, overriding ``network_mode``.
507+
Keys must be names in ``services``.
508+
509+
Returns:
510+
The configured task.
511+
512+
Raises:
513+
ValueError: On an unknown mode, an unknown service name, an empty
514+
``services``, or a ``network_mode`` fully shadowed by
515+
``service_network_modes``.
424516
"""
425517
if services is None:
426518
services = ["default"]
427519

520+
resolved_modes = _resolve_network_modes(
521+
services, network_mode, service_network_modes
522+
)
523+
428524
compose: dict[str, Any] = {"services": {}}
429525

430526
for service_name in services:
@@ -433,16 +529,14 @@ def network_sandbox(
433529
"entrypoint": ["python", "-m", "http.server", "8000"],
434530
}
435531

436-
if network_mode is None or network_mode == "none":
437-
service_config["network_mode"] = "none"
438-
elif network_mode == "bridge":
439-
service_config["network_mode"] = "bridge"
440-
elif network_mode == "bridge_network_pattern":
532+
if resolved_modes[service_name] == "bridge_network_pattern":
441533
service_config["networks"] = ["shared"]
534+
else:
535+
service_config["network_mode"] = resolved_modes[service_name]
442536

443537
compose["services"][service_name] = service_config
444538

445-
if network_mode == "bridge_network_pattern":
539+
if "bridge_network_pattern" in resolved_modes.values():
446540
compose["networks"] = {"shared": {"driver": "bridge"}}
447541

448542
tmpdir = tempfile.mkdtemp(prefix="inspect_test_utils_network_sandbox_")

‎tests/test_tasks.py‎

Lines changed: 180 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,10 @@
1010
from __future__ import annotations
1111

1212
from collections.abc import Callable
13+
from typing import Any
1314

1415
import pytest
16+
import yaml
1517
from inspect_ai import Task
1618
from inspect_ai.util import CheckpointSampleConfig
1719

@@ -22,6 +24,31 @@ def _checkpoints(task: Task) -> list[CheckpointSampleConfig | None]:
2224
return [sample.checkpoint for sample in task.dataset]
2325

2426

27+
def _compose(task: Task) -> dict[str, Any]:
28+
"""Read back the compose.yaml a docker-sandboxed task wrote to a temp dir."""
29+
sandbox = task.sandbox
30+
assert sandbox is not None and sandbox.type == "docker"
31+
with open(sandbox.config, encoding="utf-8") as f:
32+
return yaml.safe_load(f)
33+
34+
35+
def _service(mode: str) -> dict[str, Any]:
36+
"""The compose service `network_sandbox` emits for a non-shared-network mode."""
37+
return {
38+
"image": "python:3.12-bookworm",
39+
"entrypoint": ["python", "-m", "http.server", "8000"],
40+
"network_mode": mode,
41+
}
42+
43+
44+
SHARED_SERVICE = {
45+
"image": "python:3.12-bookworm",
46+
"entrypoint": ["python", "-m", "http.server", "8000"],
47+
"networks": ["shared"],
48+
}
49+
SHARED_NETWORKS = {"shared": {"driver": "bridge"}}
50+
51+
2552
ROOT_DEFAULT = CheckpointSampleConfig(sandbox_paths={"default": ["/root"]})
2653

2754

@@ -65,6 +92,159 @@ def test_network_sandbox_covers_all_services() -> None:
6592
assert all(c == expected for c in _checkpoints(task))
6693

6794

95+
@pytest.mark.parametrize(
96+
("kwargs", "expected"),
97+
[
98+
pytest.param(
99+
{},
100+
{"services": {"default": _service("none")}},
101+
id="defaults",
102+
),
103+
pytest.param(
104+
{"network_mode": "none", "services": ["default", "server"]},
105+
{"services": {"default": _service("none"), "server": _service("none")}},
106+
id="uniform_none",
107+
),
108+
pytest.param(
109+
{"network_mode": "bridge", "services": ["default", "server"]},
110+
{"services": {"default": _service("bridge"), "server": _service("bridge")}},
111+
id="uniform_bridge",
112+
),
113+
pytest.param(
114+
{
115+
"network_mode": "bridge_network_pattern",
116+
"services": ["default", "server"],
117+
},
118+
{
119+
"services": {"default": SHARED_SERVICE, "server": SHARED_SERVICE},
120+
"networks": SHARED_NETWORKS,
121+
},
122+
id="uniform_bridge_network_pattern",
123+
),
124+
],
125+
)
126+
def test_network_sandbox_uniform_modes_unchanged(
127+
kwargs: dict[str, Any], expected: dict[str, Any]
128+
) -> None:
129+
# The pre-existing signature (network_mode alone, services alone, both,
130+
# neither) must keep emitting exactly the compose it emitted before
131+
# per-service modes existed -- hawk pins a released version of this package.
132+
assert _compose(tasks.network_sandbox(**kwargs)) == expected
133+
134+
135+
def test_network_sandbox_mixed_bridge_and_none() -> None:
136+
"""The case that matters: a connected service next to an isolated one."""
137+
compose = _compose(
138+
tasks.network_sandbox(
139+
services=["default", "server"],
140+
service_network_modes={"default": "bridge", "server": "none"},
141+
)
142+
)
143+
assert compose == {
144+
"services": {"default": _service("bridge"), "server": _service("none")}
145+
}
146+
147+
148+
def test_network_sandbox_network_mode_fills_unlisted_services() -> None:
149+
# Precedence: a per-service entry wins, network_mode covers the rest.
150+
compose = _compose(
151+
tasks.network_sandbox(
152+
network_mode="bridge",
153+
services=["default", "server", "solution"],
154+
service_network_modes={"solution": "none"},
155+
)
156+
)
157+
assert compose == {
158+
"services": {
159+
"default": _service("bridge"),
160+
"server": _service("bridge"),
161+
"solution": _service("none"),
162+
}
163+
}
164+
165+
166+
def test_network_sandbox_shared_network_with_isolated_service() -> None:
167+
# An isolated service must not carry a `networks` key alongside
168+
# `network_mode: none` -- hawk and inspect_k8s_sandbox both reject that
169+
# combination -- so it is simply left off the shared network.
170+
compose = _compose(
171+
tasks.network_sandbox(
172+
services=["default", "server", "solution"],
173+
service_network_modes={
174+
"default": "bridge_network_pattern",
175+
"server": "bridge_network_pattern",
176+
"solution": "none",
177+
},
178+
)
179+
)
180+
assert compose == {
181+
"services": {
182+
"default": SHARED_SERVICE,
183+
"server": SHARED_SERVICE,
184+
"solution": _service("none"),
185+
},
186+
"networks": SHARED_NETWORKS,
187+
}
188+
189+
190+
def test_network_sandbox_no_shared_block_when_nobody_joins() -> None:
191+
# Overriding every service off the shared network must not leave a dangling
192+
# top-level networks block behind.
193+
compose = _compose(
194+
tasks.network_sandbox(
195+
services=["default"],
196+
service_network_modes={"default": "none"},
197+
)
198+
)
199+
assert "networks" not in compose
200+
201+
202+
def test_network_sandbox_per_service_modes_keep_checkpoint_paths() -> None:
203+
task = tasks.network_sandbox(
204+
services=["default", "server"],
205+
service_network_modes={"default": "bridge", "server": "none"},
206+
)
207+
expected = CheckpointSampleConfig(
208+
sandbox_paths={"default": ["/root"], "server": ["/root"]}
209+
)
210+
assert all(c == expected for c in _checkpoints(task))
211+
212+
213+
@pytest.mark.parametrize(
214+
"kwargs",
215+
[
216+
pytest.param(
217+
{
218+
"services": ["default"],
219+
"service_network_modes": {"server": "none"},
220+
},
221+
id="unknown_service",
222+
),
223+
pytest.param(
224+
{"services": ["default"], "service_network_modes": {"default": "host"}},
225+
id="unknown_mode",
226+
),
227+
pytest.param(
228+
{"network_mode": "overlay"},
229+
id="unknown_uniform_mode",
230+
),
231+
pytest.param(
232+
{
233+
"network_mode": "bridge",
234+
"services": ["default", "server"],
235+
"service_network_modes": {"default": "bridge", "server": "none"},
236+
},
237+
id="uniform_mode_fully_shadowed",
238+
),
239+
pytest.param({"services": []}, id="empty_services"),
240+
],
241+
)
242+
def test_network_sandbox_rejects_contradictory_input(kwargs: dict[str, Any]) -> None:
243+
# Contradictions fail fast rather than resolving to a silently-picked winner.
244+
with pytest.raises(ValueError):
245+
tasks.network_sandbox(**kwargs)
246+
247+
68248
def test_configurable_sandbox_crash_after_arms_setup() -> None:
69249
"""``crash_after`` wires a crash injector onto the task's ``setup``, so the
70250
task crashes whichever agent an eval-set pairs with it -- no solver chaining.

0 commit comments

Comments
 (0)