test

Author: xzrderek

eval_protocol/mcp/client/connection.py

@@ -12,6 +12,7 @@
 from contextlib import AsyncExitStack
 from typing import Any, Dict, List, Optional, Tuple
 
+import httpx
 from mcp.client.session import ClientSession
 from mcp.client.streamable_http import streamablehttp_client
 
@@ -26,6 +27,9 @@ class MCPConnectionManager:
     def __init__(self):
         self._tools_cache: Dict[str, List[Dict]] = {}
         self._tools_cache_lock = asyncio.Lock()
+        # Shared HTTP client for control plane requests with high connection limits
+        self._shared_client: Optional[httpx.AsyncClient] = None
+        self._client_lock = asyncio.Lock()
 
     async def initialize_session(self, session: MCPSession) -> None:
         """
@@ -114,6 +118,12 @@ async def _prewarm_tools_cache(self, session: MCPSession) -> None:
         """
         cache_key = session.base_url
 
+        # Fast path: if cache already exists, return immediately (no lock)
+        if cache_key in self._tools_cache:
+            logger.debug(f"Tools cache already exists for {cache_key}")
+            return
+
+        # Slow path: need to create cache (use lock only for creation)
         async with self._tools_cache_lock:
             # Only fetch tools if not already cached for this base_url
             if cache_key not in self._tools_cache:
@@ -244,21 +254,33 @@ async def get_initial_state(self, session: MCPSession) -> Any:
                     # Use shorter timeout for playback mode, longer timeout for high-concurrency initialization
                     # (50+ concurrent sessions need more time for initial state setup)
                     timeout = 3.0 if hasattr(session, "_is_playback_mode") and session._is_playback_mode else 15.0
-                    async with httpx.AsyncClient(timeout=timeout) as client:
-                        initial_state_response = await client.get(
-                            f"{base_url}/control/initial_state",
-                            headers=headers,
-                            timeout=timeout,
+
+                    # TIMING: Get shared client
+                    client_start = __import__("time").time()
+                    client = await self._get_shared_client(timeout)
+                    client_time = __import__("time").time() - client_start
+                    logger.info(
+                        f"DEBUG_CLIENT: Getting shared client took {client_time:.3f}s for {session.session_id}"
+                    )
+
+                    # TIMING: HTTP request with shared client
+                    request_start = __import__("time").time()
+                    initial_state_response = await client.get(
+                        f"{base_url}/control/initial_state",
+                        headers=headers,
+                        timeout=timeout,
+                    )
+                    request_time = __import__("time").time() - request_start
+                    logger.info(f"DEBUG_REQUEST: HTTP request took {request_time:.3f}s for {session.session_id}")
+                    if initial_state_response.status_code == 200:
+                        initial_observation = initial_state_response.json()
+                        logger.info(
+                            f"Session {session.session_id}: ✅ Successfully fetched session-aware initial state from control plane endpoint"
+                        )
+                    else:
+                        logger.warning(
+                            f"Control plane initial state endpoint returned {initial_state_response.status_code}"
                         )
-                        if initial_state_response.status_code == 200:
-                            initial_observation = initial_state_response.json()
-                            logger.info(
-                                f"Session {session.session_id}: ✅ Successfully fetched session-aware initial state from control plane endpoint"
-                            )
-                        else:
-                            logger.warning(
-                                f"Control plane initial state endpoint returned {initial_state_response.status_code}"
-                            )
                 except httpx.TimeoutException:
                     logger.warning(f"Control plane initial state endpoint timed out after {timeout}s")
                 except Exception as e:
@@ -579,3 +601,47 @@ async def close_session(self, session: MCPSession) -> None:
             finally:
                 session._exit_stack = None
                 session._mcp_session = None
+
+    async def _get_shared_client(self, timeout: float) -> httpx.AsyncClient:
+        """
+        Get or create a shared HTTP client with high connection limits for concurrent requests.
+
+        Args:
+            timeout: Timeout for requests
+
+        Returns:
+            Shared httpx.AsyncClient instance
+        """
+        # Fast path: if client exists and is not closed, return it immediately
+        if self._shared_client is not None and not self._shared_client.is_closed:
+            return self._shared_client
+
+        # Slow path: need to create client (use lock only for creation)
+        async with self._client_lock:
+            # Double-check pattern: another task might have created it while we waited
+            if self._shared_client is None or self._shared_client.is_closed:
+                # Create HTTP client with high connection limits for concurrent initial state requests
+                limits = httpx.Limits(
+                    max_keepalive_connections=None,  # Unlimited keep-alive connections
+                    max_connections=None,  # Unlimited total connection pool size
+                    keepalive_expiry=30.0,  # Keep connections alive for 30s
+                )
+
+                self._shared_client = httpx.AsyncClient(
+                    timeout=timeout,
+                    limits=limits,
+                    # Enable connection pooling and keep-alive
+                    http2=False,  # Disable HTTP/2 for better connection pooling with many concurrent requests
+                )
+                logger.info(
+                    "Created shared HTTP client with unlimited connection limits for MCP control plane requests"
+                )
+
+        return self._shared_client
+
+    async def close_shared_client(self):
+        """Close the shared HTTP client when shutting down."""
+        async with self._client_lock:
+            if self._shared_client and not self._shared_client.is_closed:
+                await self._shared_client.aclose()
+                self._shared_client = None

eval_protocol/mcp/execution/manager.py

@@ -13,6 +13,7 @@
 import time
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from dataclasses import asdict, dataclass
+from datetime import datetime
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union
 
 from openai.types import CompletionUsage
@@ -188,6 +189,11 @@ async def _execute_rollout(
         """
         session = envs.sessions[rollout_idx]
         dataset_row = envs.dataset_rows[rollout_idx]
+        rollout_start = time.time()
+        elapsed_from_main_start = rollout_start - start_time
+        logger.info(
+            f"DEBUG4. Starting rollout {dataset_row.id} at {datetime.fromtimestamp(rollout_start).strftime('%H:%M:%S.%f')[:-3]} (+{elapsed_from_main_start:.3f}s from start)"
+        )
 
         # Initialize trajectory
         trajectory = Trajectory(
@@ -210,7 +216,11 @@ async def _execute_rollout(
             },
         )
 
+        temp_start = time.time()
         current_observation, tool_schema = await envs.reset(session)
+        logger.info(
+            f"DEBUG6: User simulator get_init_state took {time.time() - temp_start:.3f}s for {session.session_id}"
+        )
         system_prompt = dataset_row.system_prompt
 
         # Record initial observation
@@ -244,7 +254,9 @@ async def _execute_rollout(
 
         usage_stats_list: List[CompletionUsage] = []
 
-        logger.info(f"🎯 Starting rollout {rollout_idx} in thread {threading.current_thread().name}")
+        logger.info(
+            f"DEBUG7: 🎯 Starting rollout {dataset_row.id} in thread {threading.current_thread().name}, {datetime.fromtimestamp(time.time()).strftime('%H:%M:%S.%f')[:-3]} (+{time.time() - rollout_start:.3f}s from start)"
+        )
 
         # Run rollout loop for this specific environment
         step = 0
@@ -264,9 +276,13 @@ async def _execute_rollout(
                 # Last message was agent, simulated user response
                 if user_simulator_messages and isinstance(user_simulator_messages[-1], AssistantMessage):
                     # Generate user response using the simulator
+                    temp_start1 = time.time()
                     user_message, user_simulator_state = user_simulator.generate_next_message(
                         user_simulator_messages[-1], user_simulator_state
                     )
+                    logger.info(
+                        f"DEBUG8: User simulator generate_next_message took {time.time() - temp_start1:.3f}s for {dataset_row.id}"
+                    )
                     user_content = user_message.content if user_message.content else ""
 
                     user_prompt = envs.format_user_prompt(rollout_idx, user_content)
@@ -279,7 +295,9 @@ async def _execute_rollout(
 
             # In each turn: keep looping until assistant is ready to provide final response
             while not turn_completed and not trajectory.terminated:
+                temp_start2 = time.time()
                 tool_calls, usage_stats = await policy(tool_schema, rollout_idx, conversation_history)
+                logger.info(f"DEBUG9: Policy took {time.time() - temp_start2:.3f}s for {dataset_row.id}")
 
                 # If no tool call is generated, turn is finished
                 if len(tool_calls) == 1:
@@ -296,7 +314,9 @@ async def _execute_rollout(
                 for tool_call in tool_calls:
 
                     # Execute tool call for this environment
+                    temp_start3 = time.time()
                     observation, reward, rollout_end, info = await envs.step(rollout_idx, tool_call)
+                    logger.info(f"DEBUG10: Env step took {time.time() - temp_start3:.3f}s for {dataset_row.id}")
 
                     tool_response = envs.format_tool_response(observation)
 
@@ -444,6 +464,9 @@ async def _execute_rollout(
         logger.info(
             f"✅ Rollout {rollout_idx} completed: {trajectory.steps} steps, reward: {trajectory.total_reward:.2f}, termination: {trajectory.termination_reason}, in thread {threading.current_thread().name}"
         )
+        logger.info(
+            f"DEBUG11: Rollout {dataset_row.id} completed at {datetime.fromtimestamp(time.time()).strftime('%H:%M:%S.%f')[:-3]} (+{time.time() - rollout_start:.3f}s from start)"
+        )
         return trajectory
 
     async def _get_control_plane_status(self, session) -> Optional[Dict[str, Any]]:

eval_protocol/mcp/execution/policy.py

@@ -19,6 +19,7 @@
 from .base_policy import LLMBasePolicy
 
 logger = logging.getLogger(__name__)
+litellm._turn_on_debug()
 
 
 class LiteLLMPolicy(LLMBasePolicy):

eval_protocol/pytest/default_mcp_gym_rollout_processor.py

@@ -213,11 +213,14 @@ async def default_mcp_gym_rollout_processor(
         )
 
         # Create MCP environments directly from evaluation_rows
+        print("DEBUG1", time.time())
         envs = await ep.make(
             "http://localhost:9700/mcp/",
             evaluation_rows=rows,
             model_id=policy.model_id,
         )
+        print("DEBUG2", time.time())
+        print("max_concurrent_rollouts", config.max_concurrent_rollouts)
 
         # Run rollout with environments and policy
         evaluation_rows = await ep.rollout(

monitor_connections.sh

@@ -0,0 +1,11 @@
+#!/bin/bash
+
+echo "Monitoring connections to port 9700..."
+echo "Press Ctrl+C to stop"
+
+while true; do
+    count=$(netstat -an | grep :9700 | grep ESTABLISHED | wc -l)
+    timestamp=$(date '+%H:%M:%S')
+    echo "$timestamp: $count connections to port 9700"
+    sleep 1
+done