diff --git a/src/olmo_eval/harness/sandbox/executor.py b/src/olmo_eval/harness/sandbox/executor.py index 7d79ca09f..d32d5b530 100644 --- a/src/olmo_eval/harness/sandbox/executor.py +++ b/src/olmo_eval/harness/sandbox/executor.py @@ -70,6 +70,85 @@ def _get_modal_app_name() -> str: return f"swerex-{uuid.uuid4().hex[:12]}" +def _modal_kwargs_with_encrypted_runtime_port( + modal_kwargs: dict[str, Any] | None, + runtime_port: int, +) -> dict[str, Any]: + """Return Modal sandbox kwargs that expose the SWE-ReX runtime over HTTPS.""" + resolved_kwargs = dict(modal_kwargs or {}) + resolved_kwargs.pop("unencrypted_ports", None) + + configured_ports = resolved_kwargs.pop("encrypted_ports", ()) or () + if isinstance(configured_ports, int): + encrypted_ports = [configured_ports] + else: + encrypted_ports = list(configured_ports) + + if runtime_port not in encrypted_ports: + encrypted_ports.append(runtime_port) + + resolved_kwargs["encrypted_ports"] = encrypted_ports + return resolved_kwargs + + +def _encrypted_modal_deployment_class(modal_deployment_cls: type[Any]) -> type[Any]: + class EncryptedModalDeployment(modal_deployment_cls): + async def start(self) -> None: + """Start the Modal runtime using encrypted sandbox tunnels.""" + from importlib import import_module + + modal = import_module("modal") + RemoteRuntime = import_module("swerex.runtime.remote").RemoteRuntime + + if self._runtime is not None and self._sandbox is not None: + self.logger.warning( + "Deployment is already started. Ignoring duplicate start() call." + ) + return + + self.logger.info("Starting modal sandbox") + self._hooks.on_custom_step("Starting modal sandbox") + t0 = time.time() + token = self._get_token() + modal_kwargs = _modal_kwargs_with_encrypted_runtime_port( + self._modal_kwargs, + self._port, + ) + self._sandbox = await modal.Sandbox.create.aio( + "/usr/bin/env", + "bash", + "-c", + self._start_swerex_cmd(token), + image=self._image, + timeout=int(self._deployment_timeout), + app=self._app, + **modal_kwargs, + ) + tunnels = await self._sandbox.tunnels.aio() + tunnel = tunnels[self._port] + elapsed_sandbox_creation = time.time() - t0 + self.logger.info( + f"Sandbox ({self._sandbox.object_id}) created in {elapsed_sandbox_creation:.2f}s" + ) + self.logger.info(f"Check sandbox logs at {await self.get_modal_log_url()}") + self.logger.info(f"Sandbox created with id {self._sandbox.object_id}") + await asyncio.sleep(1) + self.logger.info(f"Starting runtime at {tunnel.url}") + self._hooks.on_custom_step("Starting runtime") + self._runtime = RemoteRuntime( + host=tunnel.url, + timeout=self._runtime_timeout, + auth_token=token, + logger=self.logger, + ) + remaining_startup_timeout = max(0, self._startup_timeout - elapsed_sandbox_creation) + t1 = time.time() + await self._wait_until_alive(timeout=remaining_startup_timeout) + self.logger.info(f"Runtime started in {time.time() - t1:.2f}s") + + return EncryptedModalDeployment + + class SandboxExecutor: """Executor for sandboxed command execution via SWE-ReX. @@ -340,7 +419,8 @@ def get_deployment(self) -> Any: else: modal_image = modal.Image.from_registry(image) - return ModalDeployment( + EncryptedModalDeployment = _encrypted_modal_deployment_class(ModalDeployment) + return EncryptedModalDeployment( image=modal_image, startup_timeout=self.config.startup_timeout, runtime_timeout=self.config.runtime_timeout, diff --git a/tests/core/harness/test_sandbox_executor.py b/tests/core/harness/test_sandbox_executor.py new file mode 100644 index 000000000..b000ac5f4 --- /dev/null +++ b/tests/core/harness/test_sandbox_executor.py @@ -0,0 +1,163 @@ +"""Tests for sandbox executor deployment creation.""" + +from __future__ import annotations + +import logging +import sys +import types +from typing import Any + +import pytest + +from olmo_eval.harness.sandbox import SandboxConfig, SandboxExecutor, SandboxMode + + +class _AioMethod: + def __init__(self, result: Any) -> None: + self._result = result + + async def aio(self, *args: Any, **kwargs: Any) -> Any: + if callable(self._result): + return self._result(*args, **kwargs) + return self._result + + +class _FakeHooks: + def __init__(self) -> None: + self.steps: list[str] = [] + + def on_custom_step(self, step: str) -> None: + self.steps.append(step) + + +class _FakeModalDeployment: + def __init__( + self, + *, + image: Any, + logger: logging.Logger | None = None, + startup_timeout: float = 0.4, + runtime_timeout: float = 3600.0, + modal_sandbox_kwargs: dict[str, Any] | None = None, + install_pipx: bool = True, + deployment_timeout: float = 3600.0, + ) -> None: + self.logger = logger or logging.getLogger(__name__) + self._image = image + self._startup_timeout = startup_timeout + self._runtime_timeout = runtime_timeout + self._modal_kwargs = modal_sandbox_kwargs or {} + self._deployment_timeout = deployment_timeout + self._runtime = None + self._sandbox = None + self._port = 8000 + self._app = object() + self._hooks = _FakeHooks() + self.wait_until_alive_timeout: float | None = None + + @property + def runtime(self) -> Any: + return self._runtime + + def _get_token(self) -> str: + return "runtime-token" + + def _start_swerex_cmd(self, token: str) -> str: + return f"start-runtime {token}" + + async def get_modal_log_url(self) -> str: + return "https://modal.example/logs" + + async def _wait_until_alive(self, *, timeout: float) -> None: + self.wait_until_alive_timeout = timeout + + +class _FakeImage: + @staticmethod + def from_registry(image: str) -> str: + return f"registry:{image}" + + +class _FakeTunnel: + url = "https://runtime.example" + + +class _FakeSandbox: + object_id = "sandbox-123" + + def __init__(self) -> None: + self.tunnels = _AioMethod({8000: _FakeTunnel()}) + + +class _FakeSandboxCreate: + def __init__(self) -> None: + self.calls: list[tuple[tuple[Any, ...], dict[str, Any]]] = [] + + async def aio(self, *args: Any, **kwargs: Any) -> _FakeSandbox: + self.calls.append((args, kwargs)) + return _FakeSandbox() + + +class _FakeSandboxApi: + create = _FakeSandboxCreate() + + +class _FakeRemoteRuntime: + def __init__(self, **kwargs: Any) -> None: + self.kwargs = kwargs + + +def _install_modal_fakes(monkeypatch: pytest.MonkeyPatch) -> types.ModuleType: + fake_modal = types.ModuleType("modal") + fake_modal.Image = _FakeImage + fake_modal.Sandbox = _FakeSandboxApi + + fake_swerex = types.ModuleType("swerex") + fake_deployment = types.ModuleType("swerex.deployment") + fake_deployment_modal = types.ModuleType("swerex.deployment.modal") + fake_deployment_modal.ModalDeployment = _FakeModalDeployment + fake_runtime = types.ModuleType("swerex.runtime") + fake_runtime_remote = types.ModuleType("swerex.runtime.remote") + fake_runtime_remote.RemoteRuntime = _FakeRemoteRuntime + + monkeypatch.setitem(sys.modules, "modal", fake_modal) + monkeypatch.setitem(sys.modules, "swerex", fake_swerex) + monkeypatch.setitem(sys.modules, "swerex.deployment", fake_deployment) + monkeypatch.setitem(sys.modules, "swerex.deployment.modal", fake_deployment_modal) + monkeypatch.setitem(sys.modules, "swerex.runtime", fake_runtime) + monkeypatch.setitem(sys.modules, "swerex.runtime.remote", fake_runtime_remote) + + return fake_modal + + +@pytest.mark.anyio +async def test_modal_deployment_uses_encrypted_runtime_port( + monkeypatch: pytest.MonkeyPatch, +) -> None: + fake_modal = _install_modal_fakes(monkeypatch) + fake_modal.Sandbox.create.calls.clear() + executor = SandboxExecutor( + SandboxConfig( + image="python:3.12", + mode=SandboxMode.MODAL, + startup_timeout=10.0, + runtime_timeout=20.0, + modal_sandbox_kwargs={ + "encrypted_ports": [1234], + "unencrypted_ports": [5678], + "cloud": "aws", + }, + ) + ) + + deployment = executor.get_deployment() + await deployment.start() + + assert len(fake_modal.Sandbox.create.calls) == 1 + _, kwargs = fake_modal.Sandbox.create.calls[0] + assert "unencrypted_ports" not in kwargs + assert kwargs["encrypted_ports"] == [1234, 8000] + assert kwargs["cloud"] == "aws" + assert kwargs["image"] == "registry:python:3.12" + assert deployment.runtime.kwargs["host"] == "https://runtime.example" + assert deployment.runtime.kwargs["auth_token"] == "runtime-token"