Use encrypted Modal sandbox tunnels

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,

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"