evalscope 0.17.0__py3-none-any.whl → 0.17.1__py3-none-any.whl

evalscope/perf/http_client.py

@@ -1,13 +1,13 @@
 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.utils.logger import get_logger
+from .arguments import Arguments
+
+if TYPE_CHECKING:
+    from .plugin.api.base import ApiPluginBase
 
 logger = get_logger()
 
@@ -17,95 +17,48 @@ 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
         self.client = aiohttp.ClientSession(
             timeout=aiohttp.ClientTimeout(connect=self.connect_timeout, sock_read=self.read_timeout),
             trace_configs=[self._create_trace_config()] if args.debug else [])
 
-    def _create_trace_config(self):
-        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 __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'))
+    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 post(self, body):
-        headers = {'Content-Type': 'application/json', **self.headers}
+        """Send POST request and delegate response handling to API plugin.
+        Yields:
+            Tuple[bool, int, str]: (is_error, status_code, response_data)
+        """
         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
-        except asyncio.TimeoutError:
+            # 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
+        except asyncio.TimeoutError as e:
             logger.error(
-                f'TimeoutError: connect_timeout: {self.connect_timeout}, read_timeout: {self.read_timeout}. Please set longger timeout.'  # noqa: E501
+                f'TimeoutError: connect_timeout: {self.connect_timeout}, read_timeout: {self.read_timeout}. Please set longer timeout.'  # noqa: E501
             )
-            yield (True, None, 'Timeout')
+            yield (True, None, str(e))
         except (aiohttp.ClientConnectorError, Exception) as e:
             logger.error(e)
-            yield (True, None, e)
+            yield (True, None, str(e))
 
     @staticmethod
     async def on_request_start(session, context, params: aiohttp.TraceRequestStartParams):
@@ -136,25 +89,16 @@ class AioHttpClient:
         logger.debug(f'Request received: <{method=},  {url=}, {truncated_chunk=}>')
 
 
-async def test_connection(args: Arguments) -> bool:
+async def test_connection(args: Arguments, api_plugin: 'ApiPluginBase') -> bool:
     is_error = True
     start_time = time.perf_counter()
 
     async def attempt_connection():
-        client = AioHttpClient(args)
+        client = AioHttpClient(args, api_plugin)
         async with client:
-            if args.apply_chat_template:
-                request = {
-                    'messages': [{
-                        'role': 'user',
-                        'content': 'hello'
-                    }],
-                    'model': args.model,
-                    'max_tokens': 10,
-                    'stream': args.stream
-                }
-            else:
-                request = {'prompt': 'hello', 'model': args.model, 'max_tokens': 10}
+            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
 

evalscope/perf/main.py

@@ -57,8 +57,8 @@ def run_multi_benchmark(args: Arguments, output_path: str = None):
         results.append(metrics_result)
         # Sleep between runs to avoid overwhelming the server
         if i < len(number_list) - 1:
-            logger.info('Sleeping for 5 seconds before the next run...')
-            time.sleep(5)
+            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

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,18 @@
+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
 
 
 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:
@@ -39,6 +41,22 @@ 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]:
+        """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
+
+        Yields:
+            Tuple[bool, int, str]: (is_error, status_code, response_data)
+        """
+        raise NotImplementedError
+
     @staticmethod
     def replace_values(input_json: Any, model: str, prompt: str):
         if isinstance(input_json, dict):

evalscope/perf/plugin/api/custom_api.py

@@ -1,5 +1,6 @@
+import aiohttp
 import json
-from typing import Any, Dict, Iterator, List
+from typing import Any, AsyncGenerator, Dict, List, Tuple, Union
 
 from evalscope.perf.arguments import Arguments
 from evalscope.perf.plugin.api.base import ApiPluginBase
@@ -11,82 +12,238 @@ logger = get_logger()
 
 @register_api('custom')
 class CustomPlugin(ApiPluginBase):
-    """Support tensorrt-llm triton server
+    """Support custom API implementations.
+
+    This class serves as a template for users to implement their own API plugins.
+    By extending this class, users can connect to any LLM API with custom request
+    and response formats.
     """
 
-    def __init__(self, mode_path: str):
-        """Init the plugin
+    def __init__(self, param: Arguments):
+        """Initialize the plugin with the provided parameters.
 
         Args:
-            mode_path (str): The model path, we use the tokenizer
-                weight in the model to calculate the number of the
-                input and output tokens.
+            param (Arguments): Configuration parameters for the plugin, including:
+                - tokenizer_path: Path to the tokenizer for token counting
+                - model: Name of the model to use
+                - Other request parameters like max_tokens, temperature, etc.
         """
