From 9c775b47b58d5688de1353ca314316eca545d195 Mon Sep 17 00:00:00 2001 From: t Date: Wed, 4 Nov 2026 15:42:19 -0500 Subject: [PATCH 1/7] feat(deploy_sandbox): keep an unsized step's cpu and memory unset A deploy_sandbox step that sets no cpu or memory still reserves 1 vCPU and 8 GiB, but now keeps both as None on the step and in its stored dict instead of writing 1.0 and 8192. A size the step chose can then be told apart from one it left to the default, which a default sizer needs. Steps stored with explicit values keep them; nothing else reads these fields. FD-5242 Co-Authored-By: Claude Opus 5.5 --- .../task_step/task_steps/deploy_sandbox.py | 24 +++++++++------ .../deploy_sandbox_default_cpu_test.py | 30 ++++++++++++------- 2 files changed, 34 insertions(+), 20 deletions(-) diff --git a/src/agent_env/task_step/task_steps/deploy_sandbox.py b/src/agent_env/task_step/task_steps/deploy_sandbox.py index 6e8bce2f4..a0530fc55 100644 --- a/src/agent_env/task_step/task_steps/deploy_sandbox.py +++ b/src/agent_env/task_step/task_steps/deploy_sandbox.py @@ -25,6 +25,10 @@ class DeploySandboxTaskStep(TaskStep): type: ClassVar[str] = "deploy_sandbox" entity_refs = () field_scopes = {**TaskStep.field_scopes, "sandbox_type": FieldScope.USER} + # What an unsized step reserves. Unset stays None on the step and in its stored dict, so a size the + # step chose is told apart from one it left to this default. + DEFAULT_CPU: ClassVar[float] = 1.0 + DEFAULT_MEMORY_MB: ClassVar[int] = 8192 def __init__( self, @@ -33,8 +37,8 @@ def __init__( sandbox_name: str, sandbox_mode: str, image: Optional[str] = None, - cpu: float = 1.0, - memory_mb: int = 8192, + cpu: Optional[float] = None, + memory_mb: Optional[int] = None, disk_size_gb: float = 10, ttl_seconds: int = TaskStep.DEFAULT_TTL_SECONDS, exposed_ports: Optional[list[int]] = None, @@ -90,8 +94,8 @@ def from_dict(cls, data: dict) -> DeploySandboxTaskStep: sandbox_name=data["sandbox_name"], sandbox_mode=data["sandbox_mode"], image=data.get("image"), - cpu=data.get("cpu", 1.0), - memory_mb=data.get("memory_mb", 8192), + cpu=data.get("cpu"), + memory_mb=data.get("memory_mb"), disk_size_gb=data.get("disk_size_gb", 10), ttl_seconds=data.get("ttl_seconds", TaskStep.DEFAULT_TTL_SECONDS), exposed_ports=data.get("exposed_ports"), @@ -127,9 +131,11 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: _ttl_override = user_overrides.get("ttl_seconds") resolved_ttl = _ttl_override if _ttl_override is not None else self.ttl_seconds attribution = deploy_attribution(self, context) + cpu = self.cpu if self.cpu is not None else self.DEFAULT_CPU + memory_mb = self.memory_mb if self.memory_mb is not None else self.DEFAULT_MEMORY_MB logger.info( f"Deploying sandbox '{self.sandbox_name}' (mode={self.sandbox_mode}, " - f"image={self.image}, cpu={self.cpu}, memory_mb={self.memory_mb}, " + f"image={self.image}, cpu={cpu}, memory_mb={memory_mb}, " f"disk_size_gb={self.disk_size_gb}, ttl_seconds={resolved_ttl}, " f"sandbox_type={resolved_sandbox_type}, attribution={attribution})" ) @@ -142,8 +148,8 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: sandbox = await provider.create_vm( image=self.image, boot_mode=self.boot_mode, - cpu=self.cpu, - memory=self.memory_mb, + cpu=cpu, + memory=memory_mb, disk_size_gb=self.disk_size_gb, timeout=resolved_ttl, exposed_ports=self.exposed_ports or [], @@ -158,8 +164,8 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: image_name=self.image, port=self.port, env=self.env_vars or {}, - cpu=self.cpu, - memory=self.memory_mb, + cpu=cpu, + memory=memory_mb, disk_size_gb=self.disk_size_gb, timeout=resolved_ttl, attribution=attribution, diff --git a/tst/unit/task_step/deploy_sandbox_default_cpu_test.py b/tst/unit/task_step/deploy_sandbox_default_cpu_test.py index cffa512dd..5f4493a41 100644 --- a/tst/unit/task_step/deploy_sandbox_default_cpu_test.py +++ b/tst/unit/task_step/deploy_sandbox_default_cpu_test.py @@ -1,5 +1,5 @@ -"""A deploy_sandbox step that sets no cpu reserves one vCPU, whether built in code or loaded from a stored dict, and -a stored cpu is kept as written.""" +"""A deploy_sandbox step that sets no cpu or memory reserves one vCPU and 8 GiB, whether built in code or loaded from a +stored dict, while keeping them unset on the step and in its stored dict; a stored size is kept as written.""" from unittest.mock import AsyncMock, patch @@ -29,12 +29,16 @@ def _stored(**fields) -> dict: DeploySandboxTaskStep(id="s", version=1, sandbox_name="box", sandbox_mode="vm"), DeploySandboxTaskStep.from_dict(_stored()), ], ids=["constructed", "from-dict"]) -def test_an_unsized_step_reserves_one_cpu(step): - assert step.cpu == 1.0 +def test_an_unsized_step_stays_unsized(step): + assert (step.cpu, step.memory_mb) == (None, None) + stored = step.to_dict() + assert (stored["cpu"], stored["memory_mb"]) == (None, None) + assert DeploySandboxTaskStep.from_dict(stored).cpu is None -def test_a_stored_cpu_is_kept(): - assert DeploySandboxTaskStep.from_dict(_stored(cpu=2.0)).cpu == 2.0 +def test_a_stored_size_is_kept(): + step = DeploySandboxTaskStep.from_dict(_stored(cpu=2.0, memory_mb=8192)) + assert (step.cpu, step.memory_mb) == (2.0, 8192) @pytest.mark.asyncio @@ -42,14 +46,18 @@ def test_a_stored_cpu_is_kept(): ("vm", "create_vm", {}), ("container", "create_container", {"image": "example.com/box:v1", "port": 80}), ], ids=["vm", "container"]) -@pytest.mark.parametrize("step_cpu, requested", [(None, 1.0), (4.0, 4.0)], ids=["unsized", "sized"]) -async def test_the_provider_is_asked_for_the_steps_cpu(mode, create, extra, step_cpu, requested): - sized = {} if step_cpu is None else {"cpu": step_cpu} +@pytest.mark.parametrize("size, requested", [ + ({}, (1.0, 8192)), + ({"cpu": 4.0, "memory_mb": 2048}, (4.0, 2048)), + ({"memory_mb": 2048}, (1.0, 2048)), +], ids=["unsized", "sized", "memory-only"]) +async def test_the_provider_is_asked_for_the_steps_size(mode, create, extra, size, requested): step = DeploySandboxTaskStep(id="s", version=1, sandbox_name="box", sandbox_mode=mode, sandbox_type="any", - **sized, **extra) + **size, **extra) provider = _provider() with patch("agent_env.providers.sandbox_providers.sandbox_provider.build_sandbox_provider", return_value=provider): await step.execute(TaskStepContext()) - assert getattr(provider, create).call_args.kwargs["cpu"] == requested + kwargs = getattr(provider, create).call_args.kwargs + assert (kwargs["cpu"], kwargs["memory"]) == requested From db9c421fadeb6d30337e9827ef6949762d49dffe Mon Sep 17 00:00:00 2001 From: t Date: Wed, 4 Nov 2026 17:28:54 -0500 Subject: [PATCH 2/7] feat(sandbox): let a configured sizer choose the sizes deploy steps leave unset [sandbox.sizer] names a SandboxSizer (impl/config, as [runner]). For each of cpu and memory, deploy_env, deploy_agent and deploy_sandbox reserve the step's own value, else the sizer's, else the step's or the provider's default. The in-tree default sizer chooses nothing, so without the table nothing changes. Each sandbox's attribution gains cpu_source and memory_source (step, sizer or default), so a deployment can see which sizes it set. FD-5241 Co-Authored-By: Claude Opus 5.5 --- src/agent_env/config/describe.py | 3 +- src/agent_env/config/runtime.py | 7 ++ .../providers/sandbox_providers/sizing.py | 64 ++++++++++ .../task_step/task_steps/deploy_agent.py | 11 +- .../task_step/task_steps/deploy_env.py | 9 +- .../task_step/task_steps/deploy_sandbox.py | 7 +- .../attribution_serialization_test.py | 10 +- tst/unit/task_step/sandbox_sizer_test.py | 115 ++++++++++++++++++ 8 files changed, 210 insertions(+), 16 deletions(-) create mode 100644 src/agent_env/providers/sandbox_providers/sizing.py create mode 100644 tst/unit/task_step/sandbox_sizer_test.py diff --git a/src/agent_env/config/describe.py b/src/agent_env/config/describe.py index e8e291d7f..713eb3b83 100644 --- a/src/agent_env/config/describe.py +++ b/src/agent_env/config/describe.py @@ -278,7 +278,8 @@ class FileSection: FileSection("conversations", ("conversations",), "config.runtime", _resolved_keys("conversations")), FileSection("agents", ("agents",), "config.runtime", _resolved_keys("agents"), validated_agents_section), FileSection("sandbox", ("sandbox",), "providers.sandbox_providers.sandbox_provider", - {"default": None, "agent_default": None, "attribution": None, "providers": Named(_SEAM)}), + {"default": None, "agent_default": None, "attribution": None, "providers": Named(_SEAM), + "sizer": _SEAM}), FileSection("state", ("state",), "providers.env_state.env_state_provider", {"providers": Named(_SEAM)}), FileSection("envs", ("envs",), "env.registry", {"impls": None, **_resolved_keys("envs")}), FileSection("artifacts", ("artifacts",), "artifact.registry", {"impls": None, "type_aliases": None}), diff --git a/src/agent_env/config/runtime.py b/src/agent_env/config/runtime.py index a98b17005..048bf41da 100644 --- a/src/agent_env/config/runtime.py +++ b/src/agent_env/config/runtime.py @@ -320,6 +320,13 @@ def sandbox_registry(self) -> dict[str, dict]: return self._registry("sandbox", sandbox_provider._build_registry) + def sandbox_sizer(self): + """The sizer ``[sandbox.sizer]`` names, which chooses sizes deploy steps leave unset; one that chooses none + when the table is absent.""" + from agent_env.providers.sandbox_providers import sizing + + return self._registry("sandbox_sizer", sizing._build_sizer) + def state_registry(self) -> dict[str, dict]: from agent_env.providers.env_state import env_state_provider diff --git a/src/agent_env/providers/sandbox_providers/sizing.py b/src/agent_env/providers/sandbox_providers/sizing.py new file mode 100644 index 000000000..04b7cd834 --- /dev/null +++ b/src/agent_env/providers/sandbox_providers/sizing.py @@ -0,0 +1,64 @@ +"""The cpu and memory a deploy step reserves for whatever size it leaves unset. + +A step's own value always wins. For each one it leaves unset, the sizer config.toml's ``[sandbox.sizer]`` names +(``impl`` and ``config``, as for ``[runner]``) may choose one; otherwise the step's default applies, or, with none, the +provider's. Where each value came from is added to the sandbox's attribution, so it reaches the provider with the rest. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Literal, Optional + +from agent_env.attribution import Attribution +from agent_env.config import loader, runtime + +if TYPE_CHECKING: + from agent_env.config.runtime import Config + +CPU_SOURCE_KEY = "cpu_source" +MEMORY_SOURCE_KEY = "memory_source" +Source = Literal["step", "sizer", "default"] + + +@dataclass(frozen=True) +class SandboxSize: + cpu: Optional[float] = None + memory_mb: Optional[int] = None + + +class SandboxSizer: + """Chooses a size for what a deploy step leaves unset. This base chooses nothing, so the step's or provider's + default applies; a subclass returns a value for each field it sizes and None for the rest.""" + + @classmethod + def from_config(cls, **config: Any) -> SandboxSizer: + return cls(**config) + + def size(self, step: Any, attribution: Attribution) -> SandboxSize: + return SandboxSize() + + +def _build_sizer(source: Config) -> SandboxSizer: + section = dict(source.section("sandbox", "sizer")) + if not section: + return SandboxSizer() + return loader.build_store(section, SandboxSizer, secret_resolver=source._resolve_secret) + + +def _pick(own: Any, chosen: Any, default: Any) -> tuple[Any, Source]: + if own is not None: + return own, "step" + if chosen is not None: + return chosen, "sizer" + return default, "default" + + +def resolve(step: Any, attribution: Attribution, default: SandboxSize = SandboxSize()) -> tuple[SandboxSize, Attribution]: + """The size ``step`` reserves, from its ``cpu`` and ``memory_mb``, and ``attribution`` with where each came from.""" + own = SandboxSize(step.cpu, step.memory_mb) + chosen = SandboxSize() if None not in (own.cpu, own.memory_mb) else runtime.get_config().sandbox_sizer().size( + step, attribution) + cpu, cpu_source = _pick(own.cpu, chosen.cpu, default.cpu) + memory_mb, memory_source = _pick(own.memory_mb, chosen.memory_mb, default.memory_mb) + return SandboxSize(cpu, memory_mb), {**attribution, CPU_SOURCE_KEY: cpu_source, MEMORY_SOURCE_KEY: memory_source} diff --git a/src/agent_env/task_step/task_steps/deploy_agent.py b/src/agent_env/task_step/task_steps/deploy_agent.py index cbacdc804..21fbc2e20 100644 --- a/src/agent_env/task_step/task_steps/deploy_agent.py +++ b/src/agent_env/task_step/task_steps/deploy_agent.py @@ -30,6 +30,7 @@ ) from agent_env.a2a_agent.staging import StagedObjectStore, transfer_store from agent_env.config import get_config +from agent_env.providers.sandbox_providers import sizing from agent_env.env.env import DeployedEnv, DeployedSandboxEnv, Env from agent_env.env.gateway import AGENT_ENV_ROLE_HEADER from agent_env.providers.sandbox_providers.local_docker.provider import host_url_for @@ -401,11 +402,11 @@ async def _execute(self, context: TaskStepContext, deployed_agents: list) -> Tas deploy_kwargs["sandbox_type"] = resolved_sandbox_type resolved_ttl = self._resolved_ttl_seconds(context) deploy_kwargs["ttl_seconds"] = resolved_ttl - if self.cpu is not None: - deploy_kwargs["cpu"] = self.cpu - if self.memory_mb is not None: - deploy_kwargs["memory"] = self.memory_mb - deploy_kwargs["attribution"] = deploy_attribution(self, context) + size, deploy_kwargs["attribution"] = sizing.resolve(self, deploy_attribution(self, context)) + if size.cpu is not None: + deploy_kwargs["cpu"] = size.cpu + if size.memory_mb is not None: + deploy_kwargs["memory"] = size.memory_mb _policy_override = user_overrides.get("network_policy") resolved_policy = _policy_override if _policy_override is not None else self.network_policy if resolved_policy is not None: diff --git a/src/agent_env/task_step/task_steps/deploy_env.py b/src/agent_env/task_step/task_steps/deploy_env.py index 375a4c515..db175fc5b 100644 --- a/src/agent_env/task_step/task_steps/deploy_env.py +++ b/src/agent_env/task_step/task_steps/deploy_env.py @@ -15,6 +15,7 @@ from agent_env.entity.scopes import FieldScope from agent_env.task_step.task_step import TaskStep, TaskStepDependency from agent_env.attribution import deploy_attribution +from agent_env.providers.sandbox_providers import sizing logger = logging.getLogger(__name__) @@ -148,11 +149,11 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: resolved_env_state_type, resolved_env_state_instance_id = self.run_state(user_overrides) _ttl_override = user_overrides.get("ttl_seconds") resolved_ttl = _ttl_override if _ttl_override is not None else self.ttl_seconds - attribution = deploy_attribution(self, context) + size, attribution = sizing.resolve(self, deploy_attribution(self, context)) logger.info( f"Deploying env '{self.env_id}' (ttl={resolved_ttl}s, " f"disk_size_gb={self.disk_size_gb}, gateway_mode={self.gateway_mode}, " - f"cpu={self.cpu}, memory_mb={self.memory_mb}, " + f"cpu={size.cpu}, memory_mb={size.memory_mb}, " f"sandbox_type={resolved_sandbox_type}, env_state_type={resolved_env_state_type}, " f"env_state_instance_id={resolved_env_state_instance_id}, " f"attribution={attribution})" @@ -161,8 +162,8 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: "ttl_seconds": resolved_ttl, "disk_size_gb": self.disk_size_gb, "gateway_mode": GatewayMode(self.gateway_mode), - "cpu": self.cpu, - "memory_mb": self.memory_mb, + "cpu": size.cpu, + "memory_mb": size.memory_mb, "attribution": attribution, } if resolved_sandbox_type is not None: diff --git a/src/agent_env/task_step/task_steps/deploy_sandbox.py b/src/agent_env/task_step/task_steps/deploy_sandbox.py index a0530fc55..97e544033 100644 --- a/src/agent_env/task_step/task_steps/deploy_sandbox.py +++ b/src/agent_env/task_step/task_steps/deploy_sandbox.py @@ -12,6 +12,7 @@ from agent_env.entity.scopes import FieldScope from agent_env.task_step.task_step import TaskStep, TaskStepDependency from agent_env.attribution import deploy_attribution +from agent_env.providers.sandbox_providers import sizing if TYPE_CHECKING: from agent_env.providers.sandbox_providers.sandbox_provider import SandboxProvider @@ -130,9 +131,9 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: _ttl_override = user_overrides.get("ttl_seconds") resolved_ttl = _ttl_override if _ttl_override is not None else self.ttl_seconds - attribution = deploy_attribution(self, context) - cpu = self.cpu if self.cpu is not None else self.DEFAULT_CPU - memory_mb = self.memory_mb if self.memory_mb is not None else self.DEFAULT_MEMORY_MB + size, attribution = sizing.resolve(self, deploy_attribution(self, context), + sizing.SandboxSize(self.DEFAULT_CPU, self.DEFAULT_MEMORY_MB)) + cpu, memory_mb = size.cpu, size.memory_mb logger.info( f"Deploying sandbox '{self.sandbox_name}' (mode={self.sandbox_mode}, " f"image={self.image}, cpu={cpu}, memory_mb={memory_mb}, " diff --git a/tst/unit/task_step/attribution_serialization_test.py b/tst/unit/task_step/attribution_serialization_test.py index 57db30565..823487115 100644 --- a/tst/unit/task_step/attribution_serialization_test.py +++ b/tst/unit/task_step/attribution_serialization_test.py @@ -11,6 +11,7 @@ import pytest from agent_env.attribution import attribution_of +from agent_env.providers.sandbox_providers.sizing import CPU_SOURCE_KEY, MEMORY_SOURCE_KEY from agent_env.env.env import DeployedGatewayEnv from agent_env.task_step.context import TaskStepContext from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep @@ -78,6 +79,9 @@ def test_a_top_level_key_in_a_document_is_not_attribution(factory): # --- the dict is what reaches the env -------------------------------------------- +# What deploy_env adds to a step's attribution: where its unset cpu and memory came from. +UNSIZED = {CPU_SOURCE_KEY: "default", MEMORY_SOURCE_KEY: "default"} + @pytest.mark.asyncio async def test_deploying_an_env_passes_the_attribution_dict(): @@ -98,11 +102,11 @@ async def deploy(*, attribution=None, **kwargs): Env.get.return_value = env await step.execute(TaskStepContext()) - assert passed["attribution"] == ATTRIBUTION + assert passed["attribution"] == {**ATTRIBUTION, **UNSIZED} @pytest.mark.asyncio -async def test_deploying_an_env_with_no_attribution_passes_an_empty_dict(): +async def test_deploying_an_env_with_no_attribution_passes_only_its_size_sources(): passed = {} async def deploy(*, attribution=None, **kwargs): @@ -119,4 +123,4 @@ async def deploy(*, attribution=None, **kwargs): Env.get.return_value = env await _deploy_env().execute(TaskStepContext()) - assert passed["attribution"] == {} + assert passed["attribution"] == UNSIZED diff --git a/tst/unit/task_step/sandbox_sizer_test.py b/tst/unit/task_step/sandbox_sizer_test.py new file mode 100644 index 000000000..1476d07ff --- /dev/null +++ b/tst/unit/task_step/sandbox_sizer_test.py @@ -0,0 +1,115 @@ +"""A deploy step reserves the size it sets, else what the configured sizer chooses for what it leaves unset, else its +default or the provider's, and says on the sandbox's attribution where each value came from.""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from agent_env.config import ConfigError +from agent_env.env.env import DeployedGatewayEnv +from agent_env.providers.sandbox_providers import sizing +from agent_env.providers.sandbox_providers.sandbox import NetworkPolicy +from agent_env.providers.sandbox_providers.sizing import SandboxSize, SandboxSizer +from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.task_steps.deploy_env import DeployEnvTaskStep +from agent_env.task_step.task_steps.deploy_sandbox import DeploySandboxTaskStep +from tst.util.config import config_with_document + + +class HalfSizer(SandboxSizer): + """Sizes every step's memory to ``memory_mb``, and its cpu only when ``cpu`` is given; records what it was asked.""" + + asked: list = [] + + def __init__(self, memory_mb: int, cpu: float | None = None): + self.memory_mb, self.cpu = memory_mb, cpu + + def size(self, step, attribution): + HalfSizer.asked.append((step.id, dict(attribution))) + return SandboxSize(cpu=self.cpu, memory_mb=self.memory_mb) + + +def _config(sizer: dict | None = None): + return config_with_document({"sandbox": {"sizer": sizer}} if sizer else {}) + + +@pytest.fixture(autouse=True) +def _reset_asked(): + HalfSizer.asked = [] + + +class _Step: + id = "s" + + def __init__(self, cpu=None, memory_mb=None): + self.cpu, self.memory_mb = cpu, memory_mb + + +@pytest.mark.parametrize("sizer, step, default, expected", [ + (None, _Step(), SandboxSize(), (None, None, "default", "default")), + (None, _Step(), SandboxSize(1.0, 8192), (1.0, 8192, "default", "default")), + ({"impl": HalfSizer, "config": {"memory_mb": 2048}}, _Step(), SandboxSize(1.0, 8192), + (1.0, 2048, "default", "sizer")), + ({"impl": HalfSizer, "config": {"memory_mb": 2048, "cpu": 0.5}}, _Step(), SandboxSize(), + (0.5, 2048, "sizer", "sizer")), + ({"impl": HalfSizer, "config": {"memory_mb": 2048, "cpu": 0.5}}, _Step(memory_mb=16384), SandboxSize(), + (0.5, 16384, "sizer", "step")), + ({"impl": HalfSizer, "config": {"memory_mb": 2048, "cpu": 0.5}}, _Step(cpu=4.0, memory_mb=16384), SandboxSize(), + (4.0, 16384, "step", "step")), +], ids=["no-sizer", "no-sizer-step-default", "sizer-memory-only", "sizer-both", "step-memory-wins", "step-sets-both"]) +def test_each_value_comes_from_the_step_then_the_sizer_then_the_default(sizer, step, default, expected): + with patch("agent_env.config.runtime.get_config", return_value=_config(sizer)): + size, attribution = sizing.resolve(step, {"run_id": "r"}, default) + + assert (size.cpu, size.memory_mb, attribution["cpu_source"], attribution["memory_source"]) == expected + assert attribution["run_id"] == "r" + + +def test_a_step_that_sets_both_never_asks_the_sizer(): + with patch("agent_env.config.runtime.get_config", return_value=_config({"impl": HalfSizer, "config": {"memory_mb": 1}})): + sizing.resolve(_Step(cpu=1.0, memory_mb=1024), {}) + assert HalfSizer.asked == [] + + +def test_a_sizer_that_is_not_one_is_refused(): + with patch("agent_env.config.runtime.get_config", return_value=_config({"impl": dict})): + with pytest.raises(ConfigError, match="not a subclass of SandboxSizer"): + sizing.resolve(_Step(), {}) + + +@pytest.mark.asyncio +async def test_deploy_env_hands_the_sized_values_and_their_sources_to_the_env(): + env = MagicMock() + env.deploy = AsyncMock(return_value=DeployedGatewayEnv(env_id="e", env_version=1, gateway_url="http://gw", + mcp_url="http://mcp", db_web_url=None, sandbox_id="s1")) + step = DeployEnvTaskStep(id="deploy", version=None, env_id="e", env_version=1, cpu=2.0) + config = _config({"impl": HalfSizer, "config": {"memory_mb": 2048, "cpu": 0.5}}) + + with patch("agent_env.config.runtime.get_config", return_value=config), patch("agent_env.env.env.Env") as Env: + Env.get.return_value = env + await step.execute(TaskStepContext(instance_id="run-1")) + + kwargs = env.deploy.await_args.kwargs + assert (kwargs["cpu"], kwargs["memory_mb"]) == (2.0, 2048) + assert {k: kwargs["attribution"][k] for k in ("cpu_source", "memory_source", "run_id")} == { + "cpu_source": "step", "memory_source": "sizer", "run_id": "run-1"} + assert HalfSizer.asked == [("deploy", {"run_id": "run-1"})] + + +@pytest.mark.asyncio +async def test_deploy_sandbox_falls_back_to_its_own_default_when_the_sizer_chooses_nothing(): + sandbox = AsyncMock() + sandbox.sandbox_id, sandbox.mode, sandbox.type = "sb-1", "vm", "modal_vm" + sandbox.tunnel_urls, sandbox.vnc_url, sandbox.network_policy = {}, None, NetworkPolicy() + provider = AsyncMock(ON_THIS_MACHINE=False) + provider.links = (provider,) + provider.create_vm.return_value = sandbox + step = DeploySandboxTaskStep(id="s", version=1, sandbox_name="box", sandbox_mode="vm", sandbox_type="any") + + with patch("agent_env.config.runtime.get_config", return_value=_config({"impl": SandboxSizer})), \ + patch("agent_env.providers.sandbox_providers.sandbox_provider.build_sandbox_provider", return_value=provider): + await step.execute(TaskStepContext()) + + kwargs = provider.create_vm.call_args.kwargs + assert (kwargs["cpu"], kwargs["memory"]) == (1.0, 8192) + assert (kwargs["attribution"]["cpu_source"], kwargs["attribution"]["memory_source"]) == ("default", "default") From f4d5a02eb4b495cfc24c13a91558b2d77e83bcbd Mon Sep 17 00:00:00 2001 From: t Date: Wed, 4 Nov 2026 17:36:04 -0500 Subject: [PATCH 3/7] fix(sandbox): mark a gateway deploy's other containers as default-memory In container mode only the gateway gets the step's memory; its MCP servers, sidecars and service db are created with cpu alone, so they get the provider's memory. Their attribution said the step's memory source anyway. They now carry memory_source "default". Co-Authored-By: Claude Opus 5.5 --- .../providers/env_providers/env_gateway_provider.py | 3 ++- src/agent_env/providers/sandbox_providers/sizing.py | 8 ++++++++ tst/unit/task_step/sandbox_sizer_test.py | 11 +++++++++++ 3 files changed, 21 insertions(+), 1 deletion(-) 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 84ca8cb11..89274012b 100644 --- a/src/agent_env/providers/env_providers/env_gateway_provider.py +++ b/src/agent_env/providers/env_providers/env_gateway_provider.py @@ -56,6 +56,7 @@ _tool_names, ) from agent_env.providers.sandbox_providers.local_docker.provider import remove_local_work_dir +from agent_env.providers.sandbox_providers import sizing from agent_env.providers.sandbox_providers.sandbox import Sandbox, VmSandbox, port_bindings from agent_env.providers.sandbox_providers.sandbox_provider import ( SANDBOX_MODE_CONTAINER, @@ -968,7 +969,7 @@ async def _deploy_via_containers( deploy = _ContainerDeploy( sandbox_provider=sandbox_provider, cpu=cpu, disk_size_gb=disk_size_gb, ttl_seconds=ttl_seconds, - attribution=attribution, + attribution=sizing.with_default_memory(attribution), ) try: diff --git a/src/agent_env/providers/sandbox_providers/sizing.py b/src/agent_env/providers/sandbox_providers/sizing.py index 04b7cd834..7facea72c 100644 --- a/src/agent_env/providers/sandbox_providers/sizing.py +++ b/src/agent_env/providers/sandbox_providers/sizing.py @@ -54,6 +54,14 @@ def _pick(own: Any, chosen: Any, default: Any) -> tuple[Any, Source]: return default, "default" +def with_default_memory(attribution: Optional[Attribution]) -> Optional[Attribution]: + """``attribution`` for a sandbox created without the step's memory, which gets its provider's: its memory source + is ``default`` whatever the step's was.""" + if not attribution or MEMORY_SOURCE_KEY not in attribution: + return attribution + return {**attribution, MEMORY_SOURCE_KEY: "default"} + + def resolve(step: Any, attribution: Attribution, default: SandboxSize = SandboxSize()) -> tuple[SandboxSize, Attribution]: """The size ``step`` reserves, from its ``cpu`` and ``memory_mb``, and ``attribution`` with where each came from.""" own = SandboxSize(step.cpu, step.memory_mb) diff --git a/tst/unit/task_step/sandbox_sizer_test.py b/tst/unit/task_step/sandbox_sizer_test.py index 1476d07ff..02dd968d8 100644 --- a/tst/unit/task_step/sandbox_sizer_test.py +++ b/tst/unit/task_step/sandbox_sizer_test.py @@ -113,3 +113,14 @@ async def test_deploy_sandbox_falls_back_to_its_own_default_when_the_sizer_choos kwargs = provider.create_vm.call_args.kwargs assert (kwargs["cpu"], kwargs["memory"]) == (1.0, 8192) assert (kwargs["attribution"]["cpu_source"], kwargs["attribution"]["memory_source"]) == ("default", "default") + + +@pytest.mark.parametrize("attribution, expected", [ + ({"run_id": "r", "cpu_source": "sizer", "memory_source": "sizer"}, + {"run_id": "r", "cpu_source": "sizer", "memory_source": "default"}), + ({"run_id": "r", "memory_source": "step"}, {"run_id": "r", "memory_source": "default"}), + ({"run_id": "r"}, {"run_id": "r"}), + (None, None), +], ids=["sized", "step-memory", "no-sources", "none"]) +def test_a_sandbox_made_without_the_steps_memory_says_its_memory_is_the_default(attribution, expected): + assert sizing.with_default_memory(attribution) == expected From 6361e1e1e662dbe450237f206fc9158397fd8eb9 Mon Sep 17 00:00:00 2001 From: t Date: Fri, 6 Nov 2026 15:35:43 -0500 Subject: [PATCH 4/7] fix(sandbox): tell the sizer when a step may be a VM, and never let it fail a deploy - The sizer is passed vm=True when a step may deploy onto a VM: a vm-mode deploy_sandbox, or a deploy_env/deploy_agent whose sandbox type names a backend that creates VMs (may_create_vm, read from the classes). - A sizer that raises or returns something other than a SandboxSize is logged and taken to choose nothing. - A stored deploy_sandbox step with the old default pair (1.0, 8192), which every unsized step was saved with, reads as unsized; to_dict leaves an unset size out rather than writing null. - A gateway's containers that don't take the step's memory get the source "provider", so their events aren't read as sized. - cpu_source/memory_source stay out of the names, labels and tags boat, daytona, sail_vm, coreweave and tenki build from attribution. - [sandbox] sizer = "pkg:Cls" works at runtime, as describe already accepted. FD-5241, FD-5242 Co-Authored-By: Claude Opus 5.5 --- src/agent_env/attribution.py | 13 ++ .../sandbox_providers/boat/config.py | 3 + .../sandbox_providers/coreweave/provider.py | 3 +- .../sandbox_providers/daytona/config.py | 3 + .../sandbox_providers/sail_vm/provider.py | 3 +- .../sandbox_providers/sandbox_provider.py | 17 +++ .../providers/sandbox_providers/sizing.py | 57 ++++++-- .../sandbox_providers/tenki/provider.py | 3 +- .../task_step/task_steps/deploy_agent.py | 4 +- .../task_step/task_steps/deploy_env.py | 4 +- .../task_step/task_steps/deploy_sandbox.py | 23 ++- .../deploy_sandbox_default_cpu_test.py | 16 ++- tst/unit/task_step/sandbox_sizer_test.py | 132 ++++++++++++++++-- 13 files changed, 242 insertions(+), 39 deletions(-) diff --git a/src/agent_env/attribution.py b/src/agent_env/attribution.py index b53be48ad..151ff9a18 100644 --- a/src/agent_env/attribution.py +++ b/src/agent_env/attribution.py @@ -19,6 +19,13 @@ PIPELINE_STEP_KEY = "pipeline_step" RUN_ID_KEY = "run_id" +# Where a sandbox's cpu and memory came from (``step``, ``sizer`` or ``default``; see +# ``providers.sandbox_providers.sizing``). They describe the sandbox rather than name what it is for, so a provider +# that builds names, labels or tags from attribution leaves them out (``identifying``). +CPU_SOURCE_KEY = "cpu_source" +MEMORY_SOURCE_KEY = "memory_source" +SIZE_SOURCE_KEYS = frozenset({CPU_SOURCE_KEY, MEMORY_SOURCE_KEY}) + def attribution_of(step) -> Attribution: """A step's attribution: a copy of ``metadata["attribution"]`` (``{}`` when unset).""" @@ -39,3 +46,9 @@ def deploy_attribution(step, context) -> Attribution: if instance_id and RUN_ID_KEY not in attribution: attribution[RUN_ID_KEY] = instance_id return attribution + + +def identifying(attribution) -> Attribution: + """``attribution`` without the keys that describe a sandbox's size rather than what it is for: what a provider + builds a sandbox's name, labels or tags from, whose length and number are limited.""" + return {key: value for key, value in (attribution or {}).items() if key not in SIZE_SOURCE_KEYS} diff --git a/src/agent_env/providers/sandbox_providers/boat/config.py b/src/agent_env/providers/sandbox_providers/boat/config.py index 8ebda8570..9dbe0a747 100644 --- a/src/agent_env/providers/sandbox_providers/boat/config.py +++ b/src/agent_env/providers/sandbox_providers/boat/config.py @@ -9,6 +9,7 @@ from dataclasses import dataclass, field from typing import Any +from agent_env.attribution import identifying from agent_env.config.errors import ConfigError logger = logging.getLogger(__name__) @@ -74,6 +75,7 @@ def _label_value(value: object) -> str: def sandbox_labels(attribution: Mapping[str, Any], ports: Iterable[int], create_token: str) -> dict[str, str]: """Labels for reconnect (ports), cleanup (the create token, the managed marker) and attribution (cost joins). boat keeps at most ``MAX_LABELS``; attribution past that is logged and left off.""" + attribution = identifying(attribution) labels = {MANAGED_LABEL: "true", PORTS_LABEL: ",".join(str(port) for port in ports), CREATE_LABEL: create_token} for key, value in attribution.items(): slug = _UNSAFE_KEY.sub("-", str(key)).strip("-./") @@ -89,6 +91,7 @@ def sandbox_labels(attribution: Mapping[str, Any], ports: Iterable[int], create_ def sandbox_name(attribution: Mapping[str, Any]) -> str: """``ae-`` plus the attribution values in key order: unique, findable by people.""" + attribution = identifying(attribution) slugs = [ _UNSAFE_NAME.sub("-", str(attribution[key])).strip("-") for key in sorted(attribution) diff --git a/src/agent_env/providers/sandbox_providers/coreweave/provider.py b/src/agent_env/providers/sandbox_providers/coreweave/provider.py index df13857a1..921256d72 100644 --- a/src/agent_env/providers/sandbox_providers/coreweave/provider.py +++ b/src/agent_env/providers/sandbox_providers/coreweave/provider.py @@ -32,7 +32,7 @@ ) from agent_env.artifact.artifacts.docker_image import DockerImageArtifact -from agent_env.attribution import Attribution +from agent_env.attribution import Attribution, identifying from agent_env.config import get_config from agent_env.config.errors import ConfigError from agent_env.providers.sandbox_providers.coreweave.sandbox import ( @@ -108,6 +108,7 @@ def _tag(text: str) -> str: def sandbox_tags(attribution: Attribution, create_tag: str) -> list[str]: """``agentenv``, the create's unique tag, and one ``key.value`` tag per attribution entry.""" + attribution = identifying(attribution) tags = [SANDBOX_TAG, create_tag] tags += [tag for key in sorted(attribution) if (tag := _tag(f"{key}.{attribution[key]}"))] return list(dict.fromkeys(tags)) diff --git a/src/agent_env/providers/sandbox_providers/daytona/config.py b/src/agent_env/providers/sandbox_providers/daytona/config.py index b29b5e980..a1be3f9a7 100644 --- a/src/agent_env/providers/sandbox_providers/daytona/config.py +++ b/src/agent_env/providers/sandbox_providers/daytona/config.py @@ -11,6 +11,7 @@ from daytona import CreateSandboxFromImageParams, CreateSandboxFromSnapshotParams, Resources +from agent_env.attribution import identifying from agent_env.config.errors import ConfigError from agent_env.providers.sandbox_providers.daytona.image import DEFAULT_BASE_IMAGE, HOST_IMAGE_VERSION, host_image @@ -110,6 +111,7 @@ def sandbox_labels( ) -> dict[str, str]: """Labels for attribution (cost joins), reconnect (ports), sweeping (the managed marker) and, on a sandbox booted from the host image, the recipe version that built it.""" + attribution = identifying(attribution) labels = {MANAGED_LABEL: "true", PORTS_LABEL: ",".join(str(port) for port in ports)} if host_image: labels[HOST_IMAGE_LABEL] = HOST_IMAGE_VERSION @@ -125,6 +127,7 @@ def sandbox_labels( def sandbox_name(attribution: Mapping[str, Any]) -> str: """``ae-`` plus the attribution values in key order: unique, findable by people.""" + attribution = identifying(attribution) slugs = [ re.sub(r"[^A-Za-z0-9-]+", "-", str(attribution[key])).strip("-") for key in sorted(attribution) diff --git a/src/agent_env/providers/sandbox_providers/sail_vm/provider.py b/src/agent_env/providers/sandbox_providers/sail_vm/provider.py index 24dfd5e1d..9d957fa59 100644 --- a/src/agent_env/providers/sandbox_providers/sail_vm/provider.py +++ b/src/agent_env/providers/sandbox_providers/sail_vm/provider.py @@ -9,7 +9,7 @@ import uuid from typing import Any, Awaitable, Callable, ClassVar, Self -from agent_env.attribution import PIPELINE_STEP_KEY, RUN_ID_KEY, Attribution +from agent_env.attribution import PIPELINE_STEP_KEY, RUN_ID_KEY, Attribution, identifying from agent_env.config import get_config from agent_env.config.errors import ConfigError from agent_env.providers.sandbox_providers.sail_vm import _sdk @@ -65,6 +65,7 @@ def sailbox_shape(cpu: float, memory_mb: int, disk_size_gb: float, *, min_size: def sailbox_name(attribution: Attribution) -> str: """``ae-`` plus the attribution values in key order, slugged: for people and ``list(search=)``.""" + attribution = identifying(attribution) slugs = [re.sub(r"[^A-Za-z0-9]+", "-", str(attribution[key])).strip("-") for key in sorted(attribution)] return "-".join(["ae", uuid.uuid4().hex[:8], *filter(None, slugs)])[:_MAX_NAME_LENGTH].rstrip("-") diff --git a/src/agent_env/providers/sandbox_providers/sandbox_provider.py b/src/agent_env/providers/sandbox_providers/sandbox_provider.py index 628984426..10604a8a3 100644 --- a/src/agent_env/providers/sandbox_providers/sandbox_provider.py +++ b/src/agent_env/providers/sandbox_providers/sandbox_provider.py @@ -704,6 +704,23 @@ def _build_sandbox_provider_from_section(name: str, section: dict) -> SandboxPro return _install_sandbox_type_guard(provider, name) +def may_create_vm(spec: Optional[str] = None, *, agent: bool = False) -> bool: + """Whether a sandbox from ``spec`` (unset: ``[sandbox] default``, or ``agent_default`` with ``agent``) may be a VM: + a backend it names creates VMs, as a gateway on one deploys into a VM. Read from the backends' classes, so nothing + is built; True when a backend can't be resolved, as the safe answer for something that sizes it.""" + from agent_env.config import load_impl + + spec = spec or (_agent_sandbox_spec() if agent else _default_sandbox_spec()) + try: + registry = _get_sandbox_registry() + classes = [load_impl(registry[name]["impl"], SandboxProvider) + for name in (s.strip() for s in spec.split(",")) if name] + except Exception as e: + logger.debug(f"Could not tell whether sandbox spec {spec!r} creates VMs, so taking it that it may: {e}") + return True + return not classes or any(SANDBOX_MODE_VM in cls.SUPPORTED_SANDBOX_MODES for cls in classes) + + def build_sandbox_provider(spec: str) -> SandboxProvider: """Build a provider (or a comma-separated fallback chain) from backend names, resolved through the open registry (built-ins + config.toml [sandbox.providers]).""" diff --git a/src/agent_env/providers/sandbox_providers/sizing.py b/src/agent_env/providers/sandbox_providers/sizing.py index 7facea72c..bec7543da 100644 --- a/src/agent_env/providers/sandbox_providers/sizing.py +++ b/src/agent_env/providers/sandbox_providers/sizing.py @@ -1,24 +1,33 @@ """The cpu and memory a deploy step reserves for whatever size it leaves unset. A step's own value always wins. For each one it leaves unset, the sizer config.toml's ``[sandbox.sizer]`` names -(``impl`` and ``config``, as for ``[runner]``) may choose one; otherwise the step's default applies, or, with none, the -provider's. Where each value came from is added to the sandbox's attribution, so it reaches the provider with the rest. +(``impl`` and ``config``, as for ``[runner]``, or the ``impl`` alone as a string) may choose one; otherwise the step's +default applies, or, with none, the provider's. Where each value came from is added to the sandbox's attribution, so +it reaches the provider with the rest. + +The sizer is told whether the step may deploy onto a VM, whose memory is a hard limit on everything in it, so it can +decline to size one. A sizer that fails never fails the deploy: its error is logged and it is taken to choose nothing. """ from __future__ import annotations +import logging from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Literal, Optional -from agent_env.attribution import Attribution +from agent_env.attribution import CPU_SOURCE_KEY, MEMORY_SOURCE_KEY, Attribution from agent_env.config import loader, runtime if TYPE_CHECKING: from agent_env.config.runtime import Config -CPU_SOURCE_KEY = "cpu_source" -MEMORY_SOURCE_KEY = "memory_source" -Source = Literal["step", "sizer", "default"] +logger = logging.getLogger(__name__) + +__all__ = ["CPU_SOURCE_KEY", "MEMORY_SOURCE_KEY", "SandboxSize", "SandboxSizer", "resolve", "with_default_memory"] + +# ``provider``: the sandbox doesn't take the step's value at all (a gateway's containers share its cpu but not its +# memory), so it gets its provider's, and a size chosen for the step says nothing about it. +Source = Literal["step", "sizer", "default", "provider"] @dataclass(frozen=True) @@ -29,18 +38,22 @@ class SandboxSize: class SandboxSizer: """Chooses a size for what a deploy step leaves unset. This base chooses nothing, so the step's or provider's - default applies; a subclass returns a value for each field it sizes and None for the rest.""" + default applies; a subclass returns a value for each field it sizes and None for the rest. + + ``vm`` is True when the step may deploy onto a VM (its sandbox mode is ``vm``, or a backend its sandbox type names + creates VMs), where the memory chosen is the limit on everything the VM runs.""" @classmethod def from_config(cls, **config: Any) -> SandboxSizer: return cls(**config) - def size(self, step: Any, attribution: Attribution) -> SandboxSize: + def size(self, step: Any, attribution: Attribution, *, vm: bool) -> SandboxSize: return SandboxSize() def _build_sizer(source: Config) -> SandboxSizer: - section = dict(source.section("sandbox", "sizer")) + raw = source.section("sandbox").get("sizer") + section = {"impl": raw} if isinstance(raw, str) else dict(raw or {}) if not section: return SandboxSizer() return loader.build_store(section, SandboxSizer, secret_resolver=source._resolve_secret) @@ -56,17 +69,31 @@ def _pick(own: Any, chosen: Any, default: Any) -> tuple[Any, Source]: def with_default_memory(attribution: Optional[Attribution]) -> Optional[Attribution]: """``attribution`` for a sandbox created without the step's memory, which gets its provider's: its memory source - is ``default`` whatever the step's was.""" + is ``provider`` whatever the step's was.""" if not attribution or MEMORY_SOURCE_KEY not in attribution: return attribution - return {**attribution, MEMORY_SOURCE_KEY: "default"} + return {**attribution, MEMORY_SOURCE_KEY: "provider"} + + +def _chosen(step: Any, attribution: Attribution, vm: bool) -> SandboxSize: + sizer = runtime.get_config().sandbox_sizer() # a misconfigured sizer is a config error, raised as one + try: + chosen = sizer.size(step, dict(attribution), vm=vm) + if not isinstance(chosen, SandboxSize): + raise TypeError(f"returned {type(chosen).__name__}, not a SandboxSize") + return chosen + except Exception as e: + logger.warning(f"{type(sizer).__name__} could not size step {getattr(step, 'id', None)!r}, so it gets its " + f"default: {type(e).__name__}: {e}") + return SandboxSize() -def resolve(step: Any, attribution: Attribution, default: SandboxSize = SandboxSize()) -> tuple[SandboxSize, Attribution]: - """The size ``step`` reserves, from its ``cpu`` and ``memory_mb``, and ``attribution`` with where each came from.""" +def resolve(step: Any, attribution: Attribution, default: SandboxSize = SandboxSize(), *, + vm: bool) -> tuple[SandboxSize, Attribution]: + """The size ``step`` reserves, from its ``cpu`` and ``memory_mb``, and ``attribution`` with where each came from. + ``vm``: whether the step may deploy onto a VM, which the sizer is told.""" own = SandboxSize(step.cpu, step.memory_mb) - chosen = SandboxSize() if None not in (own.cpu, own.memory_mb) else runtime.get_config().sandbox_sizer().size( - step, attribution) + chosen = SandboxSize() if None not in (own.cpu, own.memory_mb) else _chosen(step, attribution, vm) cpu, cpu_source = _pick(own.cpu, chosen.cpu, default.cpu) memory_mb, memory_source = _pick(own.memory_mb, chosen.memory_mb, default.memory_mb) return SandboxSize(cpu, memory_mb), {**attribution, CPU_SOURCE_KEY: cpu_source, MEMORY_SOURCE_KEY: memory_source} diff --git a/src/agent_env/providers/sandbox_providers/tenki/provider.py b/src/agent_env/providers/sandbox_providers/tenki/provider.py index a52d5ce0a..42ea9b3d8 100644 --- a/src/agent_env/providers/sandbox_providers/tenki/provider.py +++ b/src/agent_env/providers/sandbox_providers/tenki/provider.py @@ -20,7 +20,7 @@ from tenki import AsyncClient, SessionNotFoundError -from agent_env.attribution import Attribution +from agent_env.attribution import Attribution, identifying from agent_env.config.errors import ConfigError from agent_env.providers.sandbox_providers.sandbox import NetworkPolicy, NetworkPolicyUnsupportedError from agent_env.providers.sandbox_providers.sandbox_provider import SandboxProvider, apply_default_attribution @@ -103,6 +103,7 @@ def egress_kwargs(policy: NetworkPolicy) -> dict[str, Any]: def session_name(attribution: Attribution) -> str: """``agentenv-`` and the attribution values, slugged into Tenki's 63-character name.""" + attribution = identifying(attribution) slugs = [re.sub(r"[^a-z0-9]+", "-", str(attribution[key]).lower()).strip("-") for key in sorted(attribution)] return "-".join(["agentenv", uuid.uuid4().hex[:8], *filter(None, slugs)])[:63].rstrip("-") diff --git a/src/agent_env/task_step/task_steps/deploy_agent.py b/src/agent_env/task_step/task_steps/deploy_agent.py index 21fbc2e20..30d2bb4ee 100644 --- a/src/agent_env/task_step/task_steps/deploy_agent.py +++ b/src/agent_env/task_step/task_steps/deploy_agent.py @@ -31,6 +31,7 @@ from agent_env.a2a_agent.staging import StagedObjectStore, transfer_store from agent_env.config import get_config from agent_env.providers.sandbox_providers import sizing +from agent_env.providers.sandbox_providers.sandbox_provider import may_create_vm from agent_env.env.env import DeployedEnv, DeployedSandboxEnv, Env from agent_env.env.gateway import AGENT_ENV_ROLE_HEADER from agent_env.providers.sandbox_providers.local_docker.provider import host_url_for @@ -402,7 +403,8 @@ async def _execute(self, context: TaskStepContext, deployed_agents: list) -> Tas deploy_kwargs["sandbox_type"] = resolved_sandbox_type resolved_ttl = self._resolved_ttl_seconds(context) deploy_kwargs["ttl_seconds"] = resolved_ttl - size, deploy_kwargs["attribution"] = sizing.resolve(self, deploy_attribution(self, context)) + size, deploy_kwargs["attribution"] = sizing.resolve( + self, deploy_attribution(self, context), vm=may_create_vm(resolved_sandbox_type, agent=True)) if size.cpu is not None: deploy_kwargs["cpu"] = size.cpu if size.memory_mb is not None: diff --git a/src/agent_env/task_step/task_steps/deploy_env.py b/src/agent_env/task_step/task_steps/deploy_env.py index db175fc5b..8bab03036 100644 --- a/src/agent_env/task_step/task_steps/deploy_env.py +++ b/src/agent_env/task_step/task_steps/deploy_env.py @@ -16,6 +16,7 @@ from agent_env.task_step.task_step import TaskStep, TaskStepDependency from agent_env.attribution import deploy_attribution from agent_env.providers.sandbox_providers import sizing +from agent_env.providers.sandbox_providers.sandbox_provider import may_create_vm logger = logging.getLogger(__name__) @@ -149,7 +150,8 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: resolved_env_state_type, resolved_env_state_instance_id = self.run_state(user_overrides) _ttl_override = user_overrides.get("ttl_seconds") resolved_ttl = _ttl_override if _ttl_override is not None else self.ttl_seconds - size, attribution = sizing.resolve(self, deploy_attribution(self, context)) + size, attribution = sizing.resolve(self, deploy_attribution(self, context), + vm=may_create_vm(resolved_sandbox_type)) logger.info( f"Deploying env '{self.env_id}' (ttl={resolved_ttl}s, " f"disk_size_gb={self.disk_size_gb}, gateway_mode={self.gateway_mode}, " diff --git a/src/agent_env/task_step/task_steps/deploy_sandbox.py b/src/agent_env/task_step/task_steps/deploy_sandbox.py index 97e544033..429f904dd 100644 --- a/src/agent_env/task_step/task_steps/deploy_sandbox.py +++ b/src/agent_env/task_step/task_steps/deploy_sandbox.py @@ -75,8 +75,11 @@ def to_dict(self) -> dict: base["sandbox_name"] = self.sandbox_name base["sandbox_mode"] = self.sandbox_mode base["image"] = self.image - base["cpu"] = self.cpu - base["memory_mb"] = self.memory_mb + # An unset size is left out, so a reader that predates it never sees a null. + if self.cpu is not None: + base["cpu"] = self.cpu + if self.memory_mb is not None: + base["memory_mb"] = self.memory_mb base["disk_size_gb"] = self.disk_size_gb base["ttl_seconds"] = self.ttl_seconds base["exposed_ports"] = self.exposed_ports @@ -88,6 +91,16 @@ def to_dict(self) -> dict: base["network_policy"] = self.network_policy return base + @classmethod + def _stored_size(cls, data: dict) -> dict: + """``cpu`` and ``memory_mb`` as ``data`` stores them. Before a step could leave them unset, every unsized step + was stored with both at the default, 1.0 and 8192, so that pair is read as unset: such a step gets the same + size unless a sizer chooses one.""" + cpu, memory_mb = data.get("cpu"), data.get("memory_mb") + if (cpu, memory_mb) == (cls.DEFAULT_CPU, cls.DEFAULT_MEMORY_MB): + return {"cpu": None, "memory_mb": None} + return {"cpu": cpu, "memory_mb": memory_mb} + @classmethod def from_dict(cls, data: dict) -> DeploySandboxTaskStep: return cls( @@ -95,8 +108,7 @@ def from_dict(cls, data: dict) -> DeploySandboxTaskStep: sandbox_name=data["sandbox_name"], sandbox_mode=data["sandbox_mode"], image=data.get("image"), - cpu=data.get("cpu"), - memory_mb=data.get("memory_mb"), + **cls._stored_size(data), disk_size_gb=data.get("disk_size_gb", 10), ttl_seconds=data.get("ttl_seconds", TaskStep.DEFAULT_TTL_SECONDS), exposed_ports=data.get("exposed_ports"), @@ -132,7 +144,8 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: _ttl_override = user_overrides.get("ttl_seconds") resolved_ttl = _ttl_override if _ttl_override is not None else self.ttl_seconds size, attribution = sizing.resolve(self, deploy_attribution(self, context), - sizing.SandboxSize(self.DEFAULT_CPU, self.DEFAULT_MEMORY_MB)) + sizing.SandboxSize(self.DEFAULT_CPU, self.DEFAULT_MEMORY_MB), + vm=self.sandbox_mode == "vm") cpu, memory_mb = size.cpu, size.memory_mb logger.info( f"Deploying sandbox '{self.sandbox_name}' (mode={self.sandbox_mode}, " diff --git a/tst/unit/task_step/deploy_sandbox_default_cpu_test.py b/tst/unit/task_step/deploy_sandbox_default_cpu_test.py index 5f4493a41..97531d0b9 100644 --- a/tst/unit/task_step/deploy_sandbox_default_cpu_test.py +++ b/tst/unit/task_step/deploy_sandbox_default_cpu_test.py @@ -32,13 +32,21 @@ def _stored(**fields) -> dict: def test_an_unsized_step_stays_unsized(step): assert (step.cpu, step.memory_mb) == (None, None) stored = step.to_dict() - assert (stored["cpu"], stored["memory_mb"]) == (None, None) + assert "cpu" not in stored and "memory_mb" not in stored # left out, so no older reader meets a null assert DeploySandboxTaskStep.from_dict(stored).cpu is None -def test_a_stored_size_is_kept(): - step = DeploySandboxTaskStep.from_dict(_stored(cpu=2.0, memory_mb=8192)) - assert (step.cpu, step.memory_mb) == (2.0, 8192) +@pytest.mark.parametrize("size", [{"cpu": 2.0, "memory_mb": 8192}, {"cpu": 1.0}, {"memory_mb": 8192}, + {"cpu": 1.0, "memory_mb": 4096}], ids=["both", "cpu", "memory", "cpu-at-default"]) +def test_a_stored_size_is_kept(size): + step = DeploySandboxTaskStep.from_dict(_stored(**size)) + assert (step.cpu, step.memory_mb) == (size.get("cpu"), size.get("memory_mb")) + + +def test_the_old_stored_default_reads_as_unsized(): + """Every unsized step was stored as 1.0 / 8192 before a step could leave its size unset.""" + step = DeploySandboxTaskStep.from_dict(_stored(cpu=1.0, memory_mb=8192)) + assert (step.cpu, step.memory_mb) == (None, None) @pytest.mark.asyncio diff --git a/tst/unit/task_step/sandbox_sizer_test.py b/tst/unit/task_step/sandbox_sizer_test.py index 02dd968d8..0b3256dca 100644 --- a/tst/unit/task_step/sandbox_sizer_test.py +++ b/tst/unit/task_step/sandbox_sizer_test.py @@ -24,8 +24,8 @@ class HalfSizer(SandboxSizer): def __init__(self, memory_mb: int, cpu: float | None = None): self.memory_mb, self.cpu = memory_mb, cpu - def size(self, step, attribution): - HalfSizer.asked.append((step.id, dict(attribution))) + def size(self, step, attribution, *, vm): + HalfSizer.asked.append((step.id, dict(attribution), vm)) return SandboxSize(cpu=self.cpu, memory_mb=self.memory_mb) @@ -59,7 +59,7 @@ def __init__(self, cpu=None, memory_mb=None): ], ids=["no-sizer", "no-sizer-step-default", "sizer-memory-only", "sizer-both", "step-memory-wins", "step-sets-both"]) def test_each_value_comes_from_the_step_then_the_sizer_then_the_default(sizer, step, default, expected): with patch("agent_env.config.runtime.get_config", return_value=_config(sizer)): - size, attribution = sizing.resolve(step, {"run_id": "r"}, default) + size, attribution = sizing.resolve(step, {"run_id": "r"}, default, vm=False) assert (size.cpu, size.memory_mb, attribution["cpu_source"], attribution["memory_source"]) == expected assert attribution["run_id"] == "r" @@ -67,14 +67,14 @@ def test_each_value_comes_from_the_step_then_the_sizer_then_the_default(sizer, s def test_a_step_that_sets_both_never_asks_the_sizer(): with patch("agent_env.config.runtime.get_config", return_value=_config({"impl": HalfSizer, "config": {"memory_mb": 1}})): - sizing.resolve(_Step(cpu=1.0, memory_mb=1024), {}) + sizing.resolve(_Step(cpu=1.0, memory_mb=1024), {}, vm=False) assert HalfSizer.asked == [] def test_a_sizer_that_is_not_one_is_refused(): with patch("agent_env.config.runtime.get_config", return_value=_config({"impl": dict})): with pytest.raises(ConfigError, match="not a subclass of SandboxSizer"): - sizing.resolve(_Step(), {}) + sizing.resolve(_Step(), {}, vm=False) @pytest.mark.asyncio @@ -85,7 +85,8 @@ async def test_deploy_env_hands_the_sized_values_and_their_sources_to_the_env(): step = DeployEnvTaskStep(id="deploy", version=None, env_id="e", env_version=1, cpu=2.0) config = _config({"impl": HalfSizer, "config": {"memory_mb": 2048, "cpu": 0.5}}) - with patch("agent_env.config.runtime.get_config", return_value=config), patch("agent_env.env.env.Env") as Env: + with patch("agent_env.config.runtime.get_config", return_value=config), patch("agent_env.env.env.Env") as Env, \ + patch("agent_env.task_step.task_steps.deploy_env.may_create_vm", return_value=False): Env.get.return_value = env await step.execute(TaskStepContext(instance_id="run-1")) @@ -93,7 +94,7 @@ async def test_deploy_env_hands_the_sized_values_and_their_sources_to_the_env(): assert (kwargs["cpu"], kwargs["memory_mb"]) == (2.0, 2048) assert {k: kwargs["attribution"][k] for k in ("cpu_source", "memory_source", "run_id")} == { "cpu_source": "step", "memory_source": "sizer", "run_id": "run-1"} - assert HalfSizer.asked == [("deploy", {"run_id": "run-1"})] + assert HalfSizer.asked == [("deploy", {"run_id": "run-1"}, False)] @pytest.mark.asyncio @@ -117,10 +118,121 @@ async def test_deploy_sandbox_falls_back_to_its_own_default_when_the_sizer_choos @pytest.mark.parametrize("attribution, expected", [ ({"run_id": "r", "cpu_source": "sizer", "memory_source": "sizer"}, - {"run_id": "r", "cpu_source": "sizer", "memory_source": "default"}), - ({"run_id": "r", "memory_source": "step"}, {"run_id": "r", "memory_source": "default"}), + {"run_id": "r", "cpu_source": "sizer", "memory_source": "provider"}), + ({"run_id": "r", "memory_source": "step"}, {"run_id": "r", "memory_source": "provider"}), ({"run_id": "r"}, {"run_id": "r"}), (None, None), ], ids=["sized", "step-memory", "no-sources", "none"]) -def test_a_sandbox_made_without_the_steps_memory_says_its_memory_is_the_default(attribution, expected): +def test_a_sandbox_made_without_the_steps_memory_says_its_provider_chose_it(attribution, expected): assert sizing.with_default_memory(attribution) == expected + + +class RaisingSizer(SandboxSizer): + def size(self, step, attribution, *, vm): + raise RuntimeError("store down") + + +class WrongShapeSizer(SandboxSizer): + def size(self, step, attribution, *, vm): + return {"cpu": 0.5} + + +@pytest.mark.parametrize("impl", [RaisingSizer, WrongShapeSizer], ids=["raises", "wrong-shape"]) +def test_a_sizer_that_fails_leaves_the_default_and_never_fails_the_deploy(impl, caplog): + with patch("agent_env.config.runtime.get_config", return_value=_config({"impl": impl})): + size, attribution = sizing.resolve(_Step(), {}, SandboxSize(1.0, 8192), vm=False) + + assert (size, attribution["cpu_source"], attribution["memory_source"]) == (SandboxSize(1.0, 8192), "default", + "default") + assert f"{impl.__name__} could not size step 's'" in caplog.text + + +def test_a_sizer_named_by_a_string_is_built(): + config = config_with_document({"sandbox": {"sizer": f"{__name__}:RaisingSizer"}}) + assert type(config.sandbox_sizer()) is RaisingSizer + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode, create, extra, vm", [ + ("vm", "create_vm", {}, True), + ("container", "create_container", {"image": "example.com/box:v1", "port": 80}, False), +], ids=["vm", "container"]) +async def test_deploy_sandbox_tells_the_sizer_whether_it_deploys_a_vm(mode, create, extra, vm): + sandbox = AsyncMock() + sandbox.sandbox_id, sandbox.mode, sandbox.type = "sb-1", mode, "modal" + sandbox.tunnel_urls, sandbox.vnc_url, sandbox.network_policy = {}, None, NetworkPolicy() + provider = AsyncMock(ON_THIS_MACHINE=False) + provider.links = (provider,) + getattr(provider, create).return_value = sandbox + step = DeploySandboxTaskStep(id="s", version=1, sandbox_name="box", sandbox_mode=mode, sandbox_type="any", **extra) + + with patch("agent_env.config.runtime.get_config", return_value=_config({"impl": HalfSizer, + "config": {"memory_mb": 2048}})), \ + patch("agent_env.providers.sandbox_providers.sandbox_provider.build_sandbox_provider", return_value=provider): + await step.execute(TaskStepContext()) + + assert HalfSizer.asked[0][2] is vm + + +@pytest.mark.parametrize("spec, vm", [ + ("modal", False), ("modal_vm", True), ("modal,modal_vm", True), ("no_such_backend", True), +], ids=["container", "vm", "chain-with-a-vm", "unknown"]) +def test_a_spec_may_create_a_vm_when_any_backend_it_names_does(spec, vm): + from agent_env.providers.sandbox_providers.sandbox_provider import may_create_vm + + assert may_create_vm(spec) is vm + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sandbox_type, vm", [("modal", False), ("modal_vm", True)], ids=["container", "vm"]) +async def test_deploy_env_tells_the_sizer_whether_its_sandbox_type_may_be_a_vm(sandbox_type, vm): + env = MagicMock() + env.deploy = AsyncMock(return_value=DeployedGatewayEnv(env_id="e", env_version=1, gateway_url="http://gw", + mcp_url="http://mcp", db_web_url=None, sandbox_id="s1")) + step = DeployEnvTaskStep(id="deploy", version=None, env_id="e", env_version=1, sandbox_type=sandbox_type) + + with patch("agent_env.config.runtime.get_config", return_value=_config({"impl": HalfSizer, + "config": {"memory_mb": 2048}})), \ + patch("agent_env.env.env.Env") as Env: + Env.get.return_value = env + await step.execute(TaskStepContext(instance_id="run-1")) + + assert HalfSizer.asked == [("deploy", {"run_id": "run-1"}, vm)] + + +class _Deployed(Exception): + """Stops a deploy_agent step once it has asked its agent to deploy, with what it asked.""" + + +@pytest.mark.asyncio +async def test_deploy_agent_hands_the_sized_values_and_their_sources_to_the_agent(): + from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep + + agent = MagicMock(metadata={}) + agent.deploy = AsyncMock(side_effect=_Deployed) + step = DeployAgentTaskStep(id="agent", version=1, env_ids=[], a2a_agent_id="solver", memory_mb=16384, + sandbox_type="modal") + config = _config({"impl": HalfSizer, "config": {"memory_mb": 2048, "cpu": 0.5}}) + + with patch("agent_env.config.runtime.get_config", return_value=config), \ + patch("agent_env.a2a_agent.A2AAgent") as A2AAgent, \ + patch.object(DeployAgentTaskStep, "_resolve_agent_ref", return_value=("solver", 1)), \ + patch.object(DeployAgentTaskStep, "_model_endpoint", return_value={}): + A2AAgent.get.return_value = agent + with pytest.raises(_Deployed): + await step.execute(TaskStepContext(instance_id="run-1")) + + kwargs = agent.deploy.await_args.kwargs + assert (kwargs["cpu"], kwargs["memory"]) == (0.5, 16384) + assert (kwargs["attribution"]["cpu_source"], kwargs["attribution"]["memory_source"]) == ("sizer", "step") + assert HalfSizer.asked == [("agent", {"run_id": "run-1"}, False)] + + +def test_where_a_size_came_from_stays_out_of_names_labels_and_tags(): + from agent_env.providers.sandbox_providers.boat import config as boat + from agent_env.providers.sandbox_providers.coreweave.provider import sandbox_tags + + attribution = {"project_id": "p", "cpu_source": "sizer", "memory_source": "default"} + assert "sizer" not in boat.sandbox_name(attribution) and "default" not in boat.sandbox_name(attribution) + assert not any("source" in key for key in boat.sandbox_labels(attribution, [80], "t")) + assert not any("source" in tag for tag in sandbox_tags(attribution, "c")) From fbd73bf83d914c924f0e84d9b8b1c972652360c9 Mon Sep 17 00:00:00 2001 From: t Date: Mon, 9 Nov 2026 13:17:31 -0500 Subject: [PATCH 5/7] fix(sandbox): treat a container on a VM backend as a VM, and keep a saved 1.0/8192 A container-mode deploy_sandbox on a VM backend (modal_vm) runs in a VM of its own whose memory is the size chosen, so the sizer is now told vm=True for it too. A step that sets cpu=1.0 and memory_mb=8192 is now saved with size_set, so a reload keeps the pair instead of reading it as a step stored before sizes could be unset. Co-Authored-By: Claude Opus 5.5 --- .../task_step/task_steps/deploy_sandbox.py | 16 +++++++++++----- .../task_step/deploy_sandbox_default_cpu_test.py | 6 ++++++ tst/unit/task_step/sandbox_sizer_test.py | 15 +++++++++------ 3 files changed, 26 insertions(+), 11 deletions(-) diff --git a/src/agent_env/task_step/task_steps/deploy_sandbox.py b/src/agent_env/task_step/task_steps/deploy_sandbox.py index 429f904dd..1c4c127fa 100644 --- a/src/agent_env/task_step/task_steps/deploy_sandbox.py +++ b/src/agent_env/task_step/task_steps/deploy_sandbox.py @@ -30,6 +30,7 @@ class DeploySandboxTaskStep(TaskStep): # step chose is told apart from one it left to this default. DEFAULT_CPU: ClassVar[float] = 1.0 DEFAULT_MEMORY_MB: ClassVar[int] = 8192 + SIZE_SET_KEY: ClassVar[str] = "size_set" def __init__( self, @@ -80,6 +81,8 @@ def to_dict(self) -> dict: base["cpu"] = self.cpu if self.memory_mb is not None: base["memory_mb"] = self.memory_mb + if (self.cpu, self.memory_mb) == (self.DEFAULT_CPU, self.DEFAULT_MEMORY_MB): + base[self.SIZE_SET_KEY] = True # so this pair isn't read back as a step stored before sizes could be unset base["disk_size_gb"] = self.disk_size_gb base["ttl_seconds"] = self.ttl_seconds base["exposed_ports"] = self.exposed_ports @@ -94,10 +97,11 @@ def to_dict(self) -> dict: @classmethod def _stored_size(cls, data: dict) -> dict: """``cpu`` and ``memory_mb`` as ``data`` stores them. Before a step could leave them unset, every unsized step - was stored with both at the default, 1.0 and 8192, so that pair is read as unset: such a step gets the same - size unless a sizer chooses one.""" + was stored with both at the default, 1.0 and 8192, so that pair is read as unset unless ``to_dict`` marked it as + set: a step stored back then gets the same size unless a sizer chooses one, and a step that sets the pair keeps + it.""" cpu, memory_mb = data.get("cpu"), data.get("memory_mb") - if (cpu, memory_mb) == (cls.DEFAULT_CPU, cls.DEFAULT_MEMORY_MB): + if (cpu, memory_mb) == (cls.DEFAULT_CPU, cls.DEFAULT_MEMORY_MB) and not data.get(cls.SIZE_SET_KEY): return {"cpu": None, "memory_mb": None} return {"cpu": cpu, "memory_mb": memory_mb} @@ -127,7 +131,8 @@ def _pull_refusal(self, link: SandboxProvider) -> str | None: return None async def execute(self, context: TaskStepContext) -> TaskStepContext: - from agent_env.providers.sandbox_providers.sandbox_provider import build_sandbox_provider, get_sandbox_provider + from agent_env.providers.sandbox_providers.sandbox_provider import ( + build_sandbox_provider, get_sandbox_provider, may_create_vm) if any(s.sandbox_name == self.sandbox_name for s in context.deployed_sandboxes): raise RuntimeError(f"Sandbox '{self.sandbox_name}' is already deployed") @@ -145,7 +150,8 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: resolved_ttl = _ttl_override if _ttl_override is not None else self.ttl_seconds size, attribution = sizing.resolve(self, deploy_attribution(self, context), sizing.SandboxSize(self.DEFAULT_CPU, self.DEFAULT_MEMORY_MB), - vm=self.sandbox_mode == "vm") + # A VM backend runs even a container in a VM of its own, sized by it. + vm=self.sandbox_mode == "vm" or may_create_vm(resolved_sandbox_type)) cpu, memory_mb = size.cpu, size.memory_mb logger.info( f"Deploying sandbox '{self.sandbox_name}' (mode={self.sandbox_mode}, " diff --git a/tst/unit/task_step/deploy_sandbox_default_cpu_test.py b/tst/unit/task_step/deploy_sandbox_default_cpu_test.py index 97531d0b9..455354306 100644 --- a/tst/unit/task_step/deploy_sandbox_default_cpu_test.py +++ b/tst/unit/task_step/deploy_sandbox_default_cpu_test.py @@ -43,6 +43,12 @@ def test_a_stored_size_is_kept(size): assert (step.cpu, step.memory_mb) == (size.get("cpu"), size.get("memory_mb")) +def test_a_step_that_sets_the_default_pair_keeps_it_through_a_save(): + step = DeploySandboxTaskStep(id="s", version=1, sandbox_name="box", sandbox_mode="vm", cpu=1.0, memory_mb=8192) + reloaded = DeploySandboxTaskStep.from_dict(step.to_dict()) + assert (reloaded.cpu, reloaded.memory_mb) == (1.0, 8192) + + def test_the_old_stored_default_reads_as_unsized(): """Every unsized step was stored as 1.0 / 8192 before a step could leave its size unset.""" step = DeploySandboxTaskStep.from_dict(_stored(cpu=1.0, memory_mb=8192)) diff --git a/tst/unit/task_step/sandbox_sizer_test.py b/tst/unit/task_step/sandbox_sizer_test.py index 0b3256dca..a837d3749 100644 --- a/tst/unit/task_step/sandbox_sizer_test.py +++ b/tst/unit/task_step/sandbox_sizer_test.py @@ -153,18 +153,21 @@ def test_a_sizer_named_by_a_string_is_built(): @pytest.mark.asyncio -@pytest.mark.parametrize("mode, create, extra, vm", [ - ("vm", "create_vm", {}, True), - ("container", "create_container", {"image": "example.com/box:v1", "port": 80}, False), -], ids=["vm", "container"]) -async def test_deploy_sandbox_tells_the_sizer_whether_it_deploys_a_vm(mode, create, extra, vm): +@pytest.mark.parametrize("mode, create, extra, sandbox_type, vm", [ + ("vm", "create_vm", {}, "modal", True), + ("container", "create_container", {"image": "example.com/box:v1", "port": 80}, "modal", False), + # A VM backend runs the container in a VM of its own, whose memory is the size chosen. + ("container", "create_container", {"image": "example.com/box:v1", "port": 80}, "modal_vm", True), +], ids=["vm", "container", "container-on-a-vm-backend"]) +async def test_deploy_sandbox_tells_the_sizer_whether_it_deploys_a_vm(mode, create, extra, sandbox_type, vm): sandbox = AsyncMock() sandbox.sandbox_id, sandbox.mode, sandbox.type = "sb-1", mode, "modal" sandbox.tunnel_urls, sandbox.vnc_url, sandbox.network_policy = {}, None, NetworkPolicy() provider = AsyncMock(ON_THIS_MACHINE=False) provider.links = (provider,) getattr(provider, create).return_value = sandbox - step = DeploySandboxTaskStep(id="s", version=1, sandbox_name="box", sandbox_mode=mode, sandbox_type="any", **extra) + step = DeploySandboxTaskStep(id="s", version=1, sandbox_name="box", sandbox_mode=mode, sandbox_type=sandbox_type, + **extra) with patch("agent_env.config.runtime.get_config", return_value=_config({"impl": HalfSizer, "config": {"memory_mb": 2048}})), \ From ed3cb7fb50e529ad4026eee087ca54618d8ad03a Mon Sep 17 00:00:00 2001 From: t Date: Mon, 9 Nov 2026 14:38:19 -0500 Subject: [PATCH 6/7] fix(sandbox): keep every stored deploy_sandbox size as written A stored cpu=1.0 / memory_mb=8192 was read as unset when it carried no size_set marker, so a step saved with that size before this change could be resized by a sizer. A stored size is now always the step's own; only a field the dict leaves out is unset. The size_set marker is no longer needed and is gone. Co-Authored-By: Claude Opus 5.5 --- .../task_step/task_steps/deploy_sandbox.py | 18 +++--------------- .../deploy_sandbox_default_cpu_test.py | 19 +++++++++---------- 2 files changed, 12 insertions(+), 25 deletions(-) diff --git a/src/agent_env/task_step/task_steps/deploy_sandbox.py b/src/agent_env/task_step/task_steps/deploy_sandbox.py index 1c4c127fa..de7d23321 100644 --- a/src/agent_env/task_step/task_steps/deploy_sandbox.py +++ b/src/agent_env/task_step/task_steps/deploy_sandbox.py @@ -30,7 +30,6 @@ class DeploySandboxTaskStep(TaskStep): # step chose is told apart from one it left to this default. DEFAULT_CPU: ClassVar[float] = 1.0 DEFAULT_MEMORY_MB: ClassVar[int] = 8192 - SIZE_SET_KEY: ClassVar[str] = "size_set" def __init__( self, @@ -81,8 +80,6 @@ def to_dict(self) -> dict: base["cpu"] = self.cpu if self.memory_mb is not None: base["memory_mb"] = self.memory_mb - if (self.cpu, self.memory_mb) == (self.DEFAULT_CPU, self.DEFAULT_MEMORY_MB): - base[self.SIZE_SET_KEY] = True # so this pair isn't read back as a step stored before sizes could be unset base["disk_size_gb"] = self.disk_size_gb base["ttl_seconds"] = self.ttl_seconds base["exposed_ports"] = self.exposed_ports @@ -94,17 +91,6 @@ def to_dict(self) -> dict: base["network_policy"] = self.network_policy return base - @classmethod - def _stored_size(cls, data: dict) -> dict: - """``cpu`` and ``memory_mb`` as ``data`` stores them. Before a step could leave them unset, every unsized step - was stored with both at the default, 1.0 and 8192, so that pair is read as unset unless ``to_dict`` marked it as - set: a step stored back then gets the same size unless a sizer chooses one, and a step that sets the pair keeps - it.""" - cpu, memory_mb = data.get("cpu"), data.get("memory_mb") - if (cpu, memory_mb) == (cls.DEFAULT_CPU, cls.DEFAULT_MEMORY_MB) and not data.get(cls.SIZE_SET_KEY): - return {"cpu": None, "memory_mb": None} - return {"cpu": cpu, "memory_mb": memory_mb} - @classmethod def from_dict(cls, data: dict) -> DeploySandboxTaskStep: return cls( @@ -112,7 +98,9 @@ def from_dict(cls, data: dict) -> DeploySandboxTaskStep: sandbox_name=data["sandbox_name"], sandbox_mode=data["sandbox_mode"], image=data.get("image"), - **cls._stored_size(data), + # A stored size is kept as written, whatever it is: only a field the dict leaves out is unset. + cpu=data.get("cpu"), + memory_mb=data.get("memory_mb"), disk_size_gb=data.get("disk_size_gb", 10), ttl_seconds=data.get("ttl_seconds", TaskStep.DEFAULT_TTL_SECONDS), exposed_ports=data.get("exposed_ports"), diff --git a/tst/unit/task_step/deploy_sandbox_default_cpu_test.py b/tst/unit/task_step/deploy_sandbox_default_cpu_test.py index 455354306..543f986bd 100644 --- a/tst/unit/task_step/deploy_sandbox_default_cpu_test.py +++ b/tst/unit/task_step/deploy_sandbox_default_cpu_test.py @@ -43,16 +43,15 @@ def test_a_stored_size_is_kept(size): assert (step.cpu, step.memory_mb) == (size.get("cpu"), size.get("memory_mb")) -def test_a_step_that_sets_the_default_pair_keeps_it_through_a_save(): - step = DeploySandboxTaskStep(id="s", version=1, sandbox_name="box", sandbox_mode="vm", cpu=1.0, memory_mb=8192) - reloaded = DeploySandboxTaskStep.from_dict(step.to_dict()) - assert (reloaded.cpu, reloaded.memory_mb) == (1.0, 8192) - - -def test_the_old_stored_default_reads_as_unsized(): - """Every unsized step was stored as 1.0 / 8192 before a step could leave its size unset.""" - step = DeploySandboxTaskStep.from_dict(_stored(cpu=1.0, memory_mb=8192)) - assert (step.cpu, step.memory_mb) == (None, None) +@pytest.mark.parametrize("stored", [ + DeploySandboxTaskStep(id="s", version=1, sandbox_name="box", sandbox_mode="vm", cpu=1.0, + memory_mb=8192).to_dict(), + _stored(cpu=1.0, memory_mb=8192), +], ids=["saved-now", "saved-before-sizes-could-be-unset"]) +def test_a_saved_default_pair_is_kept(stored): + """1.0 / 8192 is the step's own size however it was saved, so no sizer replaces it.""" + step = DeploySandboxTaskStep.from_dict(stored) + assert (step.cpu, step.memory_mb) == (1.0, 8192) @pytest.mark.asyncio From 722d67bc54c2d2aaa8538115f4304cb214dc7063 Mon Sep 17 00:00:00 2001 From: t Date: Mon, 9 Nov 2026 15:18:45 -0500 Subject: [PATCH 7/7] refactor(sandbox): drop the runtime sizer; keep unset deploy_sandbox sizes Sizes now live on the task: measured usage suggests them and a person applies them as a task edit, so nothing chooses a size while a run deploys. This removes the [sandbox.sizer] seam (sizing.py, may_create_vm, the size-source attribution keys and the provider name hygiene that went with them) and leaves deploy_env and deploy_agent as on main. deploy_sandbox keeps the part that stands on its own: cpu and memory_mb may be unset (None), an unset size is left out of the stored dict, a stored size is kept as written, and an unset size reserves 1 vCPU / 8 GiB as before. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/attribution.py | 13 - src/agent_env/config/describe.py | 3 +- src/agent_env/config/runtime.py | 7 - .../env_providers/env_gateway_provider.py | 3 +- .../sandbox_providers/boat/config.py | 3 - .../sandbox_providers/coreweave/provider.py | 3 +- .../sandbox_providers/daytona/config.py | 3 - .../sandbox_providers/sail_vm/provider.py | 3 +- .../sandbox_providers/sandbox_provider.py | 17 -- .../providers/sandbox_providers/sizing.py | 99 ------- .../sandbox_providers/tenki/provider.py | 3 +- .../task_step/task_steps/deploy_agent.py | 13 +- .../task_step/task_steps/deploy_env.py | 11 +- .../task_step/task_steps/deploy_sandbox.py | 12 +- .../attribution_serialization_test.py | 10 +- tst/unit/task_step/sandbox_sizer_test.py | 241 ------------------ 16 files changed, 21 insertions(+), 423 deletions(-) delete mode 100644 src/agent_env/providers/sandbox_providers/sizing.py delete mode 100644 tst/unit/task_step/sandbox_sizer_test.py diff --git a/src/agent_env/attribution.py b/src/agent_env/attribution.py index 151ff9a18..b53be48ad 100644 --- a/src/agent_env/attribution.py +++ b/src/agent_env/attribution.py @@ -19,13 +19,6 @@ PIPELINE_STEP_KEY = "pipeline_step" RUN_ID_KEY = "run_id" -# Where a sandbox's cpu and memory came from (``step``, ``sizer`` or ``default``; see -# ``providers.sandbox_providers.sizing``). They describe the sandbox rather than name what it is for, so a provider -# that builds names, labels or tags from attribution leaves them out (``identifying``). -CPU_SOURCE_KEY = "cpu_source" -MEMORY_SOURCE_KEY = "memory_source" -SIZE_SOURCE_KEYS = frozenset({CPU_SOURCE_KEY, MEMORY_SOURCE_KEY}) - def attribution_of(step) -> Attribution: """A step's attribution: a copy of ``metadata["attribution"]`` (``{}`` when unset).""" @@ -46,9 +39,3 @@ def deploy_attribution(step, context) -> Attribution: if instance_id and RUN_ID_KEY not in attribution: attribution[RUN_ID_KEY] = instance_id return attribution - - -def identifying(attribution) -> Attribution: - """``attribution`` without the keys that describe a sandbox's size rather than what it is for: what a provider - builds a sandbox's name, labels or tags from, whose length and number are limited.""" - return {key: value for key, value in (attribution or {}).items() if key not in SIZE_SOURCE_KEYS} diff --git a/src/agent_env/config/describe.py b/src/agent_env/config/describe.py index 713eb3b83..e8e291d7f 100644 --- a/src/agent_env/config/describe.py +++ b/src/agent_env/config/describe.py @@ -278,8 +278,7 @@ class FileSection: FileSection("conversations", ("conversations",), "config.runtime", _resolved_keys("conversations")), FileSection("agents", ("agents",), "config.runtime", _resolved_keys("agents"), validated_agents_section), FileSection("sandbox", ("sandbox",), "providers.sandbox_providers.sandbox_provider", - {"default": None, "agent_default": None, "attribution": None, "providers": Named(_SEAM), - "sizer": _SEAM}), + {"default": None, "agent_default": None, "attribution": None, "providers": Named(_SEAM)}), FileSection("state", ("state",), "providers.env_state.env_state_provider", {"providers": Named(_SEAM)}), FileSection("envs", ("envs",), "env.registry", {"impls": None, **_resolved_keys("envs")}), FileSection("artifacts", ("artifacts",), "artifact.registry", {"impls": None, "type_aliases": None}), diff --git a/src/agent_env/config/runtime.py b/src/agent_env/config/runtime.py index 048bf41da..a98b17005 100644 --- a/src/agent_env/config/runtime.py +++ b/src/agent_env/config/runtime.py @@ -320,13 +320,6 @@ def sandbox_registry(self) -> dict[str, dict]: return self._registry("sandbox", sandbox_provider._build_registry) - def sandbox_sizer(self): - """The sizer ``[sandbox.sizer]`` names, which chooses sizes deploy steps leave unset; one that chooses none - when the table is absent.""" - from agent_env.providers.sandbox_providers import sizing - - return self._registry("sandbox_sizer", sizing._build_sizer) - def state_registry(self) -> dict[str, dict]: from agent_env.providers.env_state import env_state_provider 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 89274012b..84ca8cb11 100644 --- a/src/agent_env/providers/env_providers/env_gateway_provider.py +++ b/src/agent_env/providers/env_providers/env_gateway_provider.py @@ -56,7 +56,6 @@ _tool_names, ) from agent_env.providers.sandbox_providers.local_docker.provider import remove_local_work_dir -from agent_env.providers.sandbox_providers import sizing from agent_env.providers.sandbox_providers.sandbox import Sandbox, VmSandbox, port_bindings from agent_env.providers.sandbox_providers.sandbox_provider import ( SANDBOX_MODE_CONTAINER, @@ -969,7 +968,7 @@ async def _deploy_via_containers( deploy = _ContainerDeploy( sandbox_provider=sandbox_provider, cpu=cpu, disk_size_gb=disk_size_gb, ttl_seconds=ttl_seconds, - attribution=sizing.with_default_memory(attribution), + attribution=attribution, ) try: diff --git a/src/agent_env/providers/sandbox_providers/boat/config.py b/src/agent_env/providers/sandbox_providers/boat/config.py index 9dbe0a747..8ebda8570 100644 --- a/src/agent_env/providers/sandbox_providers/boat/config.py +++ b/src/agent_env/providers/sandbox_providers/boat/config.py @@ -9,7 +9,6 @@ from dataclasses import dataclass, field from typing import Any -from agent_env.attribution import identifying from agent_env.config.errors import ConfigError logger = logging.getLogger(__name__) @@ -75,7 +74,6 @@ def _label_value(value: object) -> str: def sandbox_labels(attribution: Mapping[str, Any], ports: Iterable[int], create_token: str) -> dict[str, str]: """Labels for reconnect (ports), cleanup (the create token, the managed marker) and attribution (cost joins). boat keeps at most ``MAX_LABELS``; attribution past that is logged and left off.""" - attribution = identifying(attribution) labels = {MANAGED_LABEL: "true", PORTS_LABEL: ",".join(str(port) for port in ports), CREATE_LABEL: create_token} for key, value in attribution.items(): slug = _UNSAFE_KEY.sub("-", str(key)).strip("-./") @@ -91,7 +89,6 @@ def sandbox_labels(attribution: Mapping[str, Any], ports: Iterable[int], create_ def sandbox_name(attribution: Mapping[str, Any]) -> str: """``ae-`` plus the attribution values in key order: unique, findable by people.""" - attribution = identifying(attribution) slugs = [ _UNSAFE_NAME.sub("-", str(attribution[key])).strip("-") for key in sorted(attribution) diff --git a/src/agent_env/providers/sandbox_providers/coreweave/provider.py b/src/agent_env/providers/sandbox_providers/coreweave/provider.py index 921256d72..df13857a1 100644 --- a/src/agent_env/providers/sandbox_providers/coreweave/provider.py +++ b/src/agent_env/providers/sandbox_providers/coreweave/provider.py @@ -32,7 +32,7 @@ ) from agent_env.artifact.artifacts.docker_image import DockerImageArtifact -from agent_env.attribution import Attribution, identifying +from agent_env.attribution import Attribution from agent_env.config import get_config from agent_env.config.errors import ConfigError from agent_env.providers.sandbox_providers.coreweave.sandbox import ( @@ -108,7 +108,6 @@ def _tag(text: str) -> str: def sandbox_tags(attribution: Attribution, create_tag: str) -> list[str]: """``agentenv``, the create's unique tag, and one ``key.value`` tag per attribution entry.""" - attribution = identifying(attribution) tags = [SANDBOX_TAG, create_tag] tags += [tag for key in sorted(attribution) if (tag := _tag(f"{key}.{attribution[key]}"))] return list(dict.fromkeys(tags)) diff --git a/src/agent_env/providers/sandbox_providers/daytona/config.py b/src/agent_env/providers/sandbox_providers/daytona/config.py index a1be3f9a7..b29b5e980 100644 --- a/src/agent_env/providers/sandbox_providers/daytona/config.py +++ b/src/agent_env/providers/sandbox_providers/daytona/config.py @@ -11,7 +11,6 @@ from daytona import CreateSandboxFromImageParams, CreateSandboxFromSnapshotParams, Resources -from agent_env.attribution import identifying from agent_env.config.errors import ConfigError from agent_env.providers.sandbox_providers.daytona.image import DEFAULT_BASE_IMAGE, HOST_IMAGE_VERSION, host_image @@ -111,7 +110,6 @@ def sandbox_labels( ) -> dict[str, str]: """Labels for attribution (cost joins), reconnect (ports), sweeping (the managed marker) and, on a sandbox booted from the host image, the recipe version that built it.""" - attribution = identifying(attribution) labels = {MANAGED_LABEL: "true", PORTS_LABEL: ",".join(str(port) for port in ports)} if host_image: labels[HOST_IMAGE_LABEL] = HOST_IMAGE_VERSION @@ -127,7 +125,6 @@ def sandbox_labels( def sandbox_name(attribution: Mapping[str, Any]) -> str: """``ae-`` plus the attribution values in key order: unique, findable by people.""" - attribution = identifying(attribution) slugs = [ re.sub(r"[^A-Za-z0-9-]+", "-", str(attribution[key])).strip("-") for key in sorted(attribution) diff --git a/src/agent_env/providers/sandbox_providers/sail_vm/provider.py b/src/agent_env/providers/sandbox_providers/sail_vm/provider.py index 9d957fa59..24dfd5e1d 100644 --- a/src/agent_env/providers/sandbox_providers/sail_vm/provider.py +++ b/src/agent_env/providers/sandbox_providers/sail_vm/provider.py @@ -9,7 +9,7 @@ import uuid from typing import Any, Awaitable, Callable, ClassVar, Self -from agent_env.attribution import PIPELINE_STEP_KEY, RUN_ID_KEY, Attribution, identifying +from agent_env.attribution import PIPELINE_STEP_KEY, RUN_ID_KEY, Attribution from agent_env.config import get_config from agent_env.config.errors import ConfigError from agent_env.providers.sandbox_providers.sail_vm import _sdk @@ -65,7 +65,6 @@ def sailbox_shape(cpu: float, memory_mb: int, disk_size_gb: float, *, min_size: def sailbox_name(attribution: Attribution) -> str: """``ae-`` plus the attribution values in key order, slugged: for people and ``list(search=)``.""" - attribution = identifying(attribution) slugs = [re.sub(r"[^A-Za-z0-9]+", "-", str(attribution[key])).strip("-") for key in sorted(attribution)] return "-".join(["ae", uuid.uuid4().hex[:8], *filter(None, slugs)])[:_MAX_NAME_LENGTH].rstrip("-") diff --git a/src/agent_env/providers/sandbox_providers/sandbox_provider.py b/src/agent_env/providers/sandbox_providers/sandbox_provider.py index 10604a8a3..628984426 100644 --- a/src/agent_env/providers/sandbox_providers/sandbox_provider.py +++ b/src/agent_env/providers/sandbox_providers/sandbox_provider.py @@ -704,23 +704,6 @@ def _build_sandbox_provider_from_section(name: str, section: dict) -> SandboxPro return _install_sandbox_type_guard(provider, name) -def may_create_vm(spec: Optional[str] = None, *, agent: bool = False) -> bool: - """Whether a sandbox from ``spec`` (unset: ``[sandbox] default``, or ``agent_default`` with ``agent``) may be a VM: - a backend it names creates VMs, as a gateway on one deploys into a VM. Read from the backends' classes, so nothing - is built; True when a backend can't be resolved, as the safe answer for something that sizes it.""" - from agent_env.config import load_impl - - spec = spec or (_agent_sandbox_spec() if agent else _default_sandbox_spec()) - try: - registry = _get_sandbox_registry() - classes = [load_impl(registry[name]["impl"], SandboxProvider) - for name in (s.strip() for s in spec.split(",")) if name] - except Exception as e: - logger.debug(f"Could not tell whether sandbox spec {spec!r} creates VMs, so taking it that it may: {e}") - return True - return not classes or any(SANDBOX_MODE_VM in cls.SUPPORTED_SANDBOX_MODES for cls in classes) - - def build_sandbox_provider(spec: str) -> SandboxProvider: """Build a provider (or a comma-separated fallback chain) from backend names, resolved through the open registry (built-ins + config.toml [sandbox.providers]).""" diff --git a/src/agent_env/providers/sandbox_providers/sizing.py b/src/agent_env/providers/sandbox_providers/sizing.py deleted file mode 100644 index bec7543da..000000000 --- a/src/agent_env/providers/sandbox_providers/sizing.py +++ /dev/null @@ -1,99 +0,0 @@ -"""The cpu and memory a deploy step reserves for whatever size it leaves unset. - -A step's own value always wins. For each one it leaves unset, the sizer config.toml's ``[sandbox.sizer]`` names -(``impl`` and ``config``, as for ``[runner]``, or the ``impl`` alone as a string) may choose one; otherwise the step's -default applies, or, with none, the provider's. Where each value came from is added to the sandbox's attribution, so -it reaches the provider with the rest. - -The sizer is told whether the step may deploy onto a VM, whose memory is a hard limit on everything in it, so it can -decline to size one. A sizer that fails never fails the deploy: its error is logged and it is taken to choose nothing. -""" - -from __future__ import annotations - -import logging -from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Literal, Optional - -from agent_env.attribution import CPU_SOURCE_KEY, MEMORY_SOURCE_KEY, Attribution -from agent_env.config import loader, runtime - -if TYPE_CHECKING: - from agent_env.config.runtime import Config - -logger = logging.getLogger(__name__) - -__all__ = ["CPU_SOURCE_KEY", "MEMORY_SOURCE_KEY", "SandboxSize", "SandboxSizer", "resolve", "with_default_memory"] - -# ``provider``: the sandbox doesn't take the step's value at all (a gateway's containers share its cpu but not its -# memory), so it gets its provider's, and a size chosen for the step says nothing about it. -Source = Literal["step", "sizer", "default", "provider"] - - -@dataclass(frozen=True) -class SandboxSize: - cpu: Optional[float] = None - memory_mb: Optional[int] = None - - -class SandboxSizer: - """Chooses a size for what a deploy step leaves unset. This base chooses nothing, so the step's or provider's - default applies; a subclass returns a value for each field it sizes and None for the rest. - - ``vm`` is True when the step may deploy onto a VM (its sandbox mode is ``vm``, or a backend its sandbox type names - creates VMs), where the memory chosen is the limit on everything the VM runs.""" - - @classmethod - def from_config(cls, **config: Any) -> SandboxSizer: - return cls(**config) - - def size(self, step: Any, attribution: Attribution, *, vm: bool) -> SandboxSize: - return SandboxSize() - - -def _build_sizer(source: Config) -> SandboxSizer: - raw = source.section("sandbox").get("sizer") - section = {"impl": raw} if isinstance(raw, str) else dict(raw or {}) - if not section: - return SandboxSizer() - return loader.build_store(section, SandboxSizer, secret_resolver=source._resolve_secret) - - -def _pick(own: Any, chosen: Any, default: Any) -> tuple[Any, Source]: - if own is not None: - return own, "step" - if chosen is not None: - return chosen, "sizer" - return default, "default" - - -def with_default_memory(attribution: Optional[Attribution]) -> Optional[Attribution]: - """``attribution`` for a sandbox created without the step's memory, which gets its provider's: its memory source - is ``provider`` whatever the step's was.""" - if not attribution or MEMORY_SOURCE_KEY not in attribution: - return attribution - return {**attribution, MEMORY_SOURCE_KEY: "provider"} - - -def _chosen(step: Any, attribution: Attribution, vm: bool) -> SandboxSize: - sizer = runtime.get_config().sandbox_sizer() # a misconfigured sizer is a config error, raised as one - try: - chosen = sizer.size(step, dict(attribution), vm=vm) - if not isinstance(chosen, SandboxSize): - raise TypeError(f"returned {type(chosen).__name__}, not a SandboxSize") - return chosen - except Exception as e: - logger.warning(f"{type(sizer).__name__} could not size step {getattr(step, 'id', None)!r}, so it gets its " - f"default: {type(e).__name__}: {e}") - return SandboxSize() - - -def resolve(step: Any, attribution: Attribution, default: SandboxSize = SandboxSize(), *, - vm: bool) -> tuple[SandboxSize, Attribution]: - """The size ``step`` reserves, from its ``cpu`` and ``memory_mb``, and ``attribution`` with where each came from. - ``vm``: whether the step may deploy onto a VM, which the sizer is told.""" - own = SandboxSize(step.cpu, step.memory_mb) - chosen = SandboxSize() if None not in (own.cpu, own.memory_mb) else _chosen(step, attribution, vm) - cpu, cpu_source = _pick(own.cpu, chosen.cpu, default.cpu) - memory_mb, memory_source = _pick(own.memory_mb, chosen.memory_mb, default.memory_mb) - return SandboxSize(cpu, memory_mb), {**attribution, CPU_SOURCE_KEY: cpu_source, MEMORY_SOURCE_KEY: memory_source} diff --git a/src/agent_env/providers/sandbox_providers/tenki/provider.py b/src/agent_env/providers/sandbox_providers/tenki/provider.py index 42ea9b3d8..a52d5ce0a 100644 --- a/src/agent_env/providers/sandbox_providers/tenki/provider.py +++ b/src/agent_env/providers/sandbox_providers/tenki/provider.py @@ -20,7 +20,7 @@ from tenki import AsyncClient, SessionNotFoundError -from agent_env.attribution import Attribution, identifying +from agent_env.attribution import Attribution from agent_env.config.errors import ConfigError from agent_env.providers.sandbox_providers.sandbox import NetworkPolicy, NetworkPolicyUnsupportedError from agent_env.providers.sandbox_providers.sandbox_provider import SandboxProvider, apply_default_attribution @@ -103,7 +103,6 @@ def egress_kwargs(policy: NetworkPolicy) -> dict[str, Any]: def session_name(attribution: Attribution) -> str: """``agentenv-`` and the attribution values, slugged into Tenki's 63-character name.""" - attribution = identifying(attribution) slugs = [re.sub(r"[^a-z0-9]+", "-", str(attribution[key]).lower()).strip("-") for key in sorted(attribution)] return "-".join(["agentenv", uuid.uuid4().hex[:8], *filter(None, slugs)])[:63].rstrip("-") diff --git a/src/agent_env/task_step/task_steps/deploy_agent.py b/src/agent_env/task_step/task_steps/deploy_agent.py index 30d2bb4ee..cbacdc804 100644 --- a/src/agent_env/task_step/task_steps/deploy_agent.py +++ b/src/agent_env/task_step/task_steps/deploy_agent.py @@ -30,8 +30,6 @@ ) from agent_env.a2a_agent.staging import StagedObjectStore, transfer_store from agent_env.config import get_config -from agent_env.providers.sandbox_providers import sizing -from agent_env.providers.sandbox_providers.sandbox_provider import may_create_vm from agent_env.env.env import DeployedEnv, DeployedSandboxEnv, Env from agent_env.env.gateway import AGENT_ENV_ROLE_HEADER from agent_env.providers.sandbox_providers.local_docker.provider import host_url_for @@ -403,12 +401,11 @@ async def _execute(self, context: TaskStepContext, deployed_agents: list) -> Tas deploy_kwargs["sandbox_type"] = resolved_sandbox_type resolved_ttl = self._resolved_ttl_seconds(context) deploy_kwargs["ttl_seconds"] = resolved_ttl - size, deploy_kwargs["attribution"] = sizing.resolve( - self, deploy_attribution(self, context), vm=may_create_vm(resolved_sandbox_type, agent=True)) - if size.cpu is not None: - deploy_kwargs["cpu"] = size.cpu - if size.memory_mb is not None: - deploy_kwargs["memory"] = size.memory_mb + if self.cpu is not None: + deploy_kwargs["cpu"] = self.cpu + if self.memory_mb is not None: + deploy_kwargs["memory"] = self.memory_mb + deploy_kwargs["attribution"] = deploy_attribution(self, context) _policy_override = user_overrides.get("network_policy") resolved_policy = _policy_override if _policy_override is not None else self.network_policy if resolved_policy is not None: diff --git a/src/agent_env/task_step/task_steps/deploy_env.py b/src/agent_env/task_step/task_steps/deploy_env.py index 8bab03036..375a4c515 100644 --- a/src/agent_env/task_step/task_steps/deploy_env.py +++ b/src/agent_env/task_step/task_steps/deploy_env.py @@ -15,8 +15,6 @@ from agent_env.entity.scopes import FieldScope from agent_env.task_step.task_step import TaskStep, TaskStepDependency from agent_env.attribution import deploy_attribution -from agent_env.providers.sandbox_providers import sizing -from agent_env.providers.sandbox_providers.sandbox_provider import may_create_vm logger = logging.getLogger(__name__) @@ -150,12 +148,11 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: resolved_env_state_type, resolved_env_state_instance_id = self.run_state(user_overrides) _ttl_override = user_overrides.get("ttl_seconds") resolved_ttl = _ttl_override if _ttl_override is not None else self.ttl_seconds - size, attribution = sizing.resolve(self, deploy_attribution(self, context), - vm=may_create_vm(resolved_sandbox_type)) + attribution = deploy_attribution(self, context) logger.info( f"Deploying env '{self.env_id}' (ttl={resolved_ttl}s, " f"disk_size_gb={self.disk_size_gb}, gateway_mode={self.gateway_mode}, " - f"cpu={size.cpu}, memory_mb={size.memory_mb}, " + f"cpu={self.cpu}, memory_mb={self.memory_mb}, " f"sandbox_type={resolved_sandbox_type}, env_state_type={resolved_env_state_type}, " f"env_state_instance_id={resolved_env_state_instance_id}, " f"attribution={attribution})" @@ -164,8 +161,8 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: "ttl_seconds": resolved_ttl, "disk_size_gb": self.disk_size_gb, "gateway_mode": GatewayMode(self.gateway_mode), - "cpu": size.cpu, - "memory_mb": size.memory_mb, + "cpu": self.cpu, + "memory_mb": self.memory_mb, "attribution": attribution, } if resolved_sandbox_type is not None: diff --git a/src/agent_env/task_step/task_steps/deploy_sandbox.py b/src/agent_env/task_step/task_steps/deploy_sandbox.py index de7d23321..4f146dcf6 100644 --- a/src/agent_env/task_step/task_steps/deploy_sandbox.py +++ b/src/agent_env/task_step/task_steps/deploy_sandbox.py @@ -12,7 +12,6 @@ from agent_env.entity.scopes import FieldScope from agent_env.task_step.task_step import TaskStep, TaskStepDependency from agent_env.attribution import deploy_attribution -from agent_env.providers.sandbox_providers import sizing if TYPE_CHECKING: from agent_env.providers.sandbox_providers.sandbox_provider import SandboxProvider @@ -119,8 +118,7 @@ def _pull_refusal(self, link: SandboxProvider) -> str | None: return None async def execute(self, context: TaskStepContext) -> TaskStepContext: - from agent_env.providers.sandbox_providers.sandbox_provider import ( - build_sandbox_provider, get_sandbox_provider, may_create_vm) + from agent_env.providers.sandbox_providers.sandbox_provider import build_sandbox_provider, get_sandbox_provider if any(s.sandbox_name == self.sandbox_name for s in context.deployed_sandboxes): raise RuntimeError(f"Sandbox '{self.sandbox_name}' is already deployed") @@ -136,11 +134,9 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: _ttl_override = user_overrides.get("ttl_seconds") resolved_ttl = _ttl_override if _ttl_override is not None else self.ttl_seconds - size, attribution = sizing.resolve(self, deploy_attribution(self, context), - sizing.SandboxSize(self.DEFAULT_CPU, self.DEFAULT_MEMORY_MB), - # A VM backend runs even a container in a VM of its own, sized by it. - vm=self.sandbox_mode == "vm" or may_create_vm(resolved_sandbox_type)) - cpu, memory_mb = size.cpu, size.memory_mb + attribution = deploy_attribution(self, context) + cpu = self.cpu if self.cpu is not None else self.DEFAULT_CPU + memory_mb = self.memory_mb if self.memory_mb is not None else self.DEFAULT_MEMORY_MB logger.info( f"Deploying sandbox '{self.sandbox_name}' (mode={self.sandbox_mode}, " f"image={self.image}, cpu={cpu}, memory_mb={memory_mb}, " diff --git a/tst/unit/task_step/attribution_serialization_test.py b/tst/unit/task_step/attribution_serialization_test.py index 823487115..57db30565 100644 --- a/tst/unit/task_step/attribution_serialization_test.py +++ b/tst/unit/task_step/attribution_serialization_test.py @@ -11,7 +11,6 @@ import pytest from agent_env.attribution import attribution_of -from agent_env.providers.sandbox_providers.sizing import CPU_SOURCE_KEY, MEMORY_SOURCE_KEY from agent_env.env.env import DeployedGatewayEnv from agent_env.task_step.context import TaskStepContext from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep @@ -79,9 +78,6 @@ def test_a_top_level_key_in_a_document_is_not_attribution(factory): # --- the dict is what reaches the env -------------------------------------------- -# What deploy_env adds to a step's attribution: where its unset cpu and memory came from. -UNSIZED = {CPU_SOURCE_KEY: "default", MEMORY_SOURCE_KEY: "default"} - @pytest.mark.asyncio async def test_deploying_an_env_passes_the_attribution_dict(): @@ -102,11 +98,11 @@ async def deploy(*, attribution=None, **kwargs): Env.get.return_value = env await step.execute(TaskStepContext()) - assert passed["attribution"] == {**ATTRIBUTION, **UNSIZED} + assert passed["attribution"] == ATTRIBUTION @pytest.mark.asyncio -async def test_deploying_an_env_with_no_attribution_passes_only_its_size_sources(): +async def test_deploying_an_env_with_no_attribution_passes_an_empty_dict(): passed = {} async def deploy(*, attribution=None, **kwargs): @@ -123,4 +119,4 @@ async def deploy(*, attribution=None, **kwargs): Env.get.return_value = env await _deploy_env().execute(TaskStepContext()) - assert passed["attribution"] == UNSIZED + assert passed["attribution"] == {} diff --git a/tst/unit/task_step/sandbox_sizer_test.py b/tst/unit/task_step/sandbox_sizer_test.py deleted file mode 100644 index a837d3749..000000000 --- a/tst/unit/task_step/sandbox_sizer_test.py +++ /dev/null @@ -1,241 +0,0 @@ -"""A deploy step reserves the size it sets, else what the configured sizer chooses for what it leaves unset, else its -default or the provider's, and says on the sandbox's attribution where each value came from.""" - -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest - -from agent_env.config import ConfigError -from agent_env.env.env import DeployedGatewayEnv -from agent_env.providers.sandbox_providers import sizing -from agent_env.providers.sandbox_providers.sandbox import NetworkPolicy -from agent_env.providers.sandbox_providers.sizing import SandboxSize, SandboxSizer -from agent_env.task_step.context import TaskStepContext -from agent_env.task_step.task_steps.deploy_env import DeployEnvTaskStep -from agent_env.task_step.task_steps.deploy_sandbox import DeploySandboxTaskStep -from tst.util.config import config_with_document - - -class HalfSizer(SandboxSizer): - """Sizes every step's memory to ``memory_mb``, and its cpu only when ``cpu`` is given; records what it was asked.""" - - asked: list = [] - - def __init__(self, memory_mb: int, cpu: float | None = None): - self.memory_mb, self.cpu = memory_mb, cpu - - def size(self, step, attribution, *, vm): - HalfSizer.asked.append((step.id, dict(attribution), vm)) - return SandboxSize(cpu=self.cpu, memory_mb=self.memory_mb) - - -def _config(sizer: dict | None = None): - return config_with_document({"sandbox": {"sizer": sizer}} if sizer else {}) - - -@pytest.fixture(autouse=True) -def _reset_asked(): - HalfSizer.asked = [] - - -class _Step: - id = "s" - - def __init__(self, cpu=None, memory_mb=None): - self.cpu, self.memory_mb = cpu, memory_mb - - -@pytest.mark.parametrize("sizer, step, default, expected", [ - (None, _Step(), SandboxSize(), (None, None, "default", "default")), - (None, _Step(), SandboxSize(1.0, 8192), (1.0, 8192, "default", "default")), - ({"impl": HalfSizer, "config": {"memory_mb": 2048}}, _Step(), SandboxSize(1.0, 8192), - (1.0, 2048, "default", "sizer")), - ({"impl": HalfSizer, "config": {"memory_mb": 2048, "cpu": 0.5}}, _Step(), SandboxSize(), - (0.5, 2048, "sizer", "sizer")), - ({"impl": HalfSizer, "config": {"memory_mb": 2048, "cpu": 0.5}}, _Step(memory_mb=16384), SandboxSize(), - (0.5, 16384, "sizer", "step")), - ({"impl": HalfSizer, "config": {"memory_mb": 2048, "cpu": 0.5}}, _Step(cpu=4.0, memory_mb=16384), SandboxSize(), - (4.0, 16384, "step", "step")), -], ids=["no-sizer", "no-sizer-step-default", "sizer-memory-only", "sizer-both", "step-memory-wins", "step-sets-both"]) -def test_each_value_comes_from_the_step_then_the_sizer_then_the_default(sizer, step, default, expected): - with patch("agent_env.config.runtime.get_config", return_value=_config(sizer)): - size, attribution = sizing.resolve(step, {"run_id": "r"}, default, vm=False) - - assert (size.cpu, size.memory_mb, attribution["cpu_source"], attribution["memory_source"]) == expected - assert attribution["run_id"] == "r" - - -def test_a_step_that_sets_both_never_asks_the_sizer(): - with patch("agent_env.config.runtime.get_config", return_value=_config({"impl": HalfSizer, "config": {"memory_mb": 1}})): - sizing.resolve(_Step(cpu=1.0, memory_mb=1024), {}, vm=False) - assert HalfSizer.asked == [] - - -def test_a_sizer_that_is_not_one_is_refused(): - with patch("agent_env.config.runtime.get_config", return_value=_config({"impl": dict})): - with pytest.raises(ConfigError, match="not a subclass of SandboxSizer"): - sizing.resolve(_Step(), {}, vm=False) - - -@pytest.mark.asyncio -async def test_deploy_env_hands_the_sized_values_and_their_sources_to_the_env(): - env = MagicMock() - env.deploy = AsyncMock(return_value=DeployedGatewayEnv(env_id="e", env_version=1, gateway_url="http://gw", - mcp_url="http://mcp", db_web_url=None, sandbox_id="s1")) - step = DeployEnvTaskStep(id="deploy", version=None, env_id="e", env_version=1, cpu=2.0) - config = _config({"impl": HalfSizer, "config": {"memory_mb": 2048, "cpu": 0.5}}) - - with patch("agent_env.config.runtime.get_config", return_value=config), patch("agent_env.env.env.Env") as Env, \ - patch("agent_env.task_step.task_steps.deploy_env.may_create_vm", return_value=False): - Env.get.return_value = env - await step.execute(TaskStepContext(instance_id="run-1")) - - kwargs = env.deploy.await_args.kwargs - assert (kwargs["cpu"], kwargs["memory_mb"]) == (2.0, 2048) - assert {k: kwargs["attribution"][k] for k in ("cpu_source", "memory_source", "run_id")} == { - "cpu_source": "step", "memory_source": "sizer", "run_id": "run-1"} - assert HalfSizer.asked == [("deploy", {"run_id": "run-1"}, False)] - - -@pytest.mark.asyncio -async def test_deploy_sandbox_falls_back_to_its_own_default_when_the_sizer_chooses_nothing(): - sandbox = AsyncMock() - sandbox.sandbox_id, sandbox.mode, sandbox.type = "sb-1", "vm", "modal_vm" - sandbox.tunnel_urls, sandbox.vnc_url, sandbox.network_policy = {}, None, NetworkPolicy() - provider = AsyncMock(ON_THIS_MACHINE=False) - provider.links = (provider,) - provider.create_vm.return_value = sandbox - step = DeploySandboxTaskStep(id="s", version=1, sandbox_name="box", sandbox_mode="vm", sandbox_type="any") - - with patch("agent_env.config.runtime.get_config", return_value=_config({"impl": SandboxSizer})), \ - patch("agent_env.providers.sandbox_providers.sandbox_provider.build_sandbox_provider", return_value=provider): - await step.execute(TaskStepContext()) - - kwargs = provider.create_vm.call_args.kwargs - assert (kwargs["cpu"], kwargs["memory"]) == (1.0, 8192) - assert (kwargs["attribution"]["cpu_source"], kwargs["attribution"]["memory_source"]) == ("default", "default") - - -@pytest.mark.parametrize("attribution, expected", [ - ({"run_id": "r", "cpu_source": "sizer", "memory_source": "sizer"}, - {"run_id": "r", "cpu_source": "sizer", "memory_source": "provider"}), - ({"run_id": "r", "memory_source": "step"}, {"run_id": "r", "memory_source": "provider"}), - ({"run_id": "r"}, {"run_id": "r"}), - (None, None), -], ids=["sized", "step-memory", "no-sources", "none"]) -def test_a_sandbox_made_without_the_steps_memory_says_its_provider_chose_it(attribution, expected): - assert sizing.with_default_memory(attribution) == expected - - -class RaisingSizer(SandboxSizer): - def size(self, step, attribution, *, vm): - raise RuntimeError("store down") - - -class WrongShapeSizer(SandboxSizer): - def size(self, step, attribution, *, vm): - return {"cpu": 0.5} - - -@pytest.mark.parametrize("impl", [RaisingSizer, WrongShapeSizer], ids=["raises", "wrong-shape"]) -def test_a_sizer_that_fails_leaves_the_default_and_never_fails_the_deploy(impl, caplog): - with patch("agent_env.config.runtime.get_config", return_value=_config({"impl": impl})): - size, attribution = sizing.resolve(_Step(), {}, SandboxSize(1.0, 8192), vm=False) - - assert (size, attribution["cpu_source"], attribution["memory_source"]) == (SandboxSize(1.0, 8192), "default", - "default") - assert f"{impl.__name__} could not size step 's'" in caplog.text - - -def test_a_sizer_named_by_a_string_is_built(): - config = config_with_document({"sandbox": {"sizer": f"{__name__}:RaisingSizer"}}) - assert type(config.sandbox_sizer()) is RaisingSizer - - -@pytest.mark.asyncio -@pytest.mark.parametrize("mode, create, extra, sandbox_type, vm", [ - ("vm", "create_vm", {}, "modal", True), - ("container", "create_container", {"image": "example.com/box:v1", "port": 80}, "modal", False), - # A VM backend runs the container in a VM of its own, whose memory is the size chosen. - ("container", "create_container", {"image": "example.com/box:v1", "port": 80}, "modal_vm", True), -], ids=["vm", "container", "container-on-a-vm-backend"]) -async def test_deploy_sandbox_tells_the_sizer_whether_it_deploys_a_vm(mode, create, extra, sandbox_type, vm): - sandbox = AsyncMock() - sandbox.sandbox_id, sandbox.mode, sandbox.type = "sb-1", mode, "modal" - sandbox.tunnel_urls, sandbox.vnc_url, sandbox.network_policy = {}, None, NetworkPolicy() - provider = AsyncMock(ON_THIS_MACHINE=False) - provider.links = (provider,) - getattr(provider, create).return_value = sandbox - step = DeploySandboxTaskStep(id="s", version=1, sandbox_name="box", sandbox_mode=mode, sandbox_type=sandbox_type, - **extra) - - with patch("agent_env.config.runtime.get_config", return_value=_config({"impl": HalfSizer, - "config": {"memory_mb": 2048}})), \ - patch("agent_env.providers.sandbox_providers.sandbox_provider.build_sandbox_provider", return_value=provider): - await step.execute(TaskStepContext()) - - assert HalfSizer.asked[0][2] is vm - - -@pytest.mark.parametrize("spec, vm", [ - ("modal", False), ("modal_vm", True), ("modal,modal_vm", True), ("no_such_backend", True), -], ids=["container", "vm", "chain-with-a-vm", "unknown"]) -def test_a_spec_may_create_a_vm_when_any_backend_it_names_does(spec, vm): - from agent_env.providers.sandbox_providers.sandbox_provider import may_create_vm - - assert may_create_vm(spec) is vm - - -@pytest.mark.asyncio -@pytest.mark.parametrize("sandbox_type, vm", [("modal", False), ("modal_vm", True)], ids=["container", "vm"]) -async def test_deploy_env_tells_the_sizer_whether_its_sandbox_type_may_be_a_vm(sandbox_type, vm): - env = MagicMock() - env.deploy = AsyncMock(return_value=DeployedGatewayEnv(env_id="e", env_version=1, gateway_url="http://gw", - mcp_url="http://mcp", db_web_url=None, sandbox_id="s1")) - step = DeployEnvTaskStep(id="deploy", version=None, env_id="e", env_version=1, sandbox_type=sandbox_type) - - with patch("agent_env.config.runtime.get_config", return_value=_config({"impl": HalfSizer, - "config": {"memory_mb": 2048}})), \ - patch("agent_env.env.env.Env") as Env: - Env.get.return_value = env - await step.execute(TaskStepContext(instance_id="run-1")) - - assert HalfSizer.asked == [("deploy", {"run_id": "run-1"}, vm)] - - -class _Deployed(Exception): - """Stops a deploy_agent step once it has asked its agent to deploy, with what it asked.""" - - -@pytest.mark.asyncio -async def test_deploy_agent_hands_the_sized_values_and_their_sources_to_the_agent(): - from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep - - agent = MagicMock(metadata={}) - agent.deploy = AsyncMock(side_effect=_Deployed) - step = DeployAgentTaskStep(id="agent", version=1, env_ids=[], a2a_agent_id="solver", memory_mb=16384, - sandbox_type="modal") - config = _config({"impl": HalfSizer, "config": {"memory_mb": 2048, "cpu": 0.5}}) - - with patch("agent_env.config.runtime.get_config", return_value=config), \ - patch("agent_env.a2a_agent.A2AAgent") as A2AAgent, \ - patch.object(DeployAgentTaskStep, "_resolve_agent_ref", return_value=("solver", 1)), \ - patch.object(DeployAgentTaskStep, "_model_endpoint", return_value={}): - A2AAgent.get.return_value = agent - with pytest.raises(_Deployed): - await step.execute(TaskStepContext(instance_id="run-1")) - - kwargs = agent.deploy.await_args.kwargs - assert (kwargs["cpu"], kwargs["memory"]) == (0.5, 16384) - assert (kwargs["attribution"]["cpu_source"], kwargs["attribution"]["memory_source"]) == ("sizer", "step") - assert HalfSizer.asked == [("agent", {"run_id": "run-1"}, False)] - - -def test_where_a_size_came_from_stays_out_of_names_labels_and_tags(): - from agent_env.providers.sandbox_providers.boat import config as boat - from agent_env.providers.sandbox_providers.coreweave.provider import sandbox_tags - - attribution = {"project_id": "p", "cpu_source": "sizer", "memory_source": "default"} - assert "sizer" not in boat.sandbox_name(attribution) and "default" not in boat.sandbox_name(attribution) - assert not any("source" in key for key in boat.sandbox_labels(attribution, [80], "t")) - assert not any("source" in tag for tag in sandbox_tags(attribution, "c"))