fix: poll rollout status asynchronously

Use the lightweight status endpoint from RemoteRolloutProcessor via the shared aiohttp session and avoid the logs backfill after terminal status is observed.

Made-with: Cursor

Author: xzrderek

eval_protocol/adapters/fireworks_tracing.py

@@ -375,31 +375,30 @@ async def async_search_logs(
             )
         return results
 
-    def get_status(self, rollout_id: str) -> Optional[Dict[str, Any]]:
+    async def async_get_status(self, session: aiohttp.ClientSession, rollout_id: str) -> Optional[Dict[str, Any]]:
         """Fetch rollout status from the lightweight /status endpoint.
 
         Returns the parsed JSON response or None if the status is not yet available.
         Response shape: {"rollout_id": "...", "status": {"code": ...} | null}
         """
-        from ..common_utils import get_user_agent
-
         headers = {
             "Authorization": f"Bearer {self._get_api_key()}",
             "User-Agent": get_user_agent(),
         }
         params: Dict[str, Any] = {"rollout_id": rollout_id}
+        timeout = aiohttp.ClientTimeout(total=self.timeout)
 
-        urls_to_try = [f"{self.base_url}/status", f"{self.base_url}/v1/status"]
+        urls_to_try = [f"{self.base_url}/v1/status", f"{self.base_url}/status"]
         last_error: Optional[str] = None
         for url in urls_to_try:
             try:
-                response = requests.get(url, params=params, timeout=self.timeout, headers=headers)
-                if response.status_code == 404:
-                    last_error = f"404 for {url}"
-                    continue
-                response.raise_for_status()
-                return response.json()
-            except requests.exceptions.RequestException as e:
+                async with session.get(url, params=params, headers=headers, timeout=timeout) as resp:
+                    if resp.status == 404:
+                        last_error = f"404 for {url}"
+                        continue
+                    resp.raise_for_status()
+                    return (await resp.json(content_type=None)) or {}
+            except (aiohttp.ClientError, asyncio.TimeoutError, json.JSONDecodeError) as e:
                 last_error = str(e)
                 continue
 

eval_protocol/pytest/remote_rollout_processor.py

@@ -121,43 +121,23 @@ async def _process_row(row: EvaluationRow) -> EvaluationRow:
             deadline = time.time() + timeout_seconds
 
             while time.time() < deadline:
-                # Poll status (run in thread to avoid blocking event loop)
-                status_result = await asyncio.to_thread(
-                    self._tracing_adapter.get_status, rollout_id=row.execution_metadata.rollout_id
+                session = self._get_or_create_session()
+                status_result = await self._tracing_adapter.async_get_status(
+                    session,
+                    rollout_id=row.execution_metadata.rollout_id,
                 )
                 status = (status_result or {}).get("status")
-                if status and "code" in status:
+                if isinstance(status, dict) and "code" in status:
                     status_code = status["code"]
 
-                    if status_code == Status.Code.RUNNING:
-                        await asyncio.sleep(poll_interval)
-                        continue
-
                     logger.info(
                         "Found status for rollout %s with code %s",
                         row.execution_metadata.rollout_id,
                         status_code,
                     )
 
-                    # Backfill message/details/extras from the full Logs table (one-shot)
-                    session = self._get_or_create_session()
-                    completed_logs = await self._tracing_adapter.async_search_logs(
-                        session, tags=[f"rollout_id:{row.execution_metadata.rollout_id}"],
-                    )
-                    status_message = ""
-                    status_details: list = []
-                    status_extras: dict = {}
-                    for log in completed_logs:
-                        sd = log.get("status")
-                        if sd and isinstance(sd, dict) and "code" in sd:
-                            status_message = sd.get("message", "")
-                            status_details = sd.get("details", [])
-                            raw_extras = log.get("extras") or {}
-                            status_extras = {
-                                k: v for k, v in raw_extras.items()
-                                if k not in ("logger_name", "level", "timestamp")
-                            }
-                            break
+                    status_message = status.get("message", "") or ""
+                    status_details = status.get("details", []) or []
 
                     exception = exception_for_status_code(status_code, status_message)
                     if exception is not None:
@@ -169,10 +149,12 @@ async def _process_row(row: EvaluationRow) -> EvaluationRow:
                         details=status_details,
                     )
 
-                    if row.execution_metadata.extra:
-                        row.execution_metadata.extra.update(status_extras)
-                    else:
-                        row.execution_metadata.extra = status_extras
+                    status_extras = status.get("extras")
+                    if isinstance(status_extras, dict):
+                        if row.execution_metadata.extra:
+                            row.execution_metadata.extra.update(status_extras)
+                        else:
+                            row.execution_metadata.extra = status_extras
 
                     logger.info("Stopping polling for rollout %s", row.execution_metadata.rollout_id)
                     break