-        super().__init__(model_path=mode_path)
-        if mode_path is not None:
+        super().__init__(param=param)
+        if param.tokenizer_path is not None:
             from modelscope import AutoTokenizer
-            self.tokenizer = AutoTokenizer.from_pretrained(mode_path)
+            self.tokenizer = AutoTokenizer.from_pretrained(param.tokenizer_path)
         else:
             self.tokenizer = None
 
-    def build_request(self, messages: List[Dict], param: Arguments) -> Dict:
-        """Build the openai format request based on prompt, dataset
+    def build_request(self, messages: Union[List[Dict], str], param: Arguments = None) -> Dict:
+        """Build a custom API request body based on the input messages and parameters.
 
-        Args:
-            message (Dict): The basic message to generator query.
-            param (Arguments): The query parameters.
+        This method formats the input messages into the expected request format
+        for your custom API.
 
-        Raises:
-            Exception: NotImplemented
+        Args:
+            messages (Union[List[Dict], str]): The input messages to include in the request.
+                Could be a list of message dictionaries (for chat models) or a string (for completion models).
+            param (Arguments, optional): Request parameters. Defaults to self.param.
 
         Returns:
-            Dict: The request body. None if prompt format is error.
+            Dict: A properly formatted request body for your custom API.
         """
+        param = param or self.param
         try:
-            query = json.loads(param.query_template)
-            ApiPluginBase.replace_values(query, param.model, messages[0]['content'])
-            return query
+            # Create a default query format if no template is provided
+            if isinstance(messages, str):
+                query = {'input_text': messages}
+            else:
+                query = {'messages': messages}
+
+            # Add model parameters to the request
+            return self._add_parameters_to_request(query, param)
         except Exception as e:
             logger.exception(e)
-            logger.error('Prompt: %s invalidate!' % messages)
             return None
 
-    def parse_responses(self, responses, request: Any = None, **kwargs) -> Dict:
-        """Parser responses and return number of request and response tokens.
-           sample of the output delta:
-           {"id":"4","object":"chat.completion.chunk","created":1714030870,"model":"llama3","choices":[{"index":0,"delta":{"role":"assistant","content":""},"logprobs":null,"finish_reason":null}]}
+    def _add_parameters_to_request(self, payload: Dict, param: Arguments) -> Dict:
+        """Add model parameters to the request payload.
+
+        This helper method adds various parameters like temperature, max_tokens, etc.
+        to the request based on what your custom API supports.
+
+        Args:
+            payload (Dict): The base request payload.
+            param (Arguments): The parameters to add.
+
+        Returns:
+            Dict: The request payload with added parameters.
+        """
+        # Add the model name
+        payload['model'] = param.model
+
+        # Add various parameters if they are provided
+        if param.max_tokens is not None:
+            payload['max_tokens'] = param.max_tokens
+        if param.temperature is not None:
+            payload['temperature'] = param.temperature
+        if param.top_p is not None:
+            payload['top_p'] = param.top_p
+        if param.top_k is not None:
+            payload['top_k'] = param.top_k
+        if param.stream is not None:
+            payload['stream'] = param.stream
+            payload['stream_options'] = {'include_usage': True}
 
