fix: use async connection pooling for tracing log polls

Replace per-request requests.get() calls in search_logs with a native
async aiohttp implementation that reuses the RemoteRolloutProcessor's
shared ClientSession. This eliminates the SSL EOF errors seen under high
concurrency by reusing TCP+TLS connections instead of opening a fresh
handshake on every poll.

Additional fixes:
- Remove default connection pool cap (TCPConnector limit=0) so the
  semaphore is the only concurrency bound
- Increase /init timeout from 300s to 600s
- Move common_utils import to module level

Made-with: Cursor

Author: xzrderek

eval_protocol/adapters/fireworks_tracing.py

@@ -5,7 +5,9 @@
 """
 
 from __future__ import annotations
+import asyncio
 import logging
+import aiohttp
 import requests
 from datetime import datetime
 from typing import Any, Dict, List, Optional, Protocol
@@ -14,6 +16,7 @@
 from eval_protocol.models import EvaluationRow, InputMetadata, ExecutionMetadata, Message
 from .base import BaseAdapter
 from .utils import extract_messages_from_data
+from ..common_utils import get_user_agent
 
 logger = logging.getLogger(__name__)
 
@@ -280,8 +283,6 @@ def search_logs(self, tags: List[str], limit: int = 100, hours_back: int = 24) -
         if not tags:
             raise ValueError("At least one tag is required to fetch logs")
 
-        from ..common_utils import get_user_agent
-
         headers = {
             "Authorization": f"Bearer {self._get_api_key()}",
             "User-Agent": get_user_agent(),
@@ -327,6 +328,52 @@ def search_logs(self, tags: List[str], limit: int = 100, hours_back: int = 24) -
             )
         return results
 
+    async def async_search_logs(
+        self, session: aiohttp.ClientSession, tags: List[str], limit: int = 100, hours_back: int = 24
+    ) -> List[Dict[str, Any]]:
+        """Async version of search_logs, reuses a caller-provided aiohttp session."""
+        if not tags:
+            raise ValueError("At least one tag is required to fetch logs")
+
+        params: Dict[str, Any] = {"tags": tags, "limit": limit, "hours_back": hours_back, "program": "eval_protocol"}
+        headers = {"Authorization": f"Bearer {self._get_api_key()}", "User-Agent": get_user_agent()}
+        timeout = aiohttp.ClientTimeout(total=self.timeout)
+
+        urls_to_try = [f"{self.base_url}/logs", f"{self.base_url}/v1/logs"]
+        data: Dict[str, Any] = {}
+        last_error: Optional[str] = None
+        for url in urls_to_try:
+            try:
+                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()
+                    data = (await resp.json()) or {}
+                    break
+            except (aiohttp.ClientError, asyncio.TimeoutError) as e:
+                last_error = str(e)
+                continue
+        else:
+            if last_error:
+                logger.error("Failed to fetch logs from Fireworks (tried %s): %s", urls_to_try, last_error)
+            return []
+
+        entries: List[Dict[str, Any]] = data.get("entries", []) or []
+        results: List[Dict[str, Any]] = []
+        for e in entries:
+            results.append(
+                {
+                    "timestamp": e.get("timestamp"),
+                    "message": e.get("message"),
+                    "severity": e.get("severity", "INFO"),
+                    "tags": e.get("tags", []),
+                    "status": e.get("status"),
+                    "extras": e.get("extras"),
+                }
+            )
+        return results
+
     def get_evaluation_rows(
         self,
         tags: List[str],
@@ -411,8 +458,6 @@ def get_evaluation_rows(
         else:
             url = f"{self.base_url}/v1/traces/pointwise"
 
-        from ..common_utils import get_user_agent
-
         headers = {
             "Authorization": f"Bearer {self._get_api_key()}",
             "User-Agent": get_user_agent(),

eval_protocol/pytest/remote_rollout_processor.py

@@ -52,7 +52,8 @@ def __init__(
 
     def _get_or_create_session(self) -> aiohttp.ClientSession:
         if self._session is None or self._session.closed:
-            self._session = aiohttp.ClientSession()
+            connector = aiohttp.TCPConnector(limit=0)  # no total connection limit
+            self._session = aiohttp.ClientSession(connector=connector)
         return self._session
 
     def __call__(self, rows: List[EvaluationRow], config: RolloutProcessorConfig) -> List[asyncio.Task[EvaluationRow]]:
@@ -99,7 +100,7 @@ async def _process_row(row: EvaluationRow) -> EvaluationRow:
             # Fire-and-poll
             init_url = f"{remote_base_url}/init"
 
-            timeout_init = aiohttp.ClientTimeout(total=300)
+            timeout_init = aiohttp.ClientTimeout(total=600)
 
             try:
                 session = self._get_or_create_session()
@@ -114,15 +115,15 @@ async def _process_row(row: EvaluationRow) -> EvaluationRow:
                     await resp.read()  # Drain the response body and release the connection back to the pool
             except asyncio.TimeoutError:
                 raise TimeoutError(
-                    f"The /init endpoint tried {init_url} with {init_payload.model_dump()} but timed out after 300 seconds."
+                    f"The /init endpoint tried {init_url} with {init_payload.model_dump()} but timed out after 600 seconds."
                 )
 
             deadline = time.time() + timeout_seconds
 
             while time.time() < deadline:
-                # Search Fireworks tracing logs for completion (run in thread to avoid blocking event loop)
-                completed_logs = await asyncio.to_thread(
-                    self._tracing_adapter.search_logs, tags=[f"rollout_id:{row.execution_metadata.rollout_id}"]
+                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}"]
                 )
                 # Filter for logs that actually have status information
                 status_logs = []