feat(sandbox): let a configured sizer choose the sizes deploy steps leave unset

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,

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,27 +29,48 @@ 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
 @pytest.mark.parametrize("mode, create, extra", [
     ("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