evalscope 1.1.0__py3-none-any.whl → 1.1.1__py3-none-any.whl

evalscope/perf/benchmark.py

@@ -3,8 +3,6 @@ import json
 import numpy as np
 import platform
 import sqlite3
-import time
-from http import HTTPStatus
 from tqdm import tqdm
 from typing import TYPE_CHECKING, AsyncGenerator, Dict, List, Tuple
 
@@ -82,86 +80,58 @@ async def send_request(
     request: dict,
     benchmark_data_queue: asyncio.Queue,
     args: Arguments,
-    api_plugin: 'ApiPluginBase',
+    client: AioHttpClient,  # reuse shared client
 ):
     async with semaphore:
-        client = AioHttpClient(args, api_plugin)
-        async with client:
-            benchmark_data = BenchmarkData(request=request)
-            benchmark_data.start_time = time.perf_counter()
-            collected_messages = []
-            try:
-                async for is_error, state_code, response_data in client.post(request):
-                    if is_error or state_code != HTTPStatus.OK:
-                        error_msg = str(response_data) if response_data else 'Unknown error'
-                        logger.error(f'Request: {request} failed, state_code: {state_code}, data: {error_msg}')
-                        benchmark_data.success = False
-                        break
-                    if response_data:
-                        collected_messages.append(response_data)
-                        benchmark_data.chunk_times.append(time.perf_counter())
-                        benchmark_data.success = True
-                        benchmark_data.update_gpu_usage()
-            except Exception as e:
-                if response_data:
-                    collected_messages.append(response_data)
-                benchmark_data.success = False
-                logger.exception(e)
-                logger.error(f'Request query: {request} exception')
-            finally:
-                benchmark_data.completed_time = time.perf_counter()
-                benchmark_data.response_messages = collected_messages
-                await benchmark_data_queue.put(benchmark_data)
+        benchmark_data = await client.post(request)
+        benchmark_data.update_gpu_usage()
+        await benchmark_data_queue.put(benchmark_data)
 
 
 @exception_handler
 async def statistic_benchmark_metric(benchmark_data_queue: asyncio.Queue, args: Arguments, api_plugin: 'ApiPluginBase'):
     metrics = BenchmarkMetrics(concurrency=args.parallel)
-
     result_db_path = get_result_db_path(args)
 
-    collected_benchmark_data = []
-
-    with tqdm(desc='Processing', total=args.number) as pbar:
-        while not (data_process_completed_event.is_set() and benchmark_data_queue.empty()):
-            try:
-                # Attempt to get benchmark data from the queue with a timeout
-                benchmark_data = await asyncio.wait_for(benchmark_data_queue.get(), timeout=0.01)
-                benchmark_data_queue.task_done()
-            except asyncio.TimeoutError:
-                # If timeout, continue to the next iteration
-                continue
-
-            # Update metrics based on the benchmark data
-            metrics.update_metrics(benchmark_data, api_plugin)
-
-            # Collect benchmark data for later database insertion
-            collected_benchmark_data.append(benchmark_data)
-
-            # Create a message with the updated metrics
-            message = metrics.create_message()
+    # Stream inserts to DB to avoid accumulating all results in memory
+    commit_every = args.db_commit_interval
+    processed_since_commit = 0
 
-            # Log the message to wandb\swanlab if the api key is provided
-            if args.wandb_api_key:
-                import wandb
-                wandb.log(message)
-            if args.swanlab_api_key:
-                import swanlab
-                swanlab.log(message)
-
-            # Log the message to the logger every n queries
-            if int(metrics.n_total_queries) % args.log_every_n_query == 0:
-                msg = json.dumps(message, ensure_ascii=False, indent=2)
-                logger.info(msg)
-
-            pbar.update(1)  # Update the progress bar
-
-    # Now perform database operations after all benchmark data has been processed
     with sqlite3.connect(result_db_path) as con:
         cursor = con.cursor()
         create_result_table(cursor)
-        for benchmark_data in collected_benchmark_data:
-            insert_benchmark_data(cursor, benchmark_data)
+
+        with tqdm(desc='Processing', total=args.number) as pbar:
+            while not (data_process_completed_event.is_set() and benchmark_data_queue.empty()):
+                try:
+                    benchmark_data = await asyncio.wait_for(benchmark_data_queue.get(), timeout=0.1)
+                except asyncio.TimeoutError:
+                    continue
+
+                # Update metrics and write to DB immediately
+                metrics.update_metrics(benchmark_data, api_plugin)
+                insert_benchmark_data(cursor, benchmark_data)
+                processed_since_commit += 1
+                if processed_since_commit >= commit_every:
+                    con.commit()
+                    processed_since_commit = 0
+
+                message = metrics.create_message()
+
+                if args.wandb_api_key:
+                    import wandb
+                    wandb.log(message)
+                if args.swanlab_api_key:
+                    import swanlab
+                    swanlab.log(message)
+
+                if int(metrics.n_total_queries) % args.log_every_n_query == 0:
+                    msg = json.dumps(message, ensure_ascii=False, indent=2)
+                    logger.info(msg)
+
+                benchmark_data_queue.task_done()
+                pbar.update(1)
+
         con.commit()
 
     return metrics, result_db_path
@@ -179,31 +149,46 @@ async def benchmark(args: Arguments) -> Tuple[Dict, Dict]:
         loop = asyncio.get_running_loop()
         add_signal_handlers(loop)
 
-    # Create API plugin instance for request/response processing
     api_plugin_class = ApiRegistry.get_class(args.api)
     api_plugin = api_plugin_class(args)
 
-    # init queue
-    benchmark_data_queue = asyncio.Queue()
-    # reset event
+    benchmark_data_queue: asyncio.Queue = asyncio.Queue(maxsize=max(1, args.parallel * args.queue_size_multiplier))
     data_process_completed_event.clear()
+
     # test connection
     await connect_test(args, api_plugin)
-    # start statistic benchmark metric
-    statistic_benchmark_metric_task = asyncio.create_task(
-        statistic_benchmark_metric(benchmark_data_queue, args, api_plugin)
-    )
-    # start send request
-    semaphore = asyncio.Semaphore(args.parallel)
-    send_request_tasks: List[asyncio.Task] = []
-    async for request in get_requests(args, api_plugin):
-        task = asyncio.create_task(send_request(semaphore, request, benchmark_data_queue, args, api_plugin))
-        send_request_tasks.append(task)
-
-    await asyncio.gather(*send_request_tasks, return_exceptions=True)
-    await benchmark_data_queue.join()
-    data_process_completed_event.set()
-
-    metrics, result_db_path = await statistic_benchmark_metric_task
+
+    # Create a single shared client session for all requests
+    client = AioHttpClient(args, api_plugin)
+    async with client:
+        # start statistic benchmark metric (consumer)
+        statistic_benchmark_metric_task = asyncio.create_task(
+            statistic_benchmark_metric(benchmark_data_queue, args, api_plugin)
+        )
+
+        # start sending requests with bounded in-flight tasks
+        semaphore = asyncio.Semaphore(args.parallel)
+        in_flight: set[asyncio.Task] = set()
+        max_in_flight = args.parallel * args.in_flight_task_multiplier
+
+        async for request in get_requests(args, api_plugin):
+            # Keep the number of scheduled tasks bounded to avoid OOM
+            if len(in_flight) >= max_in_flight:
+                done, pending = await asyncio.wait(in_flight, return_when=asyncio.FIRST_COMPLETED)
+                in_flight = pending
+
+            task = asyncio.create_task(send_request(semaphore, request, benchmark_data_queue, args, client))
+            in_flight.add(task)
+
+        # Wait for remaining in-flight tasks
+        if in_flight:
+            await asyncio.gather(*in_flight, return_exceptions=True)
+
+        # Drain queue and finish
+        await benchmark_data_queue.join()
+        data_process_completed_event.set()
+
+        metrics, result_db_path = await statistic_benchmark_metric_task
+
     metrics_result, percentile_result = summary_result(args, metrics, result_db_path)
     return metrics_result, percentile_result

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, Any], 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, Any]: (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__':