Retry PiSessionFactory.create() with backoff

Mirror of the OpenCodeSessionFactory retry: session creation (sandbox + Pi
install + proxy + agent) is the flakiest step, and a single transient failure
dropped the rollout as unscorable. create() now retries create_attempts times
with exponential backoff; the per-attempt logic moved to _create_once, which
already tears its own sandbox down on failure, so a retry never leaks.

Author: sergiopaniego

envs/pi_env/harness.py

@@ -206,6 +206,8 @@ def __init__(
         verifier: Verifier | None = None,
         install_timeout_s: int = 300,
         setup_timeout_s: int = 300,
+        create_attempts: int = 3,
+        create_backoff_s: float = 2.0,
     ) -> None:
         if mode not in {"black_box", "transparent_proxy"}:
             raise ValueError(f"Unknown mode: {mode!r}")
@@ -215,12 +217,45 @@ def __init__(
         self._verifier = verifier
         self._install_timeout_s = install_timeout_s
         self._setup_timeout_s = setup_timeout_s
+        self._create_attempts = create_attempts
+        self._create_backoff_s = create_backoff_s
 
     def create(
         self,
         task: Any,
         seed: int | None = None,
         episode_id: str | None = None,
+    ) -> PiSession:
+        """Create one session, retrying with exponential backoff.
+
+        Session creation spins up a sandbox, installs Pi, and starts the proxy +
+        agent, the flakiest step in a rollout. Each failed attempt tears its own
+        sandbox down (see :meth:`_create_once`), so a retry never leaks.
+        """
+        import logging
+        import time
+
+        _log = logging.getLogger(__name__)
+        last_exc: Exception = RuntimeError("create_attempts must be >= 1")
+        for i in range(self._create_attempts):
+            try:
+                return self._create_once(task, seed=seed, episode_id=episode_id)
+            except Exception as exc:  # noqa: BLE001
+                last_exc = exc
+                if i + 1 < self._create_attempts:
+                    backoff = self._create_backoff_s * (2**i)
+                    _log.warning(
+                        "factory.create attempt %d/%d failed (%r); retrying in %.1fs",
+                        i + 1, self._create_attempts, exc, backoff,
+                    )
+                    time.sleep(backoff)
+        raise last_exc
+
+    def _create_once(
+        self,
+        task: Any,
+        seed: int | None = None,
+        episode_id: str | None = None,
     ) -> PiSession:
         import logging
         _log = logging.getLogger(__name__)

tests/envs/test_pi_factory_lifecycle.py

@@ -35,16 +35,17 @@ def create(self, **kwargs):
         return self._sandbox
 
 
-def _factory(sandbox):
+def _factory(sandbox, **overrides):
     return PiSessionFactory(
         config=PiConfig(base_url="http://localhost:8000/v1"),
         sandbox_backend=_FakeBackend(sandbox),
+        **overrides,
     )
 
 
 def test_create_kills_sandbox_when_bootstrap_fails(monkeypatch):
     sandbox = _FakeSandbox()
-    factory = _factory(sandbox)
+    factory = _factory(sandbox, create_attempts=1)
     monkeypatch.setattr(
         factory,
         "_bootstrap_sandbox",
@@ -61,7 +62,7 @@ def kill(self):
             raise RuntimeError("kill failed")
 
     sandbox = _KillRaises()
-    factory = _factory(sandbox)
+    factory = _factory(sandbox, create_attempts=1)
     monkeypatch.setattr(
         factory,
         "_bootstrap_sandbox",
@@ -70,3 +71,37 @@ def kill(self):
     # The original bootstrap failure must surface, not the cleanup error.
     with pytest.raises(RuntimeError, match="boom"):
         factory.create("write a function")
+
+
+def test_create_retries_then_succeeds(monkeypatch):
+    # create() retries a flaky _create_once and returns once it succeeds.
+    sandbox = _FakeSandbox()
+    factory = _factory(sandbox, create_attempts=3, create_backoff_s=0)
+    calls = {"n": 0}
+    session = object()
+
+    def _flaky(task, seed=None, episode_id=None):
+        calls["n"] += 1
+        if calls["n"] < 3:
+            raise RuntimeError("transient create failure")
+        return session
+
+    monkeypatch.setattr(factory, "_create_once", _flaky)
+    assert factory.create("write a function") is session
+    assert calls["n"] == 3
+
+
+def test_create_raises_after_exhausting_attempts(monkeypatch):
+    # A persistent failure is re-raised after create_attempts tries.
+    sandbox = _FakeSandbox()
+    factory = _factory(sandbox, create_attempts=3, create_backoff_s=0)
+    calls = {"n": 0}
+
+    def _always_fails(task, seed=None, episode_id=None):
+        calls["n"] += 1
+        raise RuntimeError("persistent create failure")
+
+    monkeypatch.setattr(factory, "_create_once", _always_fails)
+    with pytest.raises(RuntimeError, match="persistent create failure"):
+        factory.create("write a function")
+    assert calls["n"] == 3