diff --git a/src/agent_env/providers/sandbox_providers/e2b/provider.py b/src/agent_env/providers/sandbox_providers/e2b/provider.py index 050d74a4..c5c11b4c 100644 --- a/src/agent_env/providers/sandbox_providers/e2b/provider.py +++ b/src/agent_env/providers/sandbox_providers/e2b/provider.py @@ -381,3 +381,7 @@ async def get_sandbox(self, sandbox_id: str) -> E2BSandbox: ) sandbox.network_policy = None return sandbox + + async def terminate_sandbox(self, sandbox_id: str) -> None: + if not await self._get_sandbox_cls().kill(sandbox_id, api_key=self._api_key): + logger.info("E2B sandbox %s was already absent during terminate", sandbox_id) diff --git a/src/agent_env/providers/sandbox_providers/modal_sandbox.py b/src/agent_env/providers/sandbox_providers/modal_sandbox.py index cc647e8d..cb09bd87 100644 --- a/src/agent_env/providers/sandbox_providers/modal_sandbox.py +++ b/src/agent_env/providers/sandbox_providers/modal_sandbox.py @@ -90,6 +90,14 @@ def _tunnel_url(t) -> str: return f"https://{t.host}" + (f":{t.port}" if t.port != 443 else "") +async def _terminate_by_id(client, sandbox_id: str) -> None: + try: + sb = await modal.Sandbox.from_id.aio(sandbox_id, client=client) + await sb.terminate.aio() + except modal.exception.NotFoundError: + logger.info("Modal sandbox %s was already gone at terminate", sandbox_id) + + def _modal_network_kwargs(policy: NetworkPolicy) -> dict: """``policy`` as Modal ``_experimental_create`` kwargs. @@ -480,6 +488,9 @@ async def get_sandbox(self, sandbox_id: str) -> Sandbox: i6pn_address = await _resolve_i6pn_address(sb) return self._sandbox_cls(sb, tunnel_urls, i6pn_address=i6pn_address) + async def terminate_sandbox(self, sandbox_id: str) -> None: + await _terminate_by_id(await self._get_client(), sandbox_id) + async def _registry_image(self, image_name: str) -> modal.Image: """``image_name`` pulled from its registry, with the image store's credentials for it, if any.""" image_store = get_config().get_image_store_at(image_name) diff --git a/src/agent_env/providers/sandbox_providers/modal_vm_sandbox.py b/src/agent_env/providers/sandbox_providers/modal_vm_sandbox.py index d615d39a..0cec9499 100644 --- a/src/agent_env/providers/sandbox_providers/modal_vm_sandbox.py +++ b/src/agent_env/providers/sandbox_providers/modal_vm_sandbox.py @@ -32,6 +32,7 @@ _attribution_tags, _log_sandbox_started, _resolve_app_base_name, + _terminate_by_id, _tunnel_url, _modal_network_kwargs, ) @@ -412,6 +413,9 @@ async def get_sandbox(self, sandbox_id: str) -> ModalVmSandbox: # dockerd is already running on a live sandbox; do not re-run setup here. return ModalVmSandbox(sb, tunnel_urls) + async def terminate_sandbox(self, sandbox_id: str) -> None: + await _terminate_by_id(await self._get_client(), sandbox_id) + # ── Object downloads (internal) ──────────────────────────────────────────────────── # Object downloads run as one aria2c per object: a pool of HTTP range connections with small pieces, 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 76170d72..24dfd5e1 100644 --- a/src/agent_env/providers/sandbox_providers/sail_vm/provider.py +++ b/src/agent_env/providers/sandbox_providers/sail_vm/provider.py @@ -429,3 +429,14 @@ async def get_sandbox(self, sandbox_id: str) -> SailVmSandbox: sandbox_id, applied, ) return sandbox + + async def terminate_sandbox(self, sandbox_id: str) -> None: + """Through the handle: terminating also deletes the Sailbox's model-key policy, which only its applied + policy names.""" + sdk, _ = await self._connect() + try: + sandbox = await self.get_sandbox(sandbox_id) + except sdk.NotFoundError: + logger.info("Sailbox %s was already gone at terminate", sandbox_id) + return + await sandbox.terminate() diff --git a/src/agent_env/providers/sandbox_providers/sandbox_provider.py b/src/agent_env/providers/sandbox_providers/sandbox_provider.py index 522886ae..314de40b 100644 --- a/src/agent_env/providers/sandbox_providers/sandbox_provider.py +++ b/src/agent_env/providers/sandbox_providers/sandbox_provider.py @@ -162,6 +162,11 @@ def _container_args(self, sandbox: VmSandbox, *, image_name: str, port: int, env async def get_sandbox(self, sandbox_id: str) -> Sandbox: raise NotImplementedError(f"{type(self).__name__} does not support get_sandbox") + async def terminate_sandbox(self, sandbox_id: str) -> None: + """Terminate a sandbox this provider made, by id; one that has already finished or is gone is terminated. + A backend that can terminate without a live handle overrides this, so teardown doesn't need one.""" + await (await self.get_sandbox(sandbox_id)).terminate() + async def close(self) -> None: pass diff --git a/src/agent_env/providers/sandbox_providers/smol_vm/provider.py b/src/agent_env/providers/sandbox_providers/smol_vm/provider.py index 05f9c8f9..6c2c026d 100644 --- a/src/agent_env/providers/sandbox_providers/smol_vm/provider.py +++ b/src/agent_env/providers/sandbox_providers/smol_vm/provider.py @@ -183,6 +183,16 @@ async def get_sandbox(self, sandbox_id: str) -> SmolVmSandbox: sandbox.network_policy = self.effective_network_policy(None) return sandbox + async def terminate_sandbox(self, sandbox_id: str) -> None: + """Connecting starts a stopped machine, and the SDK deletes only a connected one, so a listed machine is + connected first; one not listed is gone.""" + conn = ConnectOptions(target="local", wait_for_ports=False) + if sandbox_id not in {machine.name for machine in await asyncio.to_thread(Machine.list, conn)}: + logger.info("SmolVM machine %s was already gone at terminate", sandbox_id) + return + machine = await asyncio.to_thread(Machine.connect, sandbox_id, conn) + await asyncio.to_thread(machine.delete) + async def _configure_host_access(self, sandbox: SmolVmSandbox) -> None: code, stdout, stderr = await sandbox.exec_with_output("getent", "ahostsv4", "host.smolvm.internal") if code or not stdout.strip(): diff --git a/src/agent_env/providers/sandbox_providers/tenki/provider.py b/src/agent_env/providers/sandbox_providers/tenki/provider.py index bf00a79d..a52d5ce0 100644 --- a/src/agent_env/providers/sandbox_providers/tenki/provider.py +++ b/src/agent_env/providers/sandbox_providers/tenki/provider.py @@ -18,13 +18,13 @@ from typing import Any, ClassVar, Self from urllib.parse import urlparse -from tenki import AsyncClient +from tenki import AsyncClient, SessionNotFoundError 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 -from agent_env.providers.sandbox_providers.tenki.sandbox import ENDED_STATES, TenkiSandbox, policy_from_info +from agent_env.providers.sandbox_providers.tenki.sandbox import ENDED_STATES, TenkiSandbox, close_session, policy_from_info logger = logging.getLogger(__name__) @@ -413,3 +413,11 @@ async def get_sandbox(self, sandbox_id: str) -> TenkiSandbox: network_policy=policy_from_info(info), command_timeout=_remaining_lifetime(info.timeout_at), ) + + async def terminate_sandbox(self, sandbox_id: str) -> None: + try: + session = await self._client().get(sandbox_id) + except SessionNotFoundError: + logger.info("Tenki session %s was already gone at terminate", sandbox_id) + return + await close_session(session) diff --git a/src/agent_env/providers/sandbox_providers/tenki/sandbox.py b/src/agent_env/providers/sandbox_providers/tenki/sandbox.py index 841997b0..5031cfb5 100644 --- a/src/agent_env/providers/sandbox_providers/tenki/sandbox.py +++ b/src/agent_env/providers/sandbox_providers/tenki/sandbox.py @@ -174,30 +174,13 @@ def __init__( self.network_policy = network_policy async def terminate(self) -> None: - """Terminate the VM. One another handle already terminated, while this handle's cached state still says - running, is not an error.""" - try: - await self._session.close_if_open() - except SessionNotFoundError: - logger.info("Tenki session %s was already gone at terminate", self.sandbox_id) - except (InvalidStateError, SessionTerminatedError): - if await self._current_state() not in ENDED_STATES: - raise - logger.info("Tenki session %s was already terminating at terminate", self.sandbox_id) - await self.release() + await close_session(self._session) async def release(self) -> None: - """Stop the session's data-plane credential renewal, which close_if_open leaves to the caller for a session - that is already ending.""" - reset = getattr(self._session, "_reset_data_plane", None) - if reset is not None: - await reset() + await release_session(self._session) async def _current_state(self) -> str: - try: - return (await self._session.refresh()).state - except SessionNotFoundError: - return "GONE" + return await _session_state(self._session) async def exec(self, *command: str) -> _CompletedProcess: """Run argv to completion as root, so a leading ``sudo`` is dropped. A lost stream is exit -1, which @@ -329,3 +312,32 @@ async def _load_tarballs(self, artifacts: list) -> None: signed_urls = await self._signed_image_urls(artifacts) reachable = [url if self._can_download(url) else None for url in signed_urls] await self._load_docker_images(artifacts, reachable) + + +async def close_session(session: Any) -> None: + """Terminate a session's VM. One another handle already terminated, while this one's cached state still says + running, is not an error.""" + try: + await session.close_if_open() + except SessionNotFoundError: + logger.info("Tenki session %s was already gone at terminate", session.id) + except (InvalidStateError, SessionTerminatedError): + if await _session_state(session) not in ENDED_STATES: + raise + logger.info("Tenki session %s was already terminating at terminate", session.id) + await release_session(session) + + +async def release_session(session: Any) -> None: + """Stop the session's data-plane credential renewal, which close_if_open leaves to the caller for a session + that is already ending.""" + reset = getattr(session, "_reset_data_plane", None) + if reset is not None: + await reset() + + +async def _session_state(session: Any) -> str: + try: + return (await session.refresh()).state + except SessionNotFoundError: + return "GONE" diff --git a/src/agent_env/providers/sandbox_providers/tensorlake/provider.py b/src/agent_env/providers/sandbox_providers/tensorlake/provider.py index 087a97c7..cbd0ed46 100644 --- a/src/agent_env/providers/sandbox_providers/tensorlake/provider.py +++ b/src/agent_env/providers/sandbox_providers/tensorlake/provider.py @@ -336,3 +336,10 @@ async def get_sandbox(self, sandbox_id: str) -> TensorlakeSandbox: "Reconnected to Tensorlake sandbox %s, but image loading will fail closed: %s", sandbox_id, exc, ) return sandbox + + async def terminate_sandbox(self, sandbox_id: str) -> None: + try: + raw_sandbox = await self._sandbox_cls.connect(sandbox_id, **self._client_kwargs()) + await raw_sandbox.terminate() + except SandboxNotFoundError: + logger.info("Tensorlake sandbox %s was already absent during terminate", sandbox_id) diff --git a/src/agent_env/task/teardown.py b/src/agent_env/task/teardown.py index 263e2f0e..f01419fe 100644 --- a/src/agent_env/task/teardown.py +++ b/src/agent_env/task/teardown.py @@ -98,14 +98,14 @@ async def take_down(sandbox: RecordedSandbox) -> None: return try: async with asyncio.timeout(timeout): - handle = await _terminate(sandbox) + on_this_machine = await _terminate(sandbox) except Exception as e: reason = f"timed out after {timeout:g}s" if isinstance(e, TimeoutError) else _described(e) logger.warning("teardown: %s (%s) was not terminated: %s", sandbox.sandbox_id, kind(sandbox), reason) failed.append((sandbox, reason)) return context.metadata[TORN_DOWN_KEY] = [*context.metadata.get(TORN_DOWN_KEY, []), sandbox.sandbox_id] - if handle.ON_THIS_MACHINE and (why := await _remove_folder(sandbox)): + if on_this_machine and (why := await _remove_folder(sandbox)): failed.append((sandbox, why)) else: terminated.append(sandbox) @@ -133,12 +133,17 @@ async def _remove_folder(sandbox: RecordedSandbox) -> str | None: return None -async def _terminate(sandbox: RecordedSandbox): - provider = (build_sandbox_provider(sandbox.sandbox_type) if sandbox.sandbox_type - else _DEFAULT_PROVIDERS[sandbox.slot]()) - handle = await provider.get_sandbox(sandbox.sandbox_id) +async def _terminate(sandbox: RecordedSandbox) -> bool: + """Terminate ``sandbox`` and say whether it ran on this machine: by id through the backend that made it, or, for + a record that doesn't name one, through a live handle, since "not found" from a provider that didn't make the + sandbox doesn't mean it's gone.""" + if sandbox.sandbox_type: + provider = build_sandbox_provider(sandbox.sandbox_type) + await provider.terminate_sandbox(sandbox.sandbox_id) + return provider.ON_THIS_MACHINE + handle = await _DEFAULT_PROVIDERS[sandbox.slot]().get_sandbox(sandbox.sandbox_id) await handle.terminate() - return handle + return handle.ON_THIS_MACHINE def _described(error: Exception) -> str: diff --git a/src/agent_env/task_step/task_steps/teardown_sandboxes.py b/src/agent_env/task_step/task_steps/teardown_sandboxes.py index 7f8f1539..e53ec765 100644 --- a/src/agent_env/task_step/task_steps/teardown_sandboxes.py +++ b/src/agent_env/task_step/task_steps/teardown_sandboxes.py @@ -166,8 +166,10 @@ def _is_backend(name: str) -> bool: async def _terminate(sandbox_id: str, sandbox_type: Optional[str], default_slot: str) -> None: - provider: SandboxProvider = ( - build_sandbox_provider(sandbox_type) if sandbox_type else _DEFAULT_PROVIDERS[default_slot]() - ) - sandbox = await provider.get_sandbox(sandbox_id) - await sandbox.terminate() + """By id through the backend that made it; a record that doesn't name one goes through a live handle, since + "not found" from a provider that didn't make the sandbox doesn't mean it's gone.""" + if sandbox_type: + await build_sandbox_provider(sandbox_type).terminate_sandbox(sandbox_id) + return + provider: SandboxProvider = _DEFAULT_PROVIDERS[default_slot]() + await (await provider.get_sandbox(sandbox_id)).terminate() diff --git a/tst/manual/smol_vm_agent_check.py b/tst/manual/smol_vm_agent_check.py index 9909d467..eda376aa 100644 --- a/tst/manual/smol_vm_agent_check.py +++ b/tst/manual/smol_vm_agent_check.py @@ -13,7 +13,7 @@ from agent_env.task_step.context import TaskStepContext from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep from agent_env.task_step.task_steps.deploy_sandbox import DeploySandboxTaskStep -from tst.integration.store.local_parallel_deploy_test import local_stack # noqa: F401 +from tst.integration.store.local_parallel_deploy_test import local_stack as local_stack from tst.util.a2a_test_agent import build_test_agent_image, put_test_agent diff --git a/tst/manual/smol_vm_gateway_check.py b/tst/manual/smol_vm_gateway_check.py index 7b8730d2..a5534e33 100644 --- a/tst/manual/smol_vm_gateway_check.py +++ b/tst/manual/smol_vm_gateway_check.py @@ -18,7 +18,7 @@ from agent_env.task_step.task_steps.deploy_sandbox import DeploySandboxTaskStep from tst.integration.env.envs.multi_env_local_test import _items_universe from tst.integration.env.envs.server_env_local_test import _put_items_image -from tst.integration.store.local_parallel_deploy_test import _seed_bootstrap_envs, local_stack # noqa: F401 +from tst.integration.store.local_parallel_deploy_test import _seed_bootstrap_envs, local_stack as local_stack from tst.util.a2a_test_agent import put_test_agent diff --git a/tst/unit/bundle/run_teardown_test.py b/tst/unit/bundle/run_teardown_test.py index 916dbaf2..40af715d 100644 --- a/tst/unit/bundle/run_teardown_test.py +++ b/tst/unit/bundle/run_teardown_test.py @@ -22,6 +22,7 @@ from agent_env.env.env import DeployedGatewayEnv from agent_env.task_step.context import DeployedAgent, DeployedSandbox, TaskStepContext from agent_env.task_step.task_step import TaskStep +from agent_env.providers.sandbox_providers.sandbox_provider import SandboxProvider from tst.unit.bundle._support import layout, local_store pytestmark = pytest.mark.usefixtures("sigint_handled") @@ -91,7 +92,7 @@ def __init__(self): def terminated(monkeypatch): log = _Terminated() - class _Provider: + class _Provider(SandboxProvider): async def get_sandbox(self, sandbox_id): return _FakeSandbox(log, sandbox_id) diff --git a/tst/unit/cli/eval_run_teardown_test.py b/tst/unit/cli/eval_run_teardown_test.py index 6a07b939..c9575317 100644 --- a/tst/unit/cli/eval_run_teardown_test.py +++ b/tst/unit/cli/eval_run_teardown_test.py @@ -15,6 +15,7 @@ from agent_env.task import Task, teardown from agent_env.task.teardown import TeardownReport from agent_env.task_step.context import DeployedSandbox +from agent_env.providers.sandbox_providers.sandbox_provider import SandboxProvider eval_run = importlib.import_module("agent_env.cli.eval.run") # the package's ``run`` is the click command pytestmark = pytest.mark.usefixtures("sigint_handled") @@ -81,7 +82,7 @@ async def terminate(self): await asyncio.sleep(0.2) terminated.append(self.sandbox_id) - class _Provider: + class _Provider(SandboxProvider): async def get_sandbox(self, sandbox_id): return _Sandbox(sandbox_id) diff --git a/tst/unit/providers/sandbox_providers/e2b_sandbox_provider_test.py b/tst/unit/providers/sandbox_providers/e2b_sandbox_provider_test.py index 7ec2b44d..1cdfc4fe 100644 --- a/tst/unit/providers/sandbox_providers/e2b_sandbox_provider_test.py +++ b/tst/unit/providers/sandbox_providers/e2b_sandbox_provider_test.py @@ -328,3 +328,17 @@ async def test_adapter_construction_failure_kills_raw_sandbox(): await provider.create_vm(exposed_ports=[8080], setup_for_gateway=False) raw.kill.assert_awaited_once_with() + + +@pytest.mark.parametrize("killed", [True, False], ids=["killed", "already absent"]) +@pytest.mark.asyncio +async def test_terminate_sandbox_kills_by_id_without_connecting(monkeypatch, killed): + kill = AsyncMock(return_value=killed) + monkeypatch.setattr(_AsyncSandboxSdk, "kill", kill, raising=False) + provider = E2BSandboxProvider(api_key="e2b-secret", base_template="agent-env-v1", template_resolver=_Resolver(), + sandbox_cls=_AsyncSandboxSdk) + + await provider.terminate_sandbox("e2b-sandbox-1") + + kill.assert_awaited_once_with("e2b-sandbox-1", api_key="e2b-secret") + _AsyncSandboxSdk.connect.assert_not_called() diff --git a/tst/unit/providers/sandbox_providers/modal_sandbox_test.py b/tst/unit/providers/sandbox_providers/modal_sandbox_test.py index 3b1e9bea..82b0b0d7 100644 --- a/tst/unit/providers/sandbox_providers/modal_sandbox_test.py +++ b/tst/unit/providers/sandbox_providers/modal_sandbox_test.py @@ -5,6 +5,7 @@ import asyncio from unittest.mock import AsyncMock, MagicMock, patch +import modal import pytest from agent_env.config import get_config @@ -437,3 +438,34 @@ async def collect(chunks, destination_path): sandbox._write_stream_via_exec = collect await sandbox.write_file_from_object(url, "/data/a.json") assert written == [(b"LOCAL", "/data/a.json")] + + +@pytest.mark.asyncio +async def test_terminate_sandbox_terminates_by_id_without_reading_its_tunnels(_patched_provider): + sb = MagicMock() + sb.tunnels.aio = AsyncMock(side_effect=modal.exception.ConflictError("Sandbox has already finished with status terminated")) + sb.terminate.aio = AsyncMock() + with patch("agent_env.providers.sandbox_providers.modal_sandbox.modal.Sandbox.from_id") as from_id: + from_id.aio = AsyncMock(return_value=sb) + await _patched_provider.terminate_sandbox("sb-finished") + + sb.terminate.aio.assert_awaited_once() + sb.tunnels.aio.assert_not_called() + sb.exec.aio.assert_not_called() + + +@pytest.mark.asyncio +async def test_terminate_sandbox_counts_a_sandbox_modal_no_longer_knows_as_gone(_patched_provider): + with patch("agent_env.providers.sandbox_providers.modal_sandbox.modal.Sandbox.from_id") as from_id: + from_id.aio = AsyncMock(side_effect=modal.exception.NotFoundError("Sandbox not found")) + await _patched_provider.terminate_sandbox("sb-gone") + + +@pytest.mark.asyncio +async def test_get_sandbox_still_raises_for_a_finished_sandbox_so_a_reattach_sees_it_dead(_patched_provider): + sb = MagicMock() + sb.tunnels.aio = AsyncMock(side_effect=modal.exception.ConflictError("Sandbox has already finished with status terminated")) + with patch("agent_env.providers.sandbox_providers.modal_sandbox.modal.Sandbox.from_id") as from_id: + from_id.aio = AsyncMock(return_value=sb) + with pytest.raises(modal.exception.ConflictError): + await _patched_provider.get_sandbox("sb-finished") diff --git a/tst/unit/providers/sandbox_providers/modal_vm_sandbox_test.py b/tst/unit/providers/sandbox_providers/modal_vm_sandbox_test.py index 514c5e29..2a8ce590 100644 --- a/tst/unit/providers/sandbox_providers/modal_vm_sandbox_test.py +++ b/tst/unit/providers/sandbox_providers/modal_vm_sandbox_test.py @@ -9,6 +9,7 @@ from unittest.mock import AsyncMock, MagicMock, patch +import modal import pytest from agent_env.providers.sandbox_providers.modal_vm_sandbox import ModalVmSandbox, ModalVmSandboxProvider @@ -395,3 +396,19 @@ async def test_apps_are_looked_up_in_the_configured_environment(): lookup.aio = AsyncMock(return_value=MagicMock()) await provider._get_app("agent-env-p") assert lookup.aio.call_args.kwargs["environment_name"] == "rl-prod" + + +@pytest.mark.parametrize("found", [True, False], ids=["live or finished", "gone"]) +@pytest.mark.asyncio +async def test_terminate_sandbox_terminates_a_vm_by_id_and_counts_one_modal_no_longer_knows_as_gone(found): + provider = ModalVmSandboxProvider() + provider._get_client = AsyncMock(return_value=MagicMock()) # type: ignore[method-assign] + sb = MagicMock() + sb.terminate.aio = AsyncMock() + with patch("agent_env.providers.sandbox_providers.modal_sandbox.modal.Sandbox.from_id") as from_id: + from_id.aio = (AsyncMock(return_value=sb) if found + else AsyncMock(side_effect=modal.exception.NotFoundError("Sandbox not found"))) + await provider.terminate_sandbox("sb-vm") + + assert sb.terminate.aio.await_count == (1 if found else 0) + sb.tunnels.aio.assert_not_called() diff --git a/tst/unit/providers/sandbox_providers/sail_vm_sandbox_provider_test.py b/tst/unit/providers/sandbox_providers/sail_vm_sandbox_provider_test.py index bce5ff88..fb89a870 100644 --- a/tst/unit/providers/sandbox_providers/sail_vm_sandbox_provider_test.py +++ b/tst/unit/providers/sandbox_providers/sail_vm_sandbox_provider_test.py @@ -407,3 +407,16 @@ def test_a_rejected_key_can_be_retried(): assert _sdk.API_KEY_ENV not in os.environ _sdk.connect("sail-secret", "agent-env", sdk=sdk) assert sdk.reset_transports.call_count == 2 + + +@pytest.mark.asyncio +async def test_terminate_sandbox_terminates_through_the_handle_and_counts_a_sailbox_gone_as_terminated(): + sailbox = _sailbox() + sdk = _fake_sdk(sailbox) + provider = _provider(sdk) + + await provider.terminate_sandbox("sb_1") + sailbox.terminate.aio.assert_awaited_once() + + sdk.Sailbox.get.aio = AsyncMock(side_effect=sdk.NotFoundError("gone")) + await provider.terminate_sandbox("sb_gone") diff --git a/tst/unit/providers/sandbox_providers/smol_vm_provider_test.py b/tst/unit/providers/sandbox_providers/smol_vm_provider_test.py index 85c0a4a6..274d504d 100644 --- a/tst/unit/providers/sandbox_providers/smol_vm_provider_test.py +++ b/tst/unit/providers/sandbox_providers/smol_vm_provider_test.py @@ -57,7 +57,7 @@ class FakeMachine: id = "agentenv-test" def endpoint(self, port): - return SimpleNamespace(http_url=f"http://127.0.0.1:44000") + return SimpleNamespace(http_url="http://127.0.0.1:44000") def write_file(self, path, data): created.append((path, data)) @@ -298,3 +298,20 @@ async def test_loopback_advertised_grant_uses_object_push(monkeypatch, tmp_path) await sandbox._download_object_to_vm(object_url, "/storage/image.tar.gz") sandbox._write_unsigned_object.assert_awaited_once_with(store, object_url, "/storage/image.tar.gz") + + +@pytest.mark.parametrize("listed", [True, False], ids=["listed", "gone"]) +@pytest.mark.asyncio +async def test_terminate_sandbox_deletes_a_listed_machine_without_setting_it_up(monkeypatch, listed): + deleted = [] + connect = MagicMock(return_value=SimpleNamespace(delete=lambda: deleted.append(True))) + machines = [SimpleNamespace(name="agentenv-vm-1")] if listed else [] + monkeypatch.setattr(module, "Machine", SimpleNamespace(list=lambda conn: machines, connect=connect)) + setup = AsyncMock() + monkeypatch.setattr(module.SmolVmSandboxProvider, "_setup_docker", setup) + + await module.SmolVmSandboxProvider().terminate_sandbox("agentenv-vm-1") + + assert deleted == ([True] if listed else []) + assert connect.called is listed + setup.assert_not_called() diff --git a/tst/unit/providers/sandbox_providers/tenki_sandbox_provider_test.py b/tst/unit/providers/sandbox_providers/tenki_sandbox_provider_test.py index 9830f0b5..8a905ec6 100644 --- a/tst/unit/providers/sandbox_providers/tenki_sandbox_provider_test.py +++ b/tst/unit/providers/sandbox_providers/tenki_sandbox_provider_test.py @@ -21,7 +21,7 @@ vm_shape, ) from agent_env.providers.sandbox_providers.tenki.sandbox import DOCKER_BOOTSTRAP, TenkiSandbox -from tenki import SandboxError +from tenki import SandboxError, SessionNotFoundError from tst.unit.providers.sandbox_providers.tenki_fakes import FakeSession, fake_client_cls, result @@ -383,3 +383,26 @@ async def test_get_sandbox_restores_ports_policy_and_lifetime(client_cls): mode=ALLOWLIST, allow_hosts=("pypi.org",), allow_cidrs=("10.0.0.0/8",), ) assert 1700 <= sandbox.command_timeout <= 1800 + + +@pytest.mark.asyncio +async def test_terminate_sandbox_closes_the_session_by_id(client_cls): + provider = _provider(client_cls) + created = await provider.create_vm(setup_for_gateway=False) + session = client_cls.instances[0].sessions[created.sandbox_id] + + await provider.terminate_sandbox(created.sandbox_id) + + assert (session.closed, session.info.state) == (1, "TERMINATED") + + +@pytest.mark.asyncio +async def test_terminate_sandbox_counts_a_session_tenki_no_longer_knows_as_gone(client_cls, monkeypatch): + provider = _provider(client_cls) + client = provider._client() + + async def gone(sandbox_id): + raise SessionNotFoundError("gone") + + monkeypatch.setattr(client, "get", gone) + await provider.terminate_sandbox("sess-gone") diff --git a/tst/unit/providers/sandbox_providers/tensorlake_sandbox_provider_test.py b/tst/unit/providers/sandbox_providers/tensorlake_sandbox_provider_test.py index 788a49d9..703fd224 100644 --- a/tst/unit/providers/sandbox_providers/tensorlake_sandbox_provider_test.py +++ b/tst/unit/providers/sandbox_providers/tensorlake_sandbox_provider_test.py @@ -451,3 +451,25 @@ async def test_get_sandbox_with_an_unknown_policy_fails_closed(): sandbox = await provider.get_sandbox("tl-test") assert sandbox.network_policy is None + + +@pytest.mark.asyncio +async def test_terminate_sandbox_connects_by_id_and_terminates_without_reading_its_info(): + raw = _raw_sandbox() + provider, sdk = _provider(raw) + + await provider.terminate_sandbox("tl-test") + + sdk.connect.assert_awaited_once() + raw.terminate.assert_awaited_once() + raw.info.assert_not_called() + + +@pytest.mark.parametrize("where", ["connect", "terminate"]) +@pytest.mark.asyncio +async def test_terminate_sandbox_counts_a_sandbox_tensorlake_no_longer_knows_as_gone(where): + raw = _raw_sandbox() + provider, sdk = _provider(raw) + (sdk.connect if where == "connect" else raw.terminate).side_effect = SandboxNotFoundError("tl-test") + + await provider.terminate_sandbox("tl-test") diff --git a/tst/unit/task/teardown_test.py b/tst/unit/task/teardown_test.py index 1c7b5aea..7ceb5c77 100644 --- a/tst/unit/task/teardown_test.py +++ b/tst/unit/task/teardown_test.py @@ -8,9 +8,13 @@ import tempfile import time +from unittest.mock import AsyncMock, MagicMock + +import modal import pytest from agent_env.env.env import DeployedEnv, DeployedGatewayEnv +from agent_env.providers.sandbox_providers.modal_sandbox import ModalSandboxProvider from agent_env.providers.sandbox_providers.local_sandbox import LocalSandbox, remove_local_work_dir from agent_env.task import teardown from agent_env.task.teardown import RecordedSandbox, TeardownReport, recorded_sandboxes, teardown_run @@ -35,6 +39,8 @@ async def terminate(self): class _FakeProvider: + ON_THIS_MACHINE = False + def __init__(self, name, log): self.name, self.log = name, log @@ -43,6 +49,9 @@ async def get_sandbox(self, sandbox_id): raise RuntimeError(f"no sandbox {sandbox_id}") return _FakeSandbox(self, sandbox_id) + async def terminate_sandbox(self, sandbox_id): + await (await self.get_sandbox(sandbox_id)).terminate() + class _Log(list): """The (provider, sandbox_id) terminate log, and how given sandboxes misbehave.""" @@ -144,6 +153,44 @@ def test_a_terminate_that_fails_is_reported_not_raised_and_the_others_still_go(t assert "sb-gw" not in context.metadata[TORN_DOWN_KEY] +def test_a_record_that_names_its_backend_is_terminated_by_id_and_one_that_doesnt_through_a_handle( + terminated, monkeypatch, +): + by_id = [] + + async def terminate_sandbox(self, sandbox_id): + by_id.append((self.name, sandbox_id)) + + monkeypatch.setattr(_FakeProvider, "terminate_sandbox", terminate_sandbox) + context = TaskStepContext(deployed_sandboxes=[ + DeployedSandbox(sandbox_name="named", sandbox_id="sb-named", sandbox_mode="vm", sandbox_type="modal"), + DeployedSandbox(sandbox_name="unnamed", sandbox_id="sb-unnamed", sandbox_mode="vm"), + ]) + + report = _run(context) + + assert by_id == [("modal", "sb-named")] + assert terminated == [("default-sandbox", "sb-unnamed")] + assert {sandbox.sandbox_id for sandbox in report.terminated} == {"sb-named", "sb-unnamed"} + + +def test_a_modal_sandbox_that_already_finished_is_counted_down_not_reported_up(monkeypatch): + sb = MagicMock() + sb.tunnels.aio = AsyncMock(side_effect=modal.exception.ConflictError("Sandbox has already finished with status terminated")) + sb.terminate.aio = AsyncMock() + monkeypatch.setattr(modal.Sandbox, "from_id", MagicMock(aio=AsyncMock(return_value=sb))) + monkeypatch.setattr(ModalSandboxProvider, "_get_client", AsyncMock(return_value=MagicMock())) + context = TaskStepContext(deployed_sandboxes=[ + DeployedSandbox(sandbox_name="box", sandbox_id="sb-finished", sandbox_mode="container", sandbox_type="modal")]) + + report = _run(context) + + assert report.terminated == (RecordedSandbox("sb-finished", "modal", "sandbox"),) + assert report.still_up == () + assert context.metadata[TORN_DOWN_KEY] == ["sb-finished"] + sb.tunnels.aio.assert_not_called() + + def test_cancelling_it_returns_what_it_reached_and_leaves_the_rest(terminated): terminated.behaviors["sb-gw"] = "hang" context = _context() diff --git a/tst/unit/task_step/test_teardown_sandboxes.py b/tst/unit/task_step/test_teardown_sandboxes.py index de32bcda..63617d85 100644 --- a/tst/unit/task_step/test_teardown_sandboxes.py +++ b/tst/unit/task_step/test_teardown_sandboxes.py @@ -34,6 +34,9 @@ def __init__(self, name, terminated, failing): async def get_sandbox(self, sandbox_id): return _FakeSandbox(self, sandbox_id) + async def terminate_sandbox(self, sandbox_id): + await (await self.get_sandbox(sandbox_id)).terminate() + class _Log(list): """The (provider, sandbox_id) terminate log, plus the sandbox ids whose terminate should fail.""" @@ -141,6 +144,20 @@ async def test_an_env_tears_down_every_sandbox_behind_it_on_its_own_backend(term ] +@pytest.mark.asyncio +async def test_a_typed_sandbox_is_terminated_by_id_without_a_handle(terminated, monkeypatch): + async def no_handle(self, sandbox_id): + raise AssertionError(f"asked {self.name} for a handle to {sandbox_id}") + + async def by_id(self, sandbox_id): + self.terminated.append((self.name, sandbox_id)) + + monkeypatch.setattr(_FakeProvider, "get_sandbox", no_handle) + monkeypatch.setattr(_FakeProvider, "terminate_sandbox", by_id) + await _step(agent_names=["solver"]).execute(_context()) + assert terminated == [("modal", "sb-solver")] + + @pytest.mark.asyncio async def test_an_untyped_sandbox_uses_its_slot_default(terminated): await _step(sandbox_names=["gpu-box"]).execute(_context())