Fix Runloop rollout setup cleanup

Author: tode-rl

.env.example

@@ -22,5 +22,8 @@ FIREWORKS_API_KEY="your_fireworks_api_key_here"
 # Runloop API Key (if hosting remote rollout servers in Runloop Devboxes)
 # RUNLOOP_API_KEY="your_runloop_api_key_here"
 
+# Optional: Runloop blueprint used by examples/runloop_remote_rollout/test_eval.py
+# RUNLOOP_BLUEPRINT_ID="your_runloop_blueprint_id_here"
+
 # Other environment variables your custom reward functions might need
 # MY_CUSTOM_SERVICE_API_KEY="some_other_key"

docs/integrations/runloop_remote_rollout.mdx

@@ -39,6 +39,19 @@ async def test_my_eval(row):
 
 The server command must bind to `0.0.0.0` on the configured port so the Runloop tunnel can reach it. The server must expose `POST /init` and should use `FireworksTracingHttpHandler` plus `RolloutIdFilter` to publish rollout completion status.
 
+## Creating A Blueprint
+
+`blueprint_id` is required when you want `RunloopRolloutProcessor` to create a fresh Devbox for each eval invocation. The blueprint should contain the rollout server code and its Python dependencies.
+
+The included example can create a blueprint for a new Runloop account:
+
+```bash
+export RUNLOOP_API_KEY=...
+eval "$(python examples/runloop_remote_rollout/create_blueprint.py)"
+```
+
+That helper uploads the current repository as a temporary Runloop build context and builds a Python image with `eval-protocol[runloop]` installed. Use the printed `RUNLOOP_BLUEPRINT_ID` with `examples/runloop_remote_rollout/test_eval.py`.
+
 ## Existing Devboxes
 
 You can attach to an existing Devbox instead of creating one from a blueprint:

eval_protocol/pytest/runloop_rollout_processor.py

@@ -99,28 +99,34 @@ def setup(self) -> None:
             )
 
         RunloopSDK = _load_runloop_sdk()
-        self._client = RunloopSDK(bearer_token=api_key)
-
-        if self._devbox_id:
-            self._devbox = self._client.devbox.from_id(self._devbox_id)
-            self._owns_devbox = False
-        else:
-            assert self._blueprint_id is not None
-            self._devbox = self._client.devbox.create_from_blueprint_id(self._blueprint_id)
-            self._owns_devbox = True
-
-        self._await_running()
-        tunnel = self._create_tunnel()
-        self._remote_base_url = self._derive_remote_base_url(tunnel)
-        self._server_execution = self._devbox.cmd.exec_async(self._server_command)
-        self._wait_for_server_startup()
-        self._remote_processor = RemoteRolloutProcessor(
-            remote_base_url=self._remote_base_url,
-            model_base_url=self._model_base_url,
-            poll_interval=self._poll_interval,
-            timeout_seconds=self._timeout_seconds,
-            include_payloads=self._include_payloads,
-        )
+        client: Any = RunloopSDK(bearer_token=api_key)
+        self._client = client
+
+        try:
+            if self._devbox_id:
+                devbox = client.devbox.from_id(self._devbox_id)
+                self._owns_devbox = False
+            else:
+                assert self._blueprint_id is not None
+                devbox = client.devbox.create_from_blueprint_id(self._blueprint_id)
+                self._owns_devbox = True
+
+            self._devbox = devbox
+            self._await_running()
+            tunnel = self._create_tunnel()
+            self._remote_base_url = self._derive_remote_base_url(tunnel)
+            self._server_execution = devbox.cmd.exec_async(self._server_command)
+            self._wait_for_server_startup()
+            self._remote_processor = RemoteRolloutProcessor(
+                remote_base_url=self._remote_base_url,
+                model_base_url=self._model_base_url,
+                poll_interval=self._poll_interval,
+                timeout_seconds=self._timeout_seconds,
+                include_payloads=self._include_payloads,
+            )
+        except Exception:
+            self._cleanup_partial_setup()
+            raise
 
     def __call__(self, rows: list[EvaluationRow], config: RolloutProcessorConfig) -> list[asyncio.Task[EvaluationRow]]:
         if self._remote_processor is None:
@@ -149,6 +155,7 @@ def _await_running(self) -> None:
         await_running()
 
     def _create_tunnel(self) -> Any:
+        assert self._devbox is not None
         net = self._devbox.net
         create_tunnel = getattr(net, "create_tunnel", None)
         if create_tunnel is not None:
@@ -190,8 +197,11 @@ def _wait_for_server_startup(self) -> None:
                 with urllib.request.urlopen(request, timeout=min(5.0, self._startup_timeout_seconds)) as response:
                     response.read(1)
                     return
-            except urllib.error.HTTPError:
-                return
+            except urllib.error.HTTPError as exc:
+                if exc.code < 500:
+                    return
+                last_error = exc
+                time.sleep(min(1.0, max(0.0, deadline - time.monotonic())))
             except Exception as exc:
                 last_error = exc
                 time.sleep(min(1.0, max(0.0, deadline - time.monotonic())))
@@ -214,3 +224,15 @@ def _shutdown_devbox(self) -> None:
             return
         self._devbox.shutdown()
         self._shutdown_complete = True
+
+    def _cleanup_partial_setup(self) -> None:
+        if self._remote_processor is not None:
+            self._remote_processor.cleanup()
+            self._remote_processor = None
+        if self._should_shutdown_devbox():
+            self._shutdown_devbox()
+        self._devbox = None
+        self._server_execution = None
+        self._remote_base_url = None
+        self._owns_devbox = False
+        self._shutdown_complete = False

examples/runloop_remote_rollout/README.md

@@ -13,10 +13,12 @@ export FIREWORKS_API_KEY=...
 Create a Runloop blueprint that contains this repository and its Python dependencies, then set `RUNLOOP_BLUEPRINT_ID`:
 
 ```bash
-export RUNLOOP_BLUEPRINT_ID=bpt_your_blueprint_id
+eval "$(python examples/runloop_remote_rollout/create_blueprint.py)"
 pytest examples/runloop_remote_rollout/test_eval.py
 ```
 