+        # Add any extra arguments passed via command line
+        if param.extra_args is not None:
+            payload.update(param.extra_args)
+
+        return payload
+
+    def parse_responses(self, responses: List[str], request: Any = 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.
+        Different APIs may return this information in different formats, or you may need
+        to calculate it using a tokenizer.
 
         Args:
-            responses (List[bytes]): List of http response body, for stream output,
-                there are multiple responses, for general only one.
-            kwargs: (Any): The command line --parameter content.
+            responses (List[str]): List of API response strings.
+            request (Any, optional): The original request, which might be needed for token calculation.
+            **kwargs: Additional arguments.
+
         Returns:
-            Tuple: Return number of prompt token and number of completion tokens.
+            Tuple[int, int]: (input_tokens, output_tokens) - The number of tokens in the prompt and completion.
+        """
+        try:
+            # Example 1: Try to get token counts from the API response
+            last_response = json.loads(responses[-1])
+
+            # If the API provides token usage information
+            if 'usage' in last_response and last_response['usage']:
+                input_tokens = last_response['usage'].get('prompt_tokens', 0)
+                output_tokens = last_response['usage'].get('completion_tokens', 0)
+                return input_tokens, output_tokens
+
+            # Example 2: Calculate tokens using the tokenizer if no usage info is available
+            if self.tokenizer is not None:
+                input_text = ''
+                output_text = ''
+
+                # Extract input text from the request
+                if request and 'messages' in request:
+                    # For chat API
+                    input_text = ' '.join([msg['content'] for msg in request['messages']])
+                elif request and 'input_text' in request:
+                    # For completion API
+                    input_text = request['input_text']
+
+                # Extract output text from the response
+                for response in responses:
+                    js = json.loads(response)
+                    if 'choices' in js:
+                        for choice in js['choices']:
+                            if 'message' in choice and 'content' in choice['message']:
+                                output_text += choice['message']['content']
+                            elif 'text' in choice:
+                                output_text += choice['text']
+
+                # Count tokens
+                input_tokens = len(self.tokenizer.encode(input_text))
+                output_tokens = len(self.tokenizer.encode(output_text))
+                return input_tokens, output_tokens
+
+            # If no usage information and no tokenizer, raise an error
+            raise ValueError(
+                'Cannot determine token counts: no usage information in response and no tokenizer provided.')
+
+        except Exception as e:
+            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]:
+        """Process the HTTP request and handle the response.
+
+        This method handles sending the request to your API and processing the response,
+        including handling streaming responses if supported.
+
+        Args:
+            client_session (aiohttp.ClientSession): The aiohttp client session.
+            url (str): The API endpoint URL.
+            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
         """
-        full_response_content = ''
-        delta_contents = {}
-        input_tokens = None
-        output_tokens = None
-        for response in responses:
-            data = json.loads(response)
-            # {"context_logits":0.0,"cum_log_probs":0.0,"generation_logits":0.0,"model_name":"ensemble",
-            # "model_version":"1","output_log_probs":[0.0,0.0,0.0,0.0,0.0],"sequence_end":false,"sequence_id":0,"sequence_start":false,"text_output":"性"}
-            if 'text_output' in data:
-                if 0 in delta_contents:
-                    delta_contents[0].append(data['text_output'])
+        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:
-                    delta_contents[0] = [data['text_output']]
-        if input_tokens is None and output_tokens is None and self.tokenizer is not None:
-            input_tokens = 0
-            output_tokens = 0
-            for _, choice_contents in delta_contents.items():
-                full_response_content = ''.join([m for m in choice_contents])
-                input_tokens += len(self.tokenizer.encode(request['text_input']))
-                output_tokens += len(self.tokenizer.encode(full_response_content))
-        elif input_tokens is None and output_tokens is None:  # no usage info get.
-            input_tokens = 0
-            output_tokens = 0
-            logger.warning('No usage info get.')
-
-        return input_tokens, output_tokens
+                    # 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))
+
+
+if __name__ == '__main__':
+    # Example usage of the CustomPlugin
+    from dotenv import dotenv_values
+    env = dotenv_values('.env')
+
+    from evalscope.perf.arguments import Arguments
+    from evalscope.perf.main import run_perf_benchmark
+
+    args = Arguments(
+        model='qwen2.5-7b-instruct',
+        url='https://dashscope.aliyuncs.com/compatible-mode/v1/chat/completions',
+        api_key=env.get('DASHSCOPE_API_KEY'),
+        api='custom',  # Use the custom API plugin registered above
+        dataset='openqa',
+        number=1,
+        max_tokens=10)
+
+    run_perf_benchmark(args)

evalscope/perf/plugin/api/dashscope_api.py

@@ -13,17 +13,10 @@ logger = get_logger()
 @register_api('dashscope')
 class DashScopeApiPlugin(ApiPluginBase):
 
-    def __init__(self, mode_path: str):
-        """Init the plugin
+    def __init__(self, param: Arguments):
+        super().__init__(param)
 
-        Args:
-            mode_path (str): The model path, we use the tokenizer
-                weight in the model to calculate the number of the
-                input and output tokens.
-        """
-        super().__init__(model_path=mode_path)
-
-    def build_request(self, messages: List[Dict], param: Arguments) -> Dict:
+    def build_request(self, messages: List[Dict], param: Arguments = None) -> Dict:
         """Build the openai format request based on prompt, dataset
 
         Args:
@@ -36,6 +29,7 @@ class DashScopeApiPlugin(ApiPluginBase):
         Returns:
             Dict: The request body. None if prompt format is error.
         """
+        param = param or self.param
         try:
             if param.query_template is not None:
                 if param.query_template.startswith('@'):