Run setup before the agent and skip verify on a non-zero agent exit

Addresses review: run_rollout ran setup after create() had already started the
agent, so setup raced the agent over the workspace. create() now takes
start_agent (default True, so the loop-owning training path is unchanged), the
server passes start_agent=False, runs setup, then calls session.start_agent(),
giving the correct setup then agent then verify order. A non-zero agent exit is
now treated like a timeout, so a crashed agent no longer runs verify or earns a
reward.

AI-assisted.

Author: sergiopaniego

envs/pi_env/harness.py

@@ -225,12 +225,16 @@ def create(
         task: Any,
         seed: int | None = None,
         episode_id: str | None = None,
+        start_agent: bool = True,
     ) -> 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.
+
+        Pass ``start_agent=False`` to return before launching the agent (for
+        example to run setup first), then call ``session.start_agent()``.
         """
         import logging
         import time
@@ -239,7 +243,9 @@ def create(
         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)
+                return self._create_once(
+                    task, seed=seed, episode_id=episode_id, start_agent=start_agent
+                )
             except Exception as exc:  # noqa: BLE001
                 last_exc = exc
                 if i + 1 < self._create_attempts:
@@ -256,6 +262,7 @@ def _create_once(
         task: Any,
         seed: int | None = None,
         episode_id: str | None = None,
+        start_agent: bool = True,
     ) -> PiSession:
         import logging
         _log = logging.getLogger(__name__)
@@ -312,7 +319,8 @@ def _create_once(
                 proxy_trace_path=proxy_trace_path_str,
                 proxy_bg_job=proxy_bg_job,
             )
-            session.start_agent()
+            if start_agent:
+                session.start_agent()
             return session
         except Exception as exc:
             _log.error("factory.create: setup failed, killing sandbox: %r", exc)

envs/pi_env/server/pi_environment.py

@@ -331,20 +331,15 @@ def _emit(msg: str) -> None:
                 f"creating HF sandbox (image={image or _DEFAULT_IMAGE}) — "
                 "this is the slow phase (Node + Pi cold-install per rollout)"
             )
-            session = factory.create(task=pi_task)
+            session = factory.create(task=pi_task, start_agent=False)
             result.sandbox_id = session.sandbox.sandbox_id
             _emit(
-                f"sandbox ready: {result.sandbox_id} — agent started "
+                f"sandbox ready: {result.sandbox_id} "
                 f"({'proxy on :7000, logprobs capturing' if mode == 'transparent_proxy' else 'direct LLM, no logprobs'})"
             )
 
-            # Run setup commands one at a time, *before* the agent starts.
-            # The factory has already started the agent in start_agent()
-            # during create(); to keep the order "setup → agent → verify"
-            # we'd need to restructure. As a pragmatic compromise we run
-            # setup IMMEDIATELY after create(), which races with the agent
-            # for ~1-2s but is fine for typical pip/git/download work
-            # because pi itself takes >=20s to make its first model call.
+            # Run setup one command at a time, before the agent starts. create()
+            # was called with start_agent=False, so the agent has not launched yet.
             for i, cmd in enumerate(setup, 1):
                 _emit(f"setup [{i}/{len(setup)}]: {cmd[:80]}")
                 cr = self._exec_command(session.sandbox, cmd, timeout=SETUP_TIMEOUT_S)
@@ -356,8 +351,9 @@ def _emit(msg: str) -> None:
                     _emit(f"setup FAILED at [{i}]: exit={cr.exit_code}")
                     break
 
-            # Block until the agent is done (or setup already failed).
+            # Setup is done and clean, so launch the agent now (deferred from create()).
             if result.error is None:
+                session.start_agent()
                 _emit(
                     f"agent running — pi CLI in sandbox "
                     f"(timeout {int(agent_timeout_s)}s)"
@@ -367,6 +363,9 @@ def _emit(msg: str) -> None:
                         timeout_s=agent_timeout_s
                     )
                     _emit(f"agent finished: exit_code={result.agent_exit_code}")
+                    if result.agent_exit_code != 0:
+                        result.error = f"agent exited non-zero ({result.agent_exit_code})"
+                        _emit(f"agent FAILED: exit_code={result.agent_exit_code}")
                 except TimeoutError as exc:
                     result.error = f"agent timeout: {exc}"
                     _emit(f"agent TIMEOUT: {exc}")

tests/envs/test_pi_factory_lifecycle.py

@@ -80,7 +80,7 @@ def test_create_retries_then_succeeds(monkeypatch):
     calls = {"n": 0}
     session = object()
 
-    def _flaky(task, seed=None, episode_id=None):
+    def _flaky(task, seed=None, episode_id=None, start_agent=True):
         calls["n"] += 1
         if calls["n"] < 3:
             raise RuntimeError("transient create failure")
@@ -97,7 +97,7 @@ def test_create_raises_after_exhausting_attempts(monkeypatch):
     factory = _factory(sandbox, create_attempts=3, create_backoff_s=0)
     calls = {"n": 0}
 
-    def _always_fails(task, seed=None, episode_id=None):
+    def _always_fails(task, seed=None, episode_id=None, start_agent=True):
         calls["n"] += 1
         raise RuntimeError("persistent create failure")