+The blueprint ID matters because `RunloopRolloutProcessor` uses it to create a Devbox that already has this repository and `eval-protocol[runloop]` installed. If you already have a suitable running Devbox, you can pass `devbox_id` to `RunloopRolloutProcessor` instead and skip `RUNLOOP_BLUEPRINT_ID`.
+
 The processor starts:
 
 ```bash

examples/runloop_remote_rollout/create_blueprint.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+import argparse
+import tarfile
+from datetime import timedelta
+from pathlib import Path
+
+from runloop_api_client import RunloopSDK
+from runloop_api_client.types.blueprint_build_parameters import BuildContext
+
+
+DEFAULT_DOCKERFILE = """\
+FROM python:3.12-slim
+WORKDIR /workspace
+COPY . /workspace
+RUN pip install --no-cache-dir ".[runloop]"
+"""
+
+IGNORED_CONTEXT_DIRS = {
+    ".git",
+    ".mypy_cache",
+    ".pytest_cache",
+    ".ruff_cache",
+    ".tox",
+    ".venv",
+    "__pycache__",
+    "node_modules",
+}
+
+
+def _ignore_build_context(member: tarfile.TarInfo) -> tarfile.TarInfo | None:
+    parts = set(Path(member.name).parts)
+    if parts & IGNORED_CONTEXT_DIRS:
+        return None
+    return member
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Create a Runloop blueprint for the remote rollout example.")
+    parser.add_argument(
+        "--repo-root",
+        type=Path,
+        default=Path(__file__).resolve().parents[2],
+        help="Path to the eval-protocol repository root.",
+    )
+    parser.add_argument(
+        "--name",
+        default="eval-protocol-runloop-remote-rollout",
+        help="Runloop blueprint name.",
+    )
+    args = parser.parse_args()
+
+    runloop = RunloopSDK()
+    build_context = runloop.storage_object.upload_from_dir(
+        args.repo_root,
+        name=f"{args.name}.tar.gz",
+        ttl=timedelta(hours=1),
+        ignore=_ignore_build_context,
+    )
+    blueprint = runloop.blueprint.create(
+        name=args.name,
+        dockerfile=DEFAULT_DOCKERFILE,
+        build_context=BuildContext(type="object", object_id=build_context.id),
+    )
+
+    print(f"export RUNLOOP_BLUEPRINT_ID={blueprint.id}")
+
+
+if __name__ == "__main__":
+    main()

tests/pytest/test_runloop_rollout_processor.py

@@ -1,5 +1,6 @@
 import asyncio
 from types import SimpleNamespace
+import urllib.error
 
 import pytest
 
@@ -211,6 +212,99 @@ def _raise_missing_dependency():
         processor.setup()
 
 
+def test_setup_cleans_up_owned_devbox_after_startup_failure(fake_runloop, monkeypatch):
+    processor = RunloopRolloutProcessor(
+        blueprint_id="bp-123",
+        server_command="python server.py",
+        startup_timeout_seconds=0,
+    )
+
+    def _fail_startup():
+        raise TimeoutError("server did not start")
+
+    monkeypatch.setattr(processor, "_wait_for_server_startup", _fail_startup)
+
+    with pytest.raises(TimeoutError, match="server did not start"):
+        processor.setup()
+
+    assert fake_runloop.shutdown_calls == ["devbox-created"]
+    assert processor.remote_base_url is None
+    assert FakeRemoteRolloutProcessor.instances == []
+
+
+def test_setup_does_not_shutdown_existing_devbox_after_startup_failure(fake_runloop, monkeypatch):
+    processor = RunloopRolloutProcessor(
+        devbox_id="devbox-existing",
+        server_command="python server.py",
+        startup_timeout_seconds=0,
+    )
+
+    def _fail_startup():
+        raise TimeoutError("server did not start")
+
+    monkeypatch.setattr(processor, "_wait_for_server_startup", _fail_startup)
+
+    with pytest.raises(TimeoutError, match="server did not start"):
+        processor.setup()
+
+    assert fake_runloop.shutdown_calls == []
+
+
+def test_startup_wait_retries_5xx_http_errors(monkeypatch):
+    processor = RunloopRolloutProcessor(
+        devbox_id="devbox-existing",
+        server_command="python server.py",
+        startup_timeout_seconds=5,
+    )
+    processor._remote_base_url = "https://8000-test-tunnel-key.tunnel.runloop.ai"
+
+    calls = []
+
+    class ReadyResponse:
+        def __enter__(self):
+            return self
+
+        def __exit__(self, exc_type, exc, traceback):
+            return None
+
+        def read(self, size):
+            return b"o"
+
+    def _urlopen(request, timeout):
+        calls.append((request.full_url, timeout))
+        if len(calls) == 1:
+            raise urllib.error.HTTPError(request.full_url, 503, "Service Unavailable", hdrs=None, fp=None)
+        return ReadyResponse()
+
+    monkeypatch.setattr(runloop_rollout_processor_module.urllib.request, "urlopen", _urlopen)
+    monkeypatch.setattr(runloop_rollout_processor_module.time, "sleep", lambda seconds: None)
+
+    processor._wait_for_server_startup()
+
+    assert len(calls) == 2
+
+
+def test_startup_wait_accepts_non_5xx_http_errors(monkeypatch):
+    processor = RunloopRolloutProcessor(
+        devbox_id="devbox-existing",
+        server_command="python server.py",
+        startup_timeout_seconds=5,
+    )
+    processor._remote_base_url = "https://8000-test-tunnel-key.tunnel.runloop.ai"
+
+    calls = []
+
+    def _urlopen(request, timeout):
+        calls.append((request.full_url, timeout))
+        raise urllib.error.HTTPError(request.full_url, 404, "Not Found", hdrs=None, fp=None)
+
+    monkeypatch.setattr(runloop_rollout_processor_module.urllib.request, "urlopen", _urlopen)
+
+    processor._wait_for_server_startup()
+
+    assert len(calls) == 1
+
+
 @pytest.mark.asyncio
 async def test_async_cleanup_closes_remote_processor_and_owned_devbox(fake_runloop):
     processor = RunloopRolloutProcessor(