initial commit

Author: SunnySoldier357

eval_protocol/adapters/fireworks_tracing.py

@@ -375,6 +375,38 @@ async def async_search_logs(
             )
         return results
 
+    def get_status(self, 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}
+
+        urls_to_try = [f"{self.base_url}/status", f"{self.base_url}/v1/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:
+                last_error = str(e)
+                continue
+
+        if last_error:
+            logger.error("Failed to fetch status from Fireworks (tried %s): %s", urls_to_try, last_error)
+        return None
+
     def get_evaluation_rows(
         self,
         tags: List[str],

eval_protocol/pytest/remote_rollout_processor.py

@@ -121,46 +121,39 @@ async def _process_row(row: EvaluationRow) -> EvaluationRow:
             deadline = time.time() + timeout_seconds
 
             while time.time() < deadline:
-                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}"]
+                # 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
                 )
-                # Filter for logs that actually have status information
-                status_logs = []
-                for log in completed_logs:
-                    status_dict = log.get("status")
-                    if status_dict and isinstance(status_dict, dict) and "code" in status_dict:
-                        status_logs.append(log)
-
-                if status_logs:
-                    if len(status_logs) > 1:
-                        logger.warning(
-                            "Found %s status logs for rollout %s; expected at most 1. Using the first one: %s",
-                            len(status_logs),
-                            row.execution_metadata.rollout_id,
-                            status_logs[0],
-                        )
-                    # Use the first log with status information
-                    status_log = status_logs[0]
-                    status_dict = status_log.get("status")
-                    raw_extras = status_log.get("extras") or {}
-                    status_extras = {
-                        k: v for k, v in raw_extras.items() if k not in ("logger_name", "level", "timestamp")
-                    }
+                if status_result and status_result.get("status"):
+                    status_code = status_result["status"]["code"]
 
                     logger.info(
-                        f"Found status log for rollout {row.execution_metadata.rollout_id}: {status_log.get('message', '')}"
+                        "Found status for rollout %s with code %s",
+                        row.execution_metadata.rollout_id,
+                        status_code,
                     )
 
-                    status_code = status_dict.get("code")
-                    status_message = status_dict.get("message", "")
-                    status_details = status_dict.get("details", [])
-
-                    logger.info(
-                        f"Found Fireworks log for rollout {row.execution_metadata.rollout_id} with status code {status_code}"
+                    # Backfill message/details/extras from the full Logs table (one-shot)
+                    completed_logs = await asyncio.to_thread(
+                        self._tracing_adapter.search_logs,
+                        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
 
-                    # Create and raise exception if appropriate, preserving original message
                     exception = exception_for_status_code(status_code, status_message)
                     if exception is not None:
                         raise exception