evalscope 1.0.0__py3-none-any.whl → 1.2.0__py3-none-any.whl

evalscope/perf/http_client.py

@@ -3,6 +3,7 @@ import asyncio
 import time
 from typing import TYPE_CHECKING, AsyncGenerator, Dict, List, Tuple
 
+from evalscope.perf.utils.benchmark_util import BenchmarkData
 from evalscope.utils.logger import get_logger
 from .arguments import Arguments
 
@@ -24,7 +25,22 @@ class AioHttpClient:
         self.read_timeout = args.read_timeout
         self.connect_timeout = args.connect_timeout
         self.api_plugin = api_plugin
+
+        # Configure connector similar to vLLM bench for better TTFT under load.
+        connector = aiohttp.TCPConnector(
+            limit=args.parallel or 0,  # 0 means no limit in aiohttp; use parallel as limit if set
+            limit_per_host=args.parallel or 0,
+            ttl_dns_cache=300,
+            use_dns_cache=True,
+            keepalive_timeout=60,
+            enable_cleanup_closed=True,
+            force_close=False,
+            ssl=('https://' in self.url),
+        )
+
         self.client = aiohttp.ClientSession(
+            connector=connector,
+            trust_env=True,
             timeout=aiohttp.ClientTimeout(connect=self.connect_timeout, sock_read=self.read_timeout),
             trace_configs=[self._create_trace_config()] if args.debug else []
         )
@@ -43,23 +59,25 @@ class AioHttpClient:
         trace_config.on_response_chunk_received.append(self.on_response_chunk_received)
         return trace_config
 
-    async def post(self, body):
-        """Send POST request and delegate response handling to API plugin.
-        Yields:
-            Tuple[bool, int, str]: (is_error, status_code, response_data)
+    async def post(self, body) -> BenchmarkData:
+        """
+        Send POST request and delegate response handling to API plugin.
+
+        Returns:
+            BenchmarkData: The benchmark data object containing request and response information.
         """
         try:
             # Delegate the request processing to the API plugin
-            async for result in self.api_plugin.process_request(self.client, self.url, self.headers, body):
-                yield result
+            output = await self.api_plugin.process_request(self.client, self.url, self.headers, body)
+            return output
         except asyncio.TimeoutError as e:
             logger.error(
                 f'TimeoutError: connect_timeout: {self.connect_timeout}, read_timeout: {self.read_timeout}. Please set longer timeout.'  # noqa: E501
             )
-            yield (True, None, str(e))
+            return BenchmarkData(success=False, error=str(e))
         except (aiohttp.ClientConnectorError, Exception) as e:
             logger.error(e)
-            yield (True, None, str(e))
+            return BenchmarkData(success=False, error=str(e))
 
     @staticmethod
     async def on_request_start(session, context, params: aiohttp.TraceRequestStartParams):
@@ -91,7 +109,6 @@ class AioHttpClient:
 
 
 async def test_connection(args: Arguments, api_plugin: 'ApiPluginBase') -> bool:
-    is_error = True
     start_time = time.perf_counter()
 
     async def attempt_connection():
@@ -100,18 +117,16 @@ async def test_connection(args: Arguments, api_plugin: 'ApiPluginBase') -> bool:
             messages = [{'role': 'user', 'content': 'hello'}] if args.apply_chat_template else 'hello'
             request = api_plugin.build_request(messages)
 
-            async for is_error, state_code, response_data in client.post(request):
-                return is_error, state_code, response_data
+            output = await client.post(request)
+            return output
 
     while True:
         try:
-            is_error, state_code, response_data = await asyncio.wait_for(
-                attempt_connection(), timeout=args.connect_timeout
-            )
-            if not is_error:
+            output = await asyncio.wait_for(attempt_connection(), timeout=args.connect_timeout)
+            if output.success:
                 logger.info('Test connection successful.')
                 return True
-            logger.warning(f'Retrying...  <{state_code}> {response_data}')
+            logger.warning(f'Retrying... <{output.error}>')
         except Exception as e:
             logger.warning(f'Retrying... <{e}>')
 

evalscope/perf/main.py

@@ -4,7 +4,9 @@ import os
 import platform
 import threading
 import time
+import warnings
 from argparse import Namespace
+from logging import warn
 
 from evalscope.perf.utils.local_server import start_app
 from evalscope.perf.utils.log_utils import init_swanlab, init_wandb
@@ -79,9 +81,20 @@ def run_perf_benchmark(args):
     configure_logging(args.debug, os.path.join(output_path, 'benchmark.log'))
 
     # Initialize wandb and swanlab
-    if args.wandb_api_key:
+    visualizer = args.visualizer
+    if visualizer is None:
+        if args.wandb_api_key is not None:
+            visualizer = 'wandb'
+            warnings.warn('--wandb-api-key is deprecated. Please use `--visualizer wandb` instead.', DeprecationWarning)
+        elif args.swanlab_api_key is not None:
+            visualizer = 'swanlab'
+            warnings.warn(
+                '--swanlab-api-key is deprecated. Please use `--visualizer swanlab` instead.', DeprecationWarning
+            )
+    args.visualizer = visualizer
+    if visualizer == 'wandb':
         init_wandb(args)
-    if args.swanlab_api_key:
+    elif visualizer == 'swanlab':
         init_swanlab(args)
 
     # Initialize local server if needed

evalscope/perf/plugin/api/base.py

@@ -3,6 +3,7 @@ from abc import abstractmethod
 from typing import Any, AsyncGenerator, Dict, List, Tuple
 
 from evalscope.perf.arguments import Arguments
+from evalscope.perf.utils.benchmark_util import BenchmarkData
 
 
 class ApiPluginBase:
@@ -28,13 +29,13 @@ class ApiPluginBase:
         raise NotImplementedError
 
     @abstractmethod
-    def parse_responses(self, responses: List, request: Any = None, **kwargs: Any) -> Tuple[int, int]:
+    def parse_responses(self, responses: List[Dict], request: str = None, **kwargs: Any) -> Tuple[int, int]:
         """Parser responses and return number of request and response tokens.
 
         Args:
-            responses (List[bytes]): List of http response body, for stream output,
+            responses (List[Dict]): List of http response body, for stream output,
                 there are multiple responses, each is bytes, for general only one.
-            request (Any): The request body.
+            request (str): The json string of request.
 
         Returns:
             Tuple: (Number of prompt_tokens and number of completion_tokens).
@@ -42,8 +43,9 @@ class ApiPluginBase:
         raise NotImplementedError
 
     @abstractmethod
-    async def process_request(self, client_session: aiohttp.ClientSession, url: str, headers: Dict,
-                              body: Dict) -> AsyncGenerator[Tuple[bool, int, str], None]:
+    async def process_request(
+        self, client_session: aiohttp.ClientSession, url: str, headers: Dict, body: Dict
+    ) -> BenchmarkData:
         """Process the HTTP request and handle the response.
 
         Args:
@@ -52,8 +54,8 @@ class ApiPluginBase:
             headers: The request headers
             body: The request body
 
-        Yields:
-            Tuple[bool, int, str]: (is_error, status_code, response_data)
+        Returns:
+            BenchmarkData: The benchmark data including response and timing info.
         """
         raise NotImplementedError
 

evalscope/perf/plugin/api/custom_api.py

@@ -5,6 +5,7 @@ from typing import Any, AsyncGenerator, Dict, List, Tuple, Union
 from evalscope.perf.arguments import Arguments
 from evalscope.perf.plugin.api.base import ApiPluginBase
 from evalscope.perf.plugin.registry import register_api
+from evalscope.perf.utils.benchmark_util import BenchmarkData
 from evalscope.utils.logger import get_logger
 
 logger = get_logger()
@@ -98,7 +99,7 @@ class CustomPlugin(ApiPluginBase):
 
         return payload
 
-    def parse_responses(self, responses: List[str], request: Any = None, **kwargs) -> Tuple[int, int]:
+    def parse_responses(self, responses: List[Dict], request: str = None, **kwargs) -> Tuple[int, int]:
         """Parse API responses and return token counts.
 
         This method extracts the number of input and output tokens from the API responses.
@@ -106,8 +107,8 @@ class CustomPlugin(ApiPluginBase):
         to calculate it using a tokenizer.
 
         Args:
-            responses (List[str]): List of API response strings.
-            request (Any, optional): The original request, which might be needed for token calculation.
+            responses (List[Dict]): List of API response strings.
+            request (str, optional): The original request, which might be needed for token calculation.
             **kwargs: Additional arguments.
 
         Returns:
@@ -160,8 +161,9 @@ class CustomPlugin(ApiPluginBase):
             logger.error(f'Error parsing responses: {e}')
             return 0, 0
 
-    async def process_request(self, client_session: aiohttp.ClientSession, url: str, headers: Dict,
-                              body: Dict) -> AsyncGenerator[Tuple[bool, int, str], None]:
+    async def process_request(
+        self, client_session: aiohttp.ClientSession, url: str, headers: Dict, body: Dict
+    ) -> BenchmarkData:
         """Process the HTTP request and handle the response.
 
         This method handles sending the request to your API and processing the response,
@@ -173,60 +175,13 @@ class CustomPlugin(ApiPluginBase):
             headers (Dict): The request headers.
             body (Dict): The request body.
 
-        Yields:
-            Tuple[bool, int, str]: (is_error, status_code, response_data)
-                - is_error: Whether the response indicates an error
-                - status_code: HTTP status code
-                - response_data: Response content
+        Returns:
+            BenchmarkData: The benchmark data including response and timing info.
         """
-        try:
-            # Set content type header
-            headers = {'Content-Type': 'application/json', **headers}
-
-            # Convert body to JSON
-            data = json.dumps(body, ensure_ascii=False)
-
-            # Send the request
-            async with client_session.request('POST', url=url, data=data, headers=headers) as response:  # noqa: E125
-                # Get the status code
-                status_code = response.status
-
-                # Check if it's a streaming response
-                if 'text/event-stream' in response.content_type:
-                    # Handle streaming response
-                    async for line in response.content:
-                        line_str = line.decode('utf-8').strip()
-                        if not line_str:
-                            continue
-
-                        # Check for data prefix in server-sent events
-                        if line_str.startswith('data: '):
-                            data = line_str[6:]  # Remove 'data: ' prefix
-
-                            # Check if it's the end of the stream
-                            if data == '[DONE]':
-                                break
-
-                            try:
-                                # Parse the JSON data
-                                parsed_data = json.loads(data)
-                                yield (False, status_code, json.dumps(parsed_data))
-                            except json.JSONDecodeError:
-                                yield (True, status_code, f'Failed to parse JSON: {data}')
-                else:
-                    # Handle regular response
-                    if 'application/json' in response.content_type:
-                        # JSON response
-                        content = await response.json()
-                        yield (status_code >= 400, status_code, json.dumps(content))
-                    else:
-                        # Text response
-                        content = await response.text()
-                        yield (status_code >= 400, status_code, content)
-
-        except Exception as e:
-            logger.error(f'Error in process_request: {e}')
-            yield (True, 500, str(e))
+        raise NotImplementedError(
+            'The `process_request` method must be implemented in a subclass. '
+            'For OpenAI-compatible APIs, consider inheriting from `DefaultApiPlugin` to reuse the default implementation.'  # noqa: E501
+        )
 
 
 if __name__ == '__main__':

evalscope/perf/plugin/api/default_api.py

@@ -1,24 +1,77 @@
 import aiohttp
+import codecs
 import json
-from http import HTTPStatus
-from typing import Any, AsyncGenerator, Dict, List, Tuple
+import sys
+import time
+import traceback
+from typing import Any, Dict
 
 from evalscope.perf.arguments import Arguments
 from evalscope.perf.plugin.api.base import ApiPluginBase
-from evalscope.perf.utils.local_server import ServerSentEvent
+from evalscope.perf.utils.benchmark_util import BenchmarkData
 from evalscope.utils.logger import get_logger
 
 logger = get_logger()
 
 
+class StreamedResponseHandler:
+    """Handles streaming HTTP responses by accumulating chunks until complete
+    messages are available."""
+
+    def __init__(self):
+        self.buffer = ''
+        # Keep decoder state across chunks to handle split multibyte sequences
+        self.decoder = codecs.getincrementaldecoder('utf-8')()
+
+    def add_chunk(self, chunk_bytes: bytes) -> list[str]:
+        """Add a chunk of bytes to the buffer and return any complete
+        messages."""
+        # Use incremental decoding so incomplete multibyte sequences don't error
+        try:
+            chunk_str = self.decoder.decode(chunk_bytes, final=False)
+        except UnicodeDecodeError:
+            # Bad bytes: drop them and reset decoder state to avoid corruption
+            self.decoder.reset()
+            chunk_str = chunk_bytes.decode('utf-8', errors='ignore')
+        self.buffer += chunk_str
+
+        messages = []
+
+        # Split by double newlines (SSE message separator)
+        while '\n\n' in self.buffer:
+            message, self.buffer = self.buffer.split('\n\n', 1)
+            message = message.strip()
+            if message:
+                messages.append(message)
+
+        # if self.buffer is not empty, check if it is a complete message
+        # by removing data: prefix and check if it is a valid JSON
+        if self.buffer.startswith('data: '):
+            message_content = self.buffer.removeprefix('data: ').strip()
+            if message_content == '[DONE]':
+                messages.append(self.buffer.strip())
+                self.buffer = ''
+            elif message_content:
+                try:
+                    json.loads(message_content)
+                    messages.append(self.buffer.strip())
+                    self.buffer = ''
+                except json.JSONDecodeError:
+                    # Incomplete JSON, wait for more chunks.
+                    pass
+
+        return messages
+
+
 class DefaultApiPlugin(ApiPluginBase):
     """Default implementation of API plugin with common HTTP handling methods."""
 
     def __init__(self, param: Arguments):
         super().__init__(param)
 
-    async def process_request(self, client_session: aiohttp.ClientSession, url: str, headers: Dict,
-                              body: Dict) -> AsyncGenerator[Tuple[bool, int, str], None]:
+    async def process_request(
+        self, client_session: aiohttp.ClientSession, url: str, headers: Dict, body: Dict
+    ) -> BenchmarkData:
         """Process the HTTP request and handle the response.
 
         Args:
@@ -27,79 +80,135 @@ class DefaultApiPlugin(ApiPluginBase):
             headers: The request headers
             body: The request body
 
-        Yields:
-            Tuple[bool, int, str]: (is_error, status_code, response_data)
-        """
-        try:
-            headers = {'Content-Type': 'application/json', **headers}
-            data = json.dumps(body, ensure_ascii=False)  # serialize to JSON
-            async with client_session.request('POST', url=url, data=data, headers=headers) as response:
-                async for result in self._handle_response(response):
-                    yield result
-        except Exception as e:
-            logger.error(f'Error in process_request: {e}')
-            yield (True, None, str(e))
-
-    async def _handle_stream(self, response: aiohttp.ClientResponse) -> AsyncGenerator[Tuple[bool, int, str], None]:
-        """Handle streaming response from server-sent events.
-
-        Args:
-            response: The aiohttp response object containing a stream
-
-        Yields:
-            Tuple[bool, int, Any]: (is_error, status_code, data)
+        Returns:
+            BenchmarkData: Aggregated benchmarking data for the request/response.
         """
+        headers = {'Content-Type': 'application/json', **headers}
+        data = json.dumps(body, ensure_ascii=False)  # serialize to JSON
+
+        output = BenchmarkData()
+        ttft = 0.0
+        generated_text = ''
+        st = time.perf_counter()
+        output.start_time = st
+        output.request = data
+        most_recent_timestamp = st
         try:
-            async for chunk_bytes in response.content:
-                chunk_bytes = chunk_bytes.strip()
-                if not chunk_bytes:
-                    continue
-                chunk_bytes = chunk_bytes.decode('utf-8')
-                # NOTE: SSE comments (often used as pings) start with a colon.
-                # These are not JSON data payload and should be skipped.
-                if chunk_bytes.startswith(':'):
-                    continue
-
-                chunk = chunk_bytes.removeprefix('data: ')
-
-                if chunk != '[DONE]':
-                    data = json.loads(chunk)
-
-                    yield False, response.status, data
-
-        except Exception as e:
-            logger.error(f'Error in _handle_stream: {e}')
-            yield True, response.status, str(e)
-
-    async def _handle_response(self, response: aiohttp.ClientResponse) -> AsyncGenerator[Tuple[bool, int, str], None]:
-        """Handle the HTTP response based on content type and status.
-
-        Args:
-            response: The aiohttp response object
-
-        Yields:
-            Tuple[bool, int, str]: (is_error, status_code, response_data)
-        """
-        response_status = response.status
-        response_content_type = response.content_type
-        content_type_json = 'application/json'
-        content_type_stream = 'text/event-stream'
-        is_success = (response_status == HTTPStatus.OK)
-
-        if is_success:
-            # Handle successful response with 'text/event-stream' content type
-            if content_type_stream in response_content_type:
-                async for is_error, response_status, content in self._handle_stream(response):
-                    yield (is_error, response_status, content)
-            # Handle successful response with 'application/json' content type
-            elif content_type_json in response_content_type:
-                content = await response.json()
-                yield (False, response_status, json.dumps(content, ensure_ascii=False))
-            # Handle other successful responses
-            else:
-                content = await response.read()
-                yield (False, response_status, content.decode('utf-8'))
-        else:
-            # error is always in JSON format
-            error = await response.json()
-            yield (True, response_status, json.dumps(error, ensure_ascii=False))
+            async with client_session.post(url=url, data=data, headers=headers) as response:
+                content_type = response.headers.get('Content-Type', '')
+                if response.status == 200:
+                    # Handle streaming responses (SSE)
+                    if 'text/event-stream' in content_type:
+                        handler = StreamedResponseHandler()
+                        async for chunk_bytes in response.content.iter_any():
+
+                            if not chunk_bytes:
+                                continue
+
+                            messages = handler.add_chunk(chunk_bytes)
+                            for message in messages:
+                                # NOTE: SSE comments (often used as pings) start with
+                                # a colon. These are not JSON data payload and should
+                                # be skipped.
+                                if message.startswith(':'):
+                                    continue
+
+                                chunk = message.removeprefix('data: ')
+
+                                if chunk != '[DONE]':
+                                    timestamp = time.perf_counter()
+                                    data = json.loads(chunk)
+
+                                    if choices := data.get('choices'):
+                                        content = choices[0]['delta'].get('content')
+                                        # First token
+                                        if ttft == 0.0:
+                                            ttft = timestamp - st
+                                            output.first_chunk_latency = ttft
+
+                                        # Decoding phase
+                                        else:
+                                            output.inter_chunk_latency.append(timestamp - most_recent_timestamp)
+
+                                        generated_text += content or ''
+                                        output.response_messages.append(data)
+                                    elif usage := data.get('usage'):
+                                        output.prompt_tokens = usage.get('prompt_tokens')
+                                        output.completion_tokens = usage.get('completion_tokens')
+
+                                    most_recent_timestamp = timestamp
+
+                        output.generated_text = generated_text
+                        output.success = True
+                        output.completed_time = most_recent_timestamp
+                        output.query_latency = most_recent_timestamp - st
+
+                    # Handle non-stream JSON responses
+                    elif 'application/json' in content_type or 'application/' in content_type:
+                        payload: Any
+                        try:
+                            payload = await response.json()
+                        except Exception:
+                            # Fallback to text if JSON parsing fails
+                            payload = await response.text()
+
+                        timestamp = time.perf_counter()
+                        output.completed_time = timestamp
+                        output.query_latency = timestamp - st
+                        # For non-stream, first chunk equals full latency
+                        output.first_chunk_latency = output.query_latency
+
+                        if isinstance(payload, dict):
+                            # Extract generated text from choices
+                            text = ''
+                            if choices := payload.get('choices'):
+                                first = choices[0] if choices else {}
+                                # Chat Completions format
+                                msg = first.get('message') or {}
+                                if isinstance(msg, dict) and msg.get('content') is not None:
+                                    text = msg.get('content') or ''
+                                else:
+                                    # Legacy Completions format
+                                    text = first.get('text') or ''
+                            generated_text = text
+
+                            # Extract usage if provided
+                            if usage := payload.get('usage'):
+                                output.prompt_tokens = usage.get('prompt_tokens')
+                                output.completion_tokens = usage.get('completion_tokens')
+
+                            output.response_messages.append(payload)
+                        else:
+                            generated_text = str(payload)
+
+                        output.generated_text = generated_text
+                        output.success = True
+
+                    else:
+                        # Unknown successful content-type: read as text
+                        raw = await response.text()
+                        timestamp = time.perf_counter()
+                        output.completed_time = timestamp
+                        output.query_latency = timestamp - st
+                        output.first_chunk_latency = output.query_latency
+                        output.generated_text = raw
+                        output.response_messages.append(raw)
+                        output.success = True
+                else:
+                    # Try to parse structured error, fallback to reason/text
+                    try:
+                        err_payload = await response.json()
+                        output.error = json.dumps(err_payload, ensure_ascii=False)
+                    except Exception:
+                        try:
+                            output.error = await response.text()
+                        except Exception:
+                            output.error = response.reason or ''
+                    output.success = False
+        except Exception:
+            output.success = False
+            exc_info = sys.exc_info()
+            output.error = ''.join(traceback.format_exception(*exc_info))
+            logger.error(output.error)
+
+        return output