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

evalscope/perf/http_client.py

@@ -1,13 +1,14 @@
 import aiohttp
 import asyncio
-import json
 import time
-from http import HTTPStatus
-from typing import AsyncGenerator, Dict, List, Tuple
+from typing import TYPE_CHECKING, AsyncGenerator, Dict, List, Tuple
 
-from evalscope.perf.arguments import Arguments
-from evalscope.perf.utils.local_server import ServerSentEvent
+from evalscope.perf.utils.benchmark_util import BenchmarkData
 from evalscope.utils.logger import get_logger
+from .arguments import Arguments
+
+if TYPE_CHECKING:
+    from .plugin.api.base import ApiPluginBase
 
 logger = get_logger()
 
@@ -17,94 +18,66 @@ class AioHttpClient:
     def __init__(
         self,
         args: Arguments,
+        api_plugin: 'ApiPluginBase',
     ):
         self.url = args.url
         self.headers = {'user-agent': 'modelscope_bench', **(args.headers or {})}
         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(
-            timeout=aiohttp.ClientTimeout(
-                total=self.read_timeout + self.connect_timeout,
-                connect=self.connect_timeout,
-                sock_read=self.read_timeout),
-            connector=aiohttp.TCPConnector(limit=1),
-            trace_configs=[self._create_trace_config()] if args.debug else [])
+            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 []
+        )
+
+    async def __aenter__(self):
+        pass
+
+    async def __aexit__(self, exc_type, exc, tb):
+        await self.client.close()
 
     def _create_trace_config(self):
+        """Create trace configuration for debugging."""
         trace_config = aiohttp.TraceConfig()
         trace_config.on_request_start.append(self.on_request_start)
         trace_config.on_request_chunk_sent.append(self.on_request_chunk_sent)
         trace_config.on_response_chunk_received.append(self.on_response_chunk_received)
         return trace_config
 
-    async def __aenter__(self):
-        pass
+    async def post(self, body) -> BenchmarkData:
+        """
+        Send POST request and delegate response handling to API plugin.
 
-    async def __aexit__(self, exc_type, exc, tb):
-        await self.client.close()
-
-    async def _handle_stream(self, response: aiohttp.ClientResponse):
-        is_error = False
-        async for line in response.content:
-            line = line.decode('utf8').rstrip('\n\r')
-            sse_msg = ServerSentEvent.decode(line)
-            if sse_msg:
-                logger.debug(f'Response recevied: {line}')
-                if sse_msg.event == 'error':
-                    is_error = True
-                if sse_msg.data:
-                    if sse_msg.data.startswith('[DONE]'):
-                        break
-                    yield is_error, response.status, sse_msg.data
-
-    async def _handle_response(self, response: aiohttp.ClientResponse) -> AsyncGenerator[Tuple[bool, int, str], None]:
-        response_status = response.status
-        response_content_type = response.content_type
-        content_type_json = 'application/json'
-        content_type_event_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_event_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()
-                if content.get('object') == 'error':
-                    yield (True, content.get('code'), content.get('message'))  # DashScope
-                else:
-                    yield (False, response_status, json.dumps(content, ensure_ascii=False))
-            # Handle other successful responses
-            else:
-                content = await response.read()
-                yield (False, response_status, content)
-        else:
-            # Handle error response with 'application/json' content type
-            if content_type_json in response_content_type:
-                error = await response.json()
-                yield (True, response_status, json.dumps(error, ensure_ascii=False))
-            # Handle error response with 'text/event-stream' content type
-            elif content_type_event_stream in response_content_type:
-                async for _, _, data in self._handle_stream(response):
-                    error = json.loads(data)
-                    yield (True, response_status, json.dumps(error, ensure_ascii=False))
-            # Handle other error responses
-            else:
-                msg = await response.read()
-                yield (True, response_status, msg.decode('utf-8'))
-
-    async def post(self, body):
-        headers = {'Content-Type': 'application/json', **self.headers}
+        Returns:
+            BenchmarkData: The benchmark data object containing request and response information.
+        """
         try:
-            data = json.dumps(body, ensure_ascii=False)  # serialize to JSON
-            async with self.client.request('POST', url=self.url, data=data, headers=headers) as response:
-                async for rsp in self._handle_response(response):
-                    yield rsp
+            # Delegate the request processing to the API plugin
+            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
+            )
+            return BenchmarkData(success=False, error=str(e))
         except (aiohttp.ClientConnectorError, Exception) as e:
             logger.error(e)
-            yield (True, None, e)
+            return BenchmarkData(success=False, error=str(e))
 
     @staticmethod
     async def on_request_start(session, context, params: aiohttp.TraceRequestStartParams):
@@ -135,28 +108,25 @@ class AioHttpClient:
         logger.debug(f'Request received: <{method=},  {url=}, {truncated_chunk=}>')
 
 
-async def test_connection(args: Arguments) -> bool:
-    is_error = True
+async def test_connection(args: Arguments, api_plugin: 'ApiPluginBase') -> bool:
     start_time = time.perf_counter()
 
     async def attempt_connection():
-        client = AioHttpClient(args)
+        client = AioHttpClient(args, api_plugin)
         async with client:
-            if 'chat/completions' in args.url:
-                request = {'messages': [{'role': 'user', 'content': 'hello'}], 'model': args.model}
-            else:
-                request = {'prompt': 'hello', 'model': args.model}
-            async for is_error, state_code, response_data in client.post(request):
-                return is_error, state_code, response_data
+            messages = [{'role': 'user', 'content': 'hello'}] if args.apply_chat_template else 'hello'
+            request = api_plugin.build_request(messages)
+
+            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:
-                logger.info('Connection successful.')
+            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

@@ -1,19 +1,73 @@
 import asyncio
+import copy
 import os
 import platform
+import threading
+import time
+import warnings
 from argparse import Namespace
+from logging import warn
 
-from evalscope.perf.arguments import Arguments, parse_args
-from evalscope.perf.benchmark import benchmark
-from evalscope.perf.utils.db_util import get_output_path
-from evalscope.perf.utils.handler import add_signal_handlers
+from evalscope.perf.utils.local_server import start_app
+from evalscope.perf.utils.log_utils import init_swanlab, init_wandb
 from evalscope.utils.logger import configure_logging, get_logger
-from evalscope.utils.utils import seed_everything
+from evalscope.utils.model_utils import seed_everything
+from .arguments import Arguments, parse_args
+from .benchmark import benchmark
+from .utils.db_util import get_output_path
+from .utils.handler import add_signal_handlers
+from .utils.rich_display import print_summary
 
 logger = get_logger()
 
 
+def run_one_benchmark(args: Arguments, output_path: str = None):
+    if isinstance(args.parallel, list):
+        args.parallel = args.parallel[0]
+    if isinstance(args.number, list):
+        args.number = args.number[0]
+
+    # Setup logger and output
+    args.outputs_dir = output_path
+
+    logger.info('Starting benchmark with args: ')
+    logger.info(args)
+
+    if platform.system() == 'Windows':
+        asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
+
+    loop = asyncio.new_event_loop()
+    if platform.system() != 'Windows':
+        add_signal_handlers(loop)
+
+    return loop.run_until_complete(benchmark(args))
+
+
+def run_multi_benchmark(args: Arguments, output_path: str = None):
+    results = []
+    number_list = copy.deepcopy(args.number)
+    parallel_list = copy.deepcopy(args.parallel)
+    for i, (number, parallel) in enumerate(zip(number_list, parallel_list)):
+        args.number = number
+        args.parallel = parallel
+        # Set up output path for each run
+        cur_output_path = os.path.join(output_path, f'parallel_{parallel}_number_{number}')
+        os.makedirs(cur_output_path, exist_ok=True)
+        # Start the benchmark
+        metrics_result = run_one_benchmark(args, output_path=cur_output_path)
+        # Save the results
+        results.append(metrics_result)
+        # Sleep between runs to avoid overwhelming the server
+        if i < len(number_list) - 1:
+            logger.info(f'Sleeping for {args.sleep_interval} seconds before the next run...')
+            time.sleep(args.sleep_interval)
+    # Analyze results
+    print_summary(results, args.model_id)
+    return results
+
+
 def run_perf_benchmark(args):
+    # Check if args is a dictionary or Namespace
     if isinstance(args, dict):
         args = Arguments(**args)
     elif isinstance(args, Namespace):
@@ -22,22 +76,41 @@ def run_perf_benchmark(args):
     if args.seed is not None:
         seed_everything(args.seed)
 
-    # Setup logger and output
-    args.outputs_dir = get_output_path(args)
-    configure_logging(args.debug, os.path.join(args.outputs_dir, 'benchmark.log'))
+    # Initialize output directory
+    output_path = get_output_path(args)
+    configure_logging(args.debug, os.path.join(output_path, 'benchmark.log'))
 
-    logger.info('Starting benchmark...')
-    logger.info(args)
+    # Initialize wandb and swanlab
+    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)
+    elif visualizer == 'swanlab':
+        init_swanlab(args)
 
-    if platform.system() == 'Windows':
-        asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
-
-    loop = asyncio.get_event_loop()
-    if platform.system() != 'Windows':
-        add_signal_handlers(loop)
-    loop.run_until_complete(benchmark(args))
+    # Initialize local server if needed
+    if args.api.startswith('local'):
+        #  start local server
+        server = threading.Thread(target=start_app, args=(copy.deepcopy(args), ), daemon=True)
+        server.start()
+    # Start benchmark
+    if len(args.number) == 1:
+        return run_one_benchmark(args, output_path=output_path)
+    else:
+        return run_multi_benchmark(args, output_path=output_path)
 
 
 if __name__ == '__main__':
     args = Arguments.from_args(parse_args())
-    run_perf_benchmark(args)
+    metrics_result, percentile_result = run_perf_benchmark(args)
+    print(metrics_result)
+    print(percentile_result)

evalscope/perf/plugin/__init__.py

@@ -1,2 +1,3 @@
-from evalscope.perf.plugin.api import *
-from evalscope.perf.plugin.datasets import *
+from .api import *
+from .datasets import *
+from .registry import ApiRegistry, DatasetRegistry

evalscope/perf/plugin/api/__init__.py

@@ -1,3 +1,4 @@
-from evalscope.perf.plugin.api.custom_api import CustomPlugin
-from evalscope.perf.plugin.api.dashscope_api import DashScopeApiPlugin
-from evalscope.perf.plugin.api.openai_api import OpenaiPlugin
+from .base import ApiPluginBase
+from .custom_api import CustomPlugin
+from .dashscope_api import DashScopeApiPlugin
+from .openai_api import OpenaiPlugin

evalscope/perf/plugin/api/base.py

@@ -1,16 +1,19 @@
+import aiohttp
 from abc import abstractmethod
-from typing import Any, Dict, List, Tuple
+from typing import Any, AsyncGenerator, Dict, List, Tuple
 
 from evalscope.perf.arguments import Arguments
+from evalscope.perf.utils.benchmark_util import BenchmarkData
 
 
 class ApiPluginBase:
 
-    def __init__(self, model_path: str) -> None:
-        self.model_path = model_path
+    def __init__(self, param: Arguments) -> None:
+        self.param = param
+        self.model_path = param.tokenizer_path
 
     @abstractmethod
-    def build_request(self, messages: List[Dict], param: Arguments) -> Dict:
+    def build_request(self, messages: List[Dict], param: Arguments = None) -> Dict:
         """Build a api request body.
 
         Args:
@@ -26,19 +29,36 @@ 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).
         """
         raise NotImplementedError
 
+    @abstractmethod
+    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:
+            client_session: The aiohttp client session
+            url: The request URL
+            headers: The request headers
+            body: The request body
+
+        Returns:
+            BenchmarkData: The benchmark data including response and timing info.
+        """
+        raise NotImplementedError
+
     @staticmethod
     def replace_values(input_json: Any, model: str, prompt: str):
         if isinstance(input_json, dict):