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..4f146dcf6 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, @@ -70,8 +74,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 @@ -90,8 +97,9 @@ 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), + # 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"), @@ -127,9 +135,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 +152,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 +168,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..543f986bd 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,29 @@ 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 "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_cpu_is_kept(): - assert DeploySandboxTaskStep.from_dict(_stored(cpu=2.0)).cpu == 2.0 +@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")) + + +@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 @@ -42,14 +59,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