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 581e15b3..ae00fc47 100644 --- a/src/agent_env/task_step/task_steps/deploy_sandbox.py +++ b/src/agent_env/task_step/task_steps/deploy_sandbox.py @@ -31,7 +31,7 @@ def __init__( sandbox_name: str, sandbox_mode: str, image: Optional[str] = None, - cpu: float = 2.0, + cpu: float = 1.0, memory_mb: int = 8192, disk_size_gb: float = 10, ttl_seconds: int = TaskStep.DEFAULT_TTL_SECONDS, @@ -88,7 +88,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", 2.0), + cpu=data.get("cpu", 1.0), memory_mb=data.get("memory_mb", 8192), disk_size_gb=data.get("disk_size_gb", 10), ttl_seconds=data.get("ttl_seconds", TaskStep.DEFAULT_TTL_SECONDS), diff --git a/tst/unit/task_step/deploy_sandbox_default_cpu_test.py b/tst/unit/task_step/deploy_sandbox_default_cpu_test.py new file mode 100644 index 00000000..cffa512d --- /dev/null +++ b/tst/unit/task_step/deploy_sandbox_default_cpu_test.py @@ -0,0 +1,55 @@ +"""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.""" + +from unittest.mock import AsyncMock, patch + +import pytest + +from agent_env.providers.sandbox_providers.sandbox import NetworkPolicy +from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.task_steps.deploy_sandbox import DeploySandboxTaskStep + + +def _provider(): + 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 + provider.create_container.return_value = sandbox + return provider + + +def _stored(**fields) -> dict: + return {"id": "s", "version": 1, "type": "deploy_sandbox", "sandbox_name": "box", "sandbox_mode": "vm", **fields} + + +@pytest.mark.parametrize("step", [ + 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_a_stored_cpu_is_kept(): + assert DeploySandboxTaskStep.from_dict(_stored(cpu=2.0)).cpu == 2.0 + + +@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} + step = DeploySandboxTaskStep(id="s", version=1, sandbox_name="box", sandbox_mode=mode, sandbox_type="any", + **sized, **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