diff --git a/.agentenv/config.example.toml b/.agentenv/config.example.toml index 30903567..f12a7cde 100644 --- a/.agentenv/config.example.toml +++ b/.agentenv/config.example.toml @@ -11,7 +11,7 @@ image = "local" # local OCI registry (registry:2 at localhost:5000) secret = "local" # process env vars, or an optional flat YAML file [sandbox] -default = "local" # envs run as containers on the host Docker daemon +default = "local" # envs run as containers on the host Docker daemon, each held to the cpu and memory it was deployed with agent_default = "local" # agents too [agents] diff --git a/src/agent_env/a2a_agent/a2a_agent.py b/src/agent_env/a2a_agent/a2a_agent.py index a5a6a862..6e35e8c6 100644 --- a/src/agent_env/a2a_agent/a2a_agent.py +++ b/src/agent_env/a2a_agent/a2a_agent.py @@ -13,7 +13,12 @@ from agent_env.config import get_config from agent_env.entity_refs import EntityRef -from agent_env.providers.sandbox_providers.sandbox import NetworkPolicy, NetworkPolicyUnsupportedError, port_bindings +from agent_env.providers.sandbox_providers.sandbox import ( + ContainerLimits, + NetworkPolicy, + NetworkPolicyUnsupportedError, + port_bindings, +) from agent_env.artifact.artifacts.docker_image import DockerImageArtifact from agent_env.a2a_agent.object_transfer import ( TRANSFER_TIMEOUT_SECONDS, @@ -509,9 +514,11 @@ async def _run_container(self, image_name: str, a2a_port: int, merged_env: dict[ agent_env = dict(merged_env) if trust_dir is None else {**LOCAL_TRUST_ENV, **merged_env} setup_script = "" network_flag = "".join(f"--add-host {entry} \\\n " for entry in self._sandbox.extra_hosts) + limits = self._sandbox.container_limits + limit_flag = "".join(f"{shlex.quote(arg)} " for arg in limits.docker_args) + "\\\n " if limits else "" if enable_docker: logger.info("enable_docker: starting rootless Docker-in-Docker sidecar for the agent (no host socket)") - setup_script = self._dind_setup_script() + setup_script = self._dind_setup_script(limits) network_flag += f"--network {_DIND_NETWORK} \\\n " agent_env["DOCKER_HOST"] = f"tcp://{_DIND_CONTAINER}:{_DIND_PORT}" @@ -527,7 +534,7 @@ async def _run_container(self, image_name: str, a2a_port: int, merged_env: dict[ {setup_script}docker {"run -d" if trust_dir is None else "create"} \\ --name {self._sandbox.container_name} \\ {publish} \\ - {network_flag}{env_str} \\ + {network_flag}{limit_flag}{env_str} \\ {image_name} > /dev/null {"sleep 2" if trust_dir is None else ""} """ @@ -536,20 +543,23 @@ async def _run_container(self, image_name: str, a2a_port: int, merged_env: dict[ await start_trusting(self._sandbox, self._sandbox.container_name, trust_dir) await asyncio.sleep(2) - def _dind_setup_script(self) -> str: + def _dind_setup_script(self, limits: ContainerLimits | None = None) -> str: # Constants become shell vars so the body stays a raw string (no f-string # brace escaping). TLS off is safe: the tcp endpoint is only reachable on - # the VM-local bridge. + # the VM-local bridge. DIND_LIMITS is expanded unquoted, so it splits into + # flags, or into nothing when the sandbox holds its containers to nothing. env = ( f"DIND_IMAGE={_DIND_IMAGE}\n" f"DIND_CONTAINER={_DIND_CONTAINER}\n" f"DIND_NETWORK={_DIND_NETWORK}\n" f"DIND_PORT={_DIND_PORT}\n" + f"DIND_LIMITS={shlex.quote(' '.join(limits.docker_args)) if limits else ''}\n" ) return env + r"""docker network create "$DIND_NETWORK" >/dev/null 2>&1 || true docker rm -f "$DIND_CONTAINER" >/dev/null 2>&1 || true docker run -d \ --name "$DIND_CONTAINER" \ + $DIND_LIMITS \ --privileged \ --network "$DIND_NETWORK" \ -e DOCKER_TLS_CERTDIR="" \ diff --git a/src/agent_env/env/envs/multi_env.py b/src/agent_env/env/envs/multi_env.py index 2da0ca9f..8c9763d5 100644 --- a/src/agent_env/env/envs/multi_env.py +++ b/src/agent_env/env/envs/multi_env.py @@ -612,6 +612,7 @@ async def _load_from_snapshot(self, snapshot) -> None: host_ips=self._sandbox.host_ips, extra_hosts=self._sandbox.extra_hosts, mcp_server_name=self._mcp_server_name or self.name, + container_limits=self._sandbox.container_limits, ) compose_content = compose_content.replace( " volumes:\n - ./init-schemas.sql:/docker-entrypoint-initdb.d/init-schemas.sql\n", "", diff --git a/src/agent_env/providers/env_providers/env_gateway_provider.py b/src/agent_env/providers/env_providers/env_gateway_provider.py index 6d8e78a6..0e3004e7 100644 --- a/src/agent_env/providers/env_providers/env_gateway_provider.py +++ b/src/agent_env/providers/env_providers/env_gateway_provider.py @@ -34,7 +34,7 @@ ) from agent_env.providers.env_providers.env_provider import _SandboxEnvironmentProvider, _size_kwargs, _tool_names from agent_env.providers.sandbox_providers.modal_sandbox import ModalSandboxProvider -from agent_env.providers.sandbox_providers.sandbox import Sandbox, VmSandbox, port_bindings +from agent_env.providers.sandbox_providers.sandbox import ContainerLimits, Sandbox, VmSandbox, port_bindings from agent_env.providers.sandbox_providers.sandbox_provider import SANDBOX_MODE_CONTAINER, SandboxProvider if TYPE_CHECKING: @@ -213,6 +213,7 @@ def create_docker_compose( host_ips: tuple[str, ...] = (), extra_hosts: tuple[str, ...] = (), mcp_server_name: str | None = None, + container_limits: ContainerLimits | None = None, ) -> str: """Generate docker-compose.yml content for a gateway deployment onto a VM/laptop/arbitrary machine. @@ -224,6 +225,8 @@ def create_docker_compose( gateway_image: Docker image for the gateway. gateway_port: Port to expose for the gateway. website_configs: Optional list of website configurations (frontend + backend pairs). + container_limits: What every service's container is held to, for a sandbox that is not a VM + bounding them all (the local one); None adds nothing. Returns: docker-compose.yml content as a string. @@ -504,6 +507,8 @@ def create_docker_compose( " env-network:", " driver: bridge", ]) + if container_limits is not None: + lines = _with_container_limits(lines, container_limits) return "\n".join(lines) async def deploy( @@ -710,6 +715,7 @@ async def _stand_up(_spec): host_ips=sandbox.host_ips, extra_hosts=sandbox.extra_hosts, mcp_server_name=mcp_server_name, + container_limits=sandbox.container_limits, ) logger.info(f"Generated docker-compose.yml:\n{_redact_compose_secrets(compose_content)}") @@ -1304,6 +1310,20 @@ def _gateway_topology(env: Env) -> _GatewayTopology: _SECRET_ENV_KEY_RE = re.compile(r"^(\s*-\s*)([A-Za-z0-9_]*(?:SECRET|PASSWORD|TOKEN|API_KEY)[A-Za-z0-9_]*)=.*$") +def _with_container_limits(lines: list[str], limits: ContainerLimits) -> list[str]: + """``lines`` with ``limits`` under every service, each a two-space key in the top-level ``services`` section.""" + limited: list[str] = [] + section = None + for line in "\n".join(lines).split("\n"): # an entry may hold several lines + limited.append(line) + if line and not line[0].isspace(): + section = line.rstrip().removesuffix(":") + elif (section == "services" and len(line) > 2 and line.startswith(" ") and not line[2].isspace() + and line.rstrip().endswith(":")): + limited.extend(f" {key}" for key in limits.compose_keys) + return limited + + def _redact_compose_secrets(compose_content: str) -> str: """Mask secret-valued env lines so the logged compose never carries a live secret.""" return "\n".join( diff --git a/src/agent_env/providers/sandbox_providers/local_sandbox.py b/src/agent_env/providers/sandbox_providers/local_sandbox.py index 256fd084..2e9baa36 100644 --- a/src/agent_env/providers/sandbox_providers/local_sandbox.py +++ b/src/agent_env/providers/sandbox_providers/local_sandbox.py @@ -10,6 +10,7 @@ import contextlib import functools import glob +import json import logging import os import platform @@ -28,7 +29,7 @@ from agent_env import config from agent_env.attribution import Attribution -from agent_env.providers.sandbox_providers.sandbox import SANDBOX_LABEL, NetworkPolicy, VmSandbox +from agent_env.providers.sandbox_providers.sandbox import SANDBOX_LABEL, ContainerLimits, NetworkPolicy, VmSandbox from agent_env.providers.sandbox_providers.sandbox_provider import ( SANDBOX_MODE_CONTAINER, SANDBOX_MODE_VM, @@ -74,6 +75,40 @@ def _runs_in_container(cmd: list[str]) -> bool: # Marker dropped in the work dir when this sandbox runs a container, so a later get_sandbox() # (post-run teardown reconstructs the sandbox from disk) knows it owns that container. _CONTAINER_MODE_MARKER = ".agent-container-mode" +# What a sandbox's containers are held to, kept in its work dir for a handle get_sandbox() rebuilds. +_LIMITS_FILE = ".container-limits.json" + + +def _engine_cpus() -> int | None: + """How many CPUs the Docker engine has now (its VM's, on Docker Desktop or Rancher Desktop), or None when Docker + can't be asked. Not cached: the engine can be resized while a process runs.""" + try: + return int(_docker("info", "--format", "{{.NCPU}}")) + except (OSError, RuntimeError, ValueError, subprocess.SubprocessError): + return None + + +def _container_limits(cpu: float, memory: int) -> ContainerLimits: + """``cpu`` and ``memory`` as what a local sandbox's containers are held to. Docker refuses a container more CPUs + than its engine has, so cpu is held to that, with a warning; memory beyond the engine's is accepted.""" + engine_cpus = _engine_cpus() if cpu > 1 else None # every engine has at least one CPU + if engine_cpus is not None and cpu > engine_cpus: + logger.warning( + "Asked for %g CPUs, but the Docker engine has %d; holding this sandbox's containers to %d", cpu, engine_cpus, engine_cpus, + ) + cpu = float(engine_cpus) + return ContainerLimits(cpus=cpu, memory_mib=memory) + + +def _read_limits(work_dir: Path) -> ContainerLimits | None: + path = work_dir / _LIMITS_FILE + if not path.exists(): + return None # a sandbox created before limits were recorded + try: + return ContainerLimits.from_dict(json.loads(path.read_text())) + except (OSError, ValueError, KeyError, TypeError) as e: + logger.warning("Ignoring unreadable container limits at %s: %s", path, e) + return None def _local_sandbox_root_path() -> Path: @@ -126,7 +161,7 @@ class LocalSandbox(VmSandbox): extra_hosts = _EXTRA_HOSTS def __init__(self, exposed_ports: list[int] | None = None, work_dir: Path | None = None, sandbox_id: str | None = None, - port_map: dict[int, int] | None = None): + port_map: dict[int, int] | None = None, container_limits: ContainerLimits | None = None): self.sandbox_id = sandbox_id or f"local-{uuid4().hex[:8]}" # Container port -> host port. Deployments share one host, so the reserved ports # can only be published once; the provider supplies a map to spread them. Callers @@ -144,6 +179,17 @@ def __init__(self, exposed_ports: list[int] | None = None, work_dir: Path | None self._work_dir = work_dir or Path( tempfile.mkdtemp(prefix=f"agent-env-{self.sandbox_id}-", dir=_local_sandbox_root()) ) + if container_limits is not None: + self.container_limits = container_limits + # Unrecorded, a handle a later step rebuilds would start containers unheld, so a failed write fails here. + try: + (self._work_dir / _LIMITS_FILE).write_text(json.dumps(container_limits.to_dict())) + except OSError: + if work_dir is None: # no handle is returned to tear down the folder just made + shutil.rmtree(self._work_dir, ignore_errors=True) + raise + elif work_dir is not None: + self.container_limits = _read_limits(work_dir) def host_port(self, port: int) -> int: """The allocated host port for a published container port (identity if unmapped).""" @@ -397,6 +443,7 @@ async def create_vm( # is I/O, and callers that only need a handle must not pay for it. sandbox = LocalSandbox( port_map={port: _free_host_port() for port in (exposed_ports or [])}, + container_limits=await asyncio.to_thread(_container_limits, cpu, memory), ) sandbox.network_policy = self.effective_network_policy(network_policy) return sandbox diff --git a/src/agent_env/providers/sandbox_providers/sandbox.py b/src/agent_env/providers/sandbox_providers/sandbox.py index 60d0873c..d0676228 100644 --- a/src/agent_env/providers/sandbox_providers/sandbox.py +++ b/src/agent_env/providers/sandbox_providers/sandbox.py @@ -134,6 +134,35 @@ def from_dict(cls, data: dict) -> "NetworkPolicy": ) +@dataclass(frozen=True) +class ContainerLimits: + """The cpu (cores) and memory (MiB) one container started on a sandbox's host is held to.""" + + cpus: float + memory_mib: int + + def __post_init__(self) -> None: + if self.cpus <= 0 or self.memory_mib <= 0: + raise ValueError(f"cpus and memory_mib must be positive; got cpus={self.cpus!r}, memory_mib={self.memory_mib!r}") + + @property + def docker_args(self) -> tuple[str, ...]: + """The ``docker run`` / ``docker create`` flags. Swap is held to the memory limit, so a container can't page past it.""" + return ("--cpus", f"{self.cpus:g}", "--memory", f"{self.memory_mib}m", "--memory-swap", f"{self.memory_mib}m") + + @property + def compose_keys(self) -> tuple[str, ...]: + """The same limits as compose service keys, which ``docker compose up`` applies outside swarm mode.""" + return (f"cpus: {self.cpus:g}", f"mem_limit: {self.memory_mib}m", f"memswap_limit: {self.memory_mib}m") + + def to_dict(self) -> dict: + return {"cpus": self.cpus, "memory_mib": self.memory_mib} + + @classmethod + def from_dict(cls, data: dict) -> "ContainerLimits": + return cls(cpus=float(data["cpus"]), memory_mib=int(data["memory_mib"])) + + class Sandbox(ABC): """Universal sandbox contract — anything that can host a process and expose ports.""" @@ -148,6 +177,8 @@ class Sandbox(ABC): host_ips: tuple[str, ...] = () # ``name:address`` entries the containers started on this sandbox add to their hosts file. extra_hosts: tuple[str, ...] = () + # What each container started on this sandbox is held to; None where the sandbox itself (a VM) is the bound. + container_limits: ContainerLimits | None = None _VM_READY_TIMEOUT = 1200 # wait_for_vm wall-clock budget (s) _VM_READY_POLL_INTERVAL = 30 # sparse polling (s) diff --git a/src/agent_env/providers/sandbox_providers/sandbox_provider.py b/src/agent_env/providers/sandbox_providers/sandbox_provider.py index 68009580..a8094d41 100644 --- a/src/agent_env/providers/sandbox_providers/sandbox_provider.py +++ b/src/agent_env/providers/sandbox_providers/sandbox_provider.py @@ -139,9 +139,11 @@ def _container_args(self, sandbox: VmSandbox, *, image_name: str, port: int, env """The ``docker run`` / ``docker create`` arguments for the sandbox's container.""" env_flags = " \\\n ".join(f"-e {shlex.quote(k)}={shlex.quote(v)}" for k, v in env.items()) extra_args = f"{self.EXTRA_CONTAINER_RUN_ARGS} " if self.EXTRA_CONTAINER_RUN_ARGS else "" + limits = sandbox.container_limits + limit_args = "".join(f"{shlex.quote(arg)} " for arg in limits.docker_args) if limits else "" publish = " ".join(f"-p {spec}" for spec in port_bindings(sandbox.host_ips, sandbox.host_port(port), port)) return ( - f"--name {shlex.quote(sandbox.container_name)} {publish} {extra_args}\\\n " + f"--name {shlex.quote(sandbox.container_name)} {publish} {extra_args}{limit_args}\\\n " f"{env_flags} \\\n {shlex.quote(image_name)}" ) diff --git a/src/agent_env/task_step/task_steps/run_docker_container.py b/src/agent_env/task_step/task_steps/run_docker_container.py index bbf42729..3a4cef46 100644 --- a/src/agent_env/task_step/task_steps/run_docker_container.py +++ b/src/agent_env/task_step/task_steps/run_docker_container.py @@ -219,10 +219,12 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: cap_flags = " ".join(f"--cap-add {shlex.quote(c)}" for c in self.cap_add) volume_flags = " ".join(f"-v {shlex.quote(v)}" for v in self.volumes) host_flags = " ".join(f"--add-host {shlex.quote(entry)}" for entry in sandbox.extra_hosts) + limits = sandbox.container_limits + limit_flags = " ".join(shlex.quote(arg) for arg in limits.docker_args) if limits else "" extra_flags = " ".join(f for f in ( "--privileged" if self.privileged else "", f"--shm-size {shlex.quote(self.shm_size)}" if self.shm_size else "", - device_flags, cap_flags, volume_flags, host_flags, + device_flags, cap_flags, volume_flags, host_flags, limit_flags, ) if f) entrypoint_flag = "" if self.command_override: diff --git a/tst/integration/task_step/test_container_limits_local.py b/tst/integration/task_step/test_container_limits_local.py new file mode 100644 index 00000000..fab2ebe0 --- /dev/null +++ b/tst/integration/task_step/test_container_limits_local.py @@ -0,0 +1,85 @@ +"""The containers a local sandbox starts are held to the cpu and memory it was created with. + +A VM's size bounds its containers; a local sandbox shares this host's Docker, so Docker holds each one. +Needs Docker and a pull of a small public image. +""" + +import shutil +import subprocess +import uuid + +import pytest + +from agent_env.artifact import FileArtifact, FileArtifactUniverse +from agent_env.artifact.store import reset_artifact_store +from agent_env.config import configure, reset_config +from agent_env.providers.sandbox_providers.local_sandbox import LocalSandboxProvider +from agent_env.task.teardown import teardown_run +from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.task_steps.deploy_sandbox import DeploySandboxTaskStep +from agent_env.task_step.task_steps.run_docker_container import RunDockerContainerTaskStep +from tst.util.capabilities import missing_capability_reason + +pytestmark = pytest.mark.integration + +_IMAGE = "mirror.gcr.io/library/nginx:1.27-bookworm" + + +@pytest.fixture +def local_backends(monkeypatch, tmp_path): + if shutil.which("docker") is None or subprocess.run(["docker", "info"], capture_output=True).returncode != 0: + pytest.skip(missing_capability_reason("docker_daemon")) + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("AGENT_ENV_DOCUMENT_STORE", "local") + monkeypatch.setenv("AGENT_ENV_OBJECT_STORE", "local") + monkeypatch.setenv("AGENT_ENV_LOCAL_SANDBOX_DIR", str(tmp_path / "sandboxes")) + configure() + reset_artifact_store() + try: + yield tmp_path + finally: + reset_artifact_store() + reset_config() + + +def _limits(container: str) -> tuple[int, int, int]: + """``(NanoCpus, Memory, MemorySwap)`` Docker holds ``container`` to.""" + out = subprocess.run( + ["docker", "inspect", container, "--format", "{{.HostConfig.NanoCpus}} {{.HostConfig.Memory}} {{.HostConfig.MemorySwap}}"], + capture_output=True, text=True, check=True, + ).stdout.split() + return int(out[0]), int(out[1]), int(out[2]) + + +@pytest.mark.asyncio +async def test_a_container_a_step_starts_on_a_local_sandbox_is_held_to_its_cpu_and_memory(local_backends): + suffix = uuid.uuid4().hex[:8] + (local_backends / "Dockerfile").write_text(f"FROM {_IMAGE}\n") + build_context = FileArtifactUniverse.put(id=f"lim-ctx-{suffix}", file_artifacts={ + "Dockerfile": FileArtifact.put(id=f"lim-df-{suffix}", description="Dockerfile", + file_path=str(local_backends / "Dockerfile")), + }) + container = f"lim-{suffix}" + context = TaskStepContext(instance_id=f"lim-{suffix}") + try: + context = await DeploySandboxTaskStep( + id="box", version=None, sandbox_name="box", sandbox_mode="vm", sandbox_type="local", + cpu=0.5, memory_mb=256, + ).execute(context) + context = await RunDockerContainerTaskStep( + id="ctr", version=None, sandbox_name="box", docker_context_artifact_id=build_context.id, + container_name=container, + ).execute(context) + + assert _limits(container) == (500_000_000, 256 * 2**20, 256 * 2**20) + finally: + await teardown_run(context) + + +@pytest.mark.asyncio +async def test_a_container_mode_local_sandbox_is_held_to_its_cpu_and_memory(local_backends): + sandbox = await LocalSandboxProvider().create_sandbox(image_name=_IMAGE, port=80, env={}, cpu=0.5, memory=256) + try: + assert _limits(sandbox.container_name) == (500_000_000, 256 * 2**20, 256 * 2**20) + finally: + await sandbox.terminate() diff --git a/tst/integration/task_step/test_verify_sandbox_local.py b/tst/integration/task_step/test_verify_sandbox_local.py index 5131b26c..94b59a83 100644 --- a/tst/integration/task_step/test_verify_sandbox_local.py +++ b/tst/integration/task_step/test_verify_sandbox_local.py @@ -38,6 +38,8 @@ _CHECK_PY = "import pathlib, sys\nsys.exit(0 if 'hello' in pathlib.Path('hello.txt').read_text() else 1)\n" # Stands in for an agent's image: a server that stays up, on a Debian userland (bash, GNU find), for any arch. _AGENT_IMAGE = "mirror.gcr.io/library/nginx:1.27-bookworm" +# All an agent's own sandbox keeps on the host: its records, never what a step stages for its container. +_SANDBOX_RECORDS = [".agent-container-mode", ".container-limits.json"] @pytest.fixture @@ -161,7 +163,7 @@ async def test_verify_sandbox_probes_a_local_agent_inside_its_container(local_ag ctx = await step.execute(context) assert [r["result"] for r in ctx.metadata["verifications"]["agent"]["results"]] == [True, True] - assert list(sandbox.work_dir.iterdir()) == [sandbox.work_dir / ".agent-container-mode"] + assert sorted(p.name for p in sandbox.work_dir.iterdir()) == _SANDBOX_RECORDS @pytest.mark.asyncio @@ -195,7 +197,7 @@ async def test_load_artifact_puts_a_universe_in_a_local_agents_container(local_a shown = subprocess.run(["docker", "exec", sandbox.container_name, "cat", "/app/greeting/hello.txt"], capture_output=True, text=True) assert (shown.returncode, shown.stdout) == (0, "hello, world\n") - assert list(sandbox.work_dir.iterdir()) == [sandbox.work_dir / ".agent-container-mode"] + assert sorted(p.name for p in sandbox.work_dir.iterdir()) == _SANDBOX_RECORDS @pytest.mark.asyncio diff --git a/tst/unit/conftest.py b/tst/unit/conftest.py index beaa828a..e8836cf5 100644 --- a/tst/unit/conftest.py +++ b/tst/unit/conftest.py @@ -50,9 +50,11 @@ def _disable_network(): @pytest.fixture(autouse=True) def _docker_answers(monkeypatch): """A bundle run asks whether Docker answers before it deploys containers locally or builds an infra env; unit - tests never reach a daemon, so it answers. Tests of the probe itself patch ``subprocess.run``.""" + tests never reach a daemon, so it answers. Tests of the probe itself patch ``subprocess.run``. A local sandbox + asks the engine how many CPUs it has; here it can't tell, so nothing is held below what was asked for.""" monkeypatch.setattr("agent_env.bundle.preflight.docker_unreachable", lambda: None) monkeypatch.setattr("agent_env.env.bootstrap.docker_unreachable", lambda: None) + monkeypatch.setattr("agent_env.providers.sandbox_providers.local_sandbox._engine_cpus", lambda: None) @pytest.fixture(autouse=True) diff --git a/tst/unit/providers/env_providers/env_gateway_provider_test.py b/tst/unit/providers/env_providers/env_gateway_provider_test.py index 6bf18f27..1881bb49 100644 --- a/tst/unit/providers/env_providers/env_gateway_provider_test.py +++ b/tst/unit/providers/env_providers/env_gateway_provider_test.py @@ -37,7 +37,7 @@ ) from agent_env.providers.sandbox_providers.chained_sandbox_provider import ChainedSandboxProvider from agent_env.providers.sandbox_providers.modal_sandbox import ModalSandboxProvider -from agent_env.providers.sandbox_providers.sandbox import Sandbox +from agent_env.providers.sandbox_providers.sandbox import ContainerLimits, Sandbox from agent_env.providers.env_state import LocalPostgresStateProvider from agent_env.config import get_config, set_document_store from agent_env.store import EcrImageStore, NotFoundError @@ -735,6 +735,35 @@ def test_compose_maps_extra_hosts_into_the_gateway(extra_hosts): assert mapped == ({GATEWAY_SERVICE_NAME: list(extra_hosts)} if extra_hosts else {}) +def test_compose_holds_every_service_to_the_sandboxs_container_limits(): + """A local sandbox has no VM bounding its stack, so every service it renders, the store's and websites' + included, carries the sandbox's cpu and memory; a VM's sandbox (no limits) renders exactly as before.""" + from agent_env.env.envs.service_db import ServiceDBConfig + + def render(**limits): + return EnvironmentGatewayProvider().create_docker_compose( + mcp_servers=[MCPServerConfig(image="mcp-slack", environment_name="slack"), + MCPServerConfig(image="mcp-crm", environment_name="crm")], + gateway_image="agent-gateway", + website_configs=[WebsiteConfig(backend_image="shop-be", frontend_image="shop-fe", environment_name="shop")], + sidecars=[_sidecar_with_secret()], + state_provider=LocalPostgresStateProvider(service_db_config=ServiceDBConfig()), + state_instance=LocalPostgresStateProvider.default_instance(), + **limits, + ) + + unlimited = render() + limited = yaml.safe_load(render(container_limits=ContainerLimits(cpus=2, memory_mib=4096))) + + assert render(container_limits=None) == unlimited + assert not any("cpus" in service or "mem_limit" in service for service in yaml.safe_load(unlimited)["services"].values()) + assert set(limited["services"]) == set(yaml.safe_load(unlimited)["services"]) + assert {name: (s["cpus"], s["mem_limit"], s["memswap_limit"]) for name, s in limited["services"].items()} == { + name: (2, "4096m", "4096m") for name in limited["services"] + } + assert limited["networks"] == yaml.safe_load(unlimited)["networks"] + + # --- sidecar rendering -------------------------------------------------------- class _FakeImageArtifact: diff --git a/tst/unit/providers/sandbox_providers/local_sandbox_container_limits_test.py b/tst/unit/providers/sandbox_providers/local_sandbox_container_limits_test.py new file mode 100644 index 00000000..e31d7a7d --- /dev/null +++ b/tst/unit/providers/sandbox_providers/local_sandbox_container_limits_test.py @@ -0,0 +1,220 @@ +"""A local sandbox holds the containers it starts to the cpu and memory it was created with.""" + +from __future__ import annotations + +import logging +from types import SimpleNamespace + +import pytest + +import agent_env.a2a_agent.a2a_agent as a2a_module +import agent_env.config as cfg +import agent_env.providers.sandbox_providers.local_sandbox as ls +from agent_env.a2a_agent.a2a_agent import A2AAgent +from agent_env.providers.sandbox_providers.local_sandbox import LocalSandbox, LocalSandboxProvider +from agent_env.providers.sandbox_providers.sandbox import ContainerLimits + +_real_engine_cpus = ls._engine_cpus # before the unit conftest answers it + + +@pytest.fixture(autouse=True) +def sandbox_root(monkeypatch, tmp_path): + monkeypatch.setenv("AGENT_ENV_LOCAL_SANDBOX_DIR", str(tmp_path)) + monkeypatch.setattr(ls, "_host_ips", lambda: ("127.0.0.1",)) + return tmp_path + + +@pytest.fixture +def no_local_ca(monkeypatch): + """Containers start with ``docker run``, not created, given the local CA's files, then started.""" + monkeypatch.setattr(ls, "local_grant_trust", lambda: None) + monkeypatch.setattr(a2a_module, "local_grant_trust", lambda: None) + + +@pytest.fixture +def engine_answers(monkeypatch): + """``docker info`` answers each item in turn: a string is its output, None a failure.""" + asked = [] + + def answer(*outputs): + replies = iter(outputs) + + def run(args, **kwargs): + asked.append(args) + out = next(replies) + return SimpleNamespace(returncode=1 if out is None else 0, stdout=out or "", stderr="daemon down") + + monkeypatch.setattr(ls.subprocess, "run", run) + return asked + + return answer + + +class _RecordingLocalSandbox(LocalSandbox): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.scripts: list[str] = [] + + async def exec_script(self, script, *, max_retries=0): + self.scripts.append(script) + return "" + + +def _agent() -> A2AAgent: + return A2AAgent(id="a", version=None, docker_image_artifact=SimpleNamespace(image_name="img:1")) + + +def test_limits_render_as_docker_flags_and_compose_keys(): + limits = ContainerLimits(cpus=1.0, memory_mib=8192) + assert limits.docker_args == ("--cpus", "1", "--memory", "8192m", "--memory-swap", "8192m") + assert limits.compose_keys == ("cpus: 1", "mem_limit: 8192m", "memswap_limit: 8192m") + assert ContainerLimits(cpus=0.5, memory_mib=512).docker_args[:2] == ("--cpus", "0.5") + + +@pytest.mark.parametrize("cpus, memory_mib", [(0, 1024), (1, 0), (-1, 1024)]) +def test_limits_must_be_positive(cpus, memory_mib): + with pytest.raises(ValueError, match="must be positive"): + ContainerLimits(cpus=cpus, memory_mib=memory_mib) + + +@pytest.mark.asyncio +async def test_a_local_vm_keeps_its_limits_for_a_handle_rebuilt_later(): + sandbox = await LocalSandboxProvider().create_vm(cpu=2, memory=4096) + + assert sandbox.container_limits == ContainerLimits(cpus=2, memory_mib=4096) + reopened = await LocalSandboxProvider().get_sandbox(sandbox.sandbox_id) + assert reopened.container_limits == ContainerLimits(cpus=2, memory_mib=4096) + + +@pytest.mark.asyncio +async def test_cpu_beyond_the_engines_is_held_to_it(monkeypatch, caplog): + """Docker refuses a container more CPUs than its engine has, so a bigger ask is held to the engine's.""" + monkeypatch.setattr(ls, "_engine_cpus", lambda: 4) + + with caplog.at_level(logging.WARNING, logger=ls.logger.name): + sandbox = await LocalSandboxProvider().create_vm(cpu=8, memory=8192) + + assert sandbox.container_limits == ContainerLimits(cpus=4, memory_mib=8192) + assert "Asked for 8 CPUs, but the Docker engine has 4" in caplog.text + + +@pytest.mark.asyncio +async def test_cpu_within_the_engines_is_kept(monkeypatch): + monkeypatch.setattr(ls, "_engine_cpus", lambda: 4) + sandbox = await LocalSandboxProvider().create_vm(cpu=1.5, memory=2048) + assert sandbox.container_limits == ContainerLimits(cpus=1.5, memory_mib=2048) + + +def test_the_engine_is_asked_each_time_so_a_resized_one_is_seen(engine_answers): + asked = engine_answers("6\n", None, "2\n") + + assert [_real_engine_cpus() for _ in range(3)] == [6, None, 2] + assert asked == [["docker", "info", "--format", "{{.NCPU}}"]] * 3 + + +@pytest.mark.asyncio +async def test_one_cpu_or_less_is_not_checked_against_the_engine(monkeypatch): + def asked(): + raise AssertionError("every engine has at least one CPU") + + monkeypatch.setattr(ls, "_engine_cpus", asked) + sandbox = await LocalSandboxProvider().create_vm(cpu=1, memory=1024) + assert sandbox.container_limits == ContainerLimits(cpus=1, memory_mib=1024) + + +def test_no_docker_means_no_cpu_count(monkeypatch, engine_answers): + def missing(args, **kwargs): + raise FileNotFoundError("docker") + + monkeypatch.setattr(ls.subprocess, "run", missing) + assert _real_engine_cpus() is None + + +def test_a_sandbox_whose_limits_cant_be_recorded_is_not_made(sandbox_root): + """A handle a later step rebuilds would start its containers unheld.""" + with pytest.raises(FileNotFoundError): + LocalSandbox(sandbox_id="local-gone", work_dir=sandbox_root / "missing", container_limits=ContainerLimits(1, 512)) + + +def test_a_new_sandbox_whose_limits_cant_be_recorded_leaves_no_folder(monkeypatch, sandbox_root): + def full(self, *args, **kwargs): + raise OSError(28, "No space left on device") + + monkeypatch.setattr(ls.Path, "write_text", full) + with pytest.raises(OSError, match="No space left"): + LocalSandbox(container_limits=ContainerLimits(1, 512)) + assert list(sandbox_root.iterdir()) == [] + + +def test_a_sandbox_from_before_limits_were_recorded_holds_nothing(sandbox_root): + work_dir = sandbox_root / "agent-env-local-old-x" + work_dir.mkdir() + assert LocalSandbox(sandbox_id="local-old", work_dir=work_dir).container_limits is None + + +def test_an_unreadable_limits_file_is_ignored(sandbox_root, caplog): + work_dir = sandbox_root / "agent-env-local-bad-x" + work_dir.mkdir() + (work_dir / ls._LIMITS_FILE).write_text("{not json") + + with caplog.at_level(logging.WARNING, logger=ls.logger.name): + assert LocalSandbox(sandbox_id="local-bad", work_dir=work_dir).container_limits is None + assert "Ignoring unreadable container limits" in caplog.text + + +@pytest.mark.asyncio +async def test_a_container_mode_sandbox_starts_its_container_held_to_the_limits(monkeypatch, no_local_ca): + """An agent's own sandbox is one container, given the cpu and memory the sandbox was asked for.""" + monkeypatch.setattr(cfg, "get_config", lambda: SimpleNamespace(get_image_store=lambda: SimpleNamespace(auth=lambda _: None))) + monkeypatch.setattr(ls, "_free_host_port", lambda: 41337) + made: list[_RecordingLocalSandbox] = [] + + class _Provider(LocalSandboxProvider): + async def create_vm(self, **kwargs): + sandbox = await LocalSandboxProvider.create_vm(self, **kwargs) + reopened = _RecordingLocalSandbox(sandbox_id=sandbox.sandbox_id, work_dir=sandbox.work_dir) + made.append(reopened) + return reopened + + await _Provider().create_sandbox(image_name="img:v1", port=8000, env={}, cpu=1.5, memory=2048) + + (run,) = [s for s in made[0].scripts if s.startswith("docker run -d")] + assert "--cpus 1.5 --memory 2048m --memory-swap 2048m " in run + + +@pytest.mark.asyncio +async def test_an_agent_placed_on_a_local_vm_is_held_to_its_limits(no_local_ca): + """An agent placed on a sandbox shares the sandbox's budget, as it would a VM's.""" + agent = _agent() + agent._sandbox = _RecordingLocalSandbox(container_limits=ContainerLimits(cpus=2, memory_mib=4096)) + + await agent._run_container("img:1", 8000, {}) + + assert "--cpus 2 --memory 4096m --memory-swap 4096m " in agent._sandbox.scripts[0] + + +@pytest.mark.asyncio +async def test_the_docker_sidecar_is_held_to_the_same_limits(no_local_ca): + agent = _agent() + agent._sandbox = _RecordingLocalSandbox(container_limits=ContainerLimits(cpus=2, memory_mib=4096)) + + await agent._run_container("img:1", 8000, {}, enable_docker=True) + + script = agent._sandbox.scripts[0] + assert "DIND_LIMITS='--cpus 2 --memory 4096m --memory-swap 4096m'\n" in script + assert '--name "$DIND_CONTAINER" \\\n $DIND_LIMITS \\\n' in script + + +def test_a_sandbox_without_limits_gives_the_docker_sidecar_none(): + assert "DIND_LIMITS=\n" in _agent()._dind_setup_script() + + +@pytest.mark.asyncio +async def test_a_sandbox_without_limits_starts_containers_unlimited(no_local_ca): + """A VM's own size bounds its containers, so a sandbox that records no limits passes no flags.""" + agent = _agent() + agent._sandbox = _RecordingLocalSandbox() + + await agent._run_container("img:1", 8000, {}) + + assert "--cpus" not in agent._sandbox.scripts[0] and "--memory" not in agent._sandbox.scripts[0] diff --git a/tst/unit/task_step/run_docker_container_test.py b/tst/unit/task_step/run_docker_container_test.py index 4b1fa420..e4775e48 100644 --- a/tst/unit/task_step/run_docker_container_test.py +++ b/tst/unit/task_step/run_docker_container_test.py @@ -12,6 +12,7 @@ from agent_env.config import set_object_store from agent_env.providers.sandbox_providers import sandbox_provider from agent_env.providers.sandbox_providers.local_sandbox import LocalSandbox +from agent_env.providers.sandbox_providers.sandbox import ContainerLimits from agent_env.task_step.context import DeployedSandbox, TaskStepContext from agent_env.task_step.task_steps.run_docker_container import ( RunDockerContainerTaskStep as Step, @@ -216,9 +217,10 @@ async def test_a_store_context_that_is_not_a_zip_is_refused(fake_store): class _RunSandbox: sandbox_id = "local-1" - def __init__(self, extra_hosts: tuple[str, ...] = ()): + def __init__(self, extra_hosts: tuple[str, ...] = (), container_limits: ContainerLimits | None = None): self.scripts: list[str] = [] self.extra_hosts = extra_hosts + self.container_limits = container_limits async def exec_script(self, script): self.scripts.append(script) @@ -259,3 +261,24 @@ async def test_what_it_starts_maps_the_sandboxs_extra_hosts(monkeypatch, extra_h (run,) = [s for s in sandbox.scripts if "docker run -d" in s] assert ("--add-host host.docker.internal:host-gateway" in run) is bool(extra_hosts) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("limits", [None, ContainerLimits(cpus=2, memory_mib=4096)]) +async def test_what_it_starts_is_held_to_the_sandboxs_container_limits(monkeypatch, limits): + """A local sandbox has no VM bounding it, so each container it starts carries the sandbox's cpu and memory.""" + sandbox = _RunSandbox(container_limits=limits) + monkeypatch.setattr(sandbox_provider, "build_sandbox_provider", + lambda _type: SimpleNamespace(get_sandbox=AsyncMock(return_value=sandbox))) + monkeypatch.setattr(Step, "_stage_from_universe", AsyncMock()) + context = TaskStepContext(deployed_sandboxes=[ + DeployedSandbox(sandbox_name="h", sandbox_id="local-1", sandbox_mode="vm", sandbox_type="local"), + ]) + + await _step().execute(context) + + (run,) = [s for s in sandbox.scripts if "docker run -d" in s] + if limits is None: + assert "--cpus" not in run and "--memory" not in run + else: + assert "--cpus 2 --memory 4096m --memory-swap 4096m" in run