[Bench] Defaults to aiohttp client, add ServerMetrics (#2527)

* [Bench] Defaults to aiohttp client

* Add ServerMetrics to summary

* Remove duplicate servermetric def

Author: yongwww

python/mlc_llm/bench/metrics.py

@@ -1,6 +1,6 @@
 """ MLC LLM bench Metrics"""
 import json
-from typing import Callable, Dict, List, Optional, Union
+from typing import Any, Callable, Dict, List, Optional, Union
 
 from pydantic import BaseModel
 
@@ -12,6 +12,19 @@
 logger = logging.getLogger(__name__)
 
 
+class ServerMetrics(BaseModel):
+    """The metrics from the server side."""
+
+    prompt_tokens: int
+    prefill_tokens: int
+    completion_tokens: int
+    decode_tokens_per_s: float
+    prefill_tokens_per_s: float
+    end_to_end_latency_s: float
+    inter_token_latency_s: float
+    ttft_s: Optional[float] = None
+
+
 class Metrics(BaseModel):
     """The list of metric keys"""
 
@@ -21,6 +34,7 @@ class Metrics(BaseModel):
     inter_token_latency_s: float
     decode_tokens_per_s: float
     ttft: Optional[float] = None
+    server_metrics: Optional[ServerMetrics] = None
 
 
 class MetricsProcessor:
@@ -87,13 +101,26 @@ def extract_metrics_from_request_records(
             assert prompt_tokens > 0 and completion_tokens >= 0, "Invalid prompt tokens"
             end_to_end_latency_s = metric.end_to_end_latency_s
             ttft = metric.ttft if metric.ttft is not None else 0
+            server_metric = None
+            if metric.server_metrics is not None:
+                server_metric = ServerMetrics(
+                    prompt_tokens=metric.server_metrics["prompt_tokens"],
+                    prefill_tokens=metric.server_metrics["prefill_tokens"],
+                    completion_tokens=metric.server_metrics["completion_tokens"],
+                    decode_tokens_per_s=metric.server_metrics["decode_tokens_per_s"],
+                    prefill_tokens_per_s=metric.server_metrics["prefill_tokens_per_s"],
+                    end_to_end_latency_s=metric.server_metrics["end_to_end_latency_s"],
+                    inter_token_latency_s=metric.server_metrics["inter_token_latency_s"],
+                    ttft_s=metric.server_metrics["ttft_s"],
+                )
             refined_metric = Metrics(
                 inter_token_latency_s=end_to_end_latency_s / completion_tokens,
-                decode_tokens_per_s=completion_tokens / (end_to_end_latency_s - ttft),
+                decode_tokens_per_s=(completion_tokens - 1) / (end_to_end_latency_s - ttft),
                 ttft=metric.ttft,
                 end_to_end_latency_s=end_to_end_latency_s,
                 prompt_tokens=prompt_tokens,
                 completion_tokens=completion_tokens,
+                server_metrics=server_metric,
             )
             result.append(refined_metric)
         return result
@@ -148,9 +175,7 @@ def criteria(metric: Metrics) -> bool:
         self.reset_metrics(filered_metrics)
         return filered_metrics
 
-    def generate_metrics_summary(
-        self, start_time: float, end_time: float
-    ) -> Dict[str, Union[int, float]]:
+    def generate_metrics_summary(self, start_time: float, end_time: float) -> Dict[str, Any]:
         """
         Computes summary statistics across all metrics collected.
 
@@ -170,16 +195,49 @@ def generate_metrics_summary(
         report : Dict
             A dictionary containing the summary statistics of the collected metrics.
         """
-        import pandas as pd  # pylint: disable=import-outside-toplevel,import-error
-
         if not self.all_metrics:
             return {}
 
-        metrics = self.all_metrics
-        df = pd.DataFrame([metric.model_dump() for metric in metrics])
+        # Generate the client metrics statistics
+        report = self._compute_metrics_statistics(self.all_metrics)
+        report["num_completed_requests"] = len(self.all_metrics)
+        total_tokens = sum(metric.completion_tokens for metric in self.all_metrics)
+        report["overall_output_throughput"] = total_tokens / (end_time - start_time)
+
+        # Generate the server metrics statistics
+        server_metrics = [
+            metric.server_metrics for metric in self.all_metrics if metric.server_metrics
+        ]
+        server_report = self._compute_metrics_statistics(server_metrics)
+        report["server_metrics"] = server_report
+
+        logger.info("Metrics Summary:\n%s", json.dumps(report, indent=4, default=str))
+        return report
+
+    def _compute_metrics_statistics(self, metrics: List[Union[Metrics, ServerMetrics]]) -> Dict:
+        """
+        Compute the statistics of the metrics.
+
+        Parameters
+        ----------
+        metrics : List[Union[Metrics, ServerMetrics]]
+            The list of metrics to get the statistics.
+
+        Returns
+        -------
+        report : Dict
+            The statistics of the metrics.
+        """
+        import pandas as pd  # pylint: disable=import-outside-toplevel,import-error
 
         report: Dict = {}
-        for key, _ in Metrics.model_fields.items():
+        if not metrics:
+            return report
+
+        df = pd.DataFrame([metric.model_dump() for metric in metrics])
+        for key, _ in metrics[0].model_fields.items():
+            if key == "server_metrics":
+                continue
             if key in df.columns:
                 series = df[key].dropna()
                 report[key] = {
@@ -192,11 +250,4 @@ def generate_metrics_summary(
                     "max": series.max(),
                     "stddev": series.std(),
                 }
-
-        report["num_completed_requests"] = len(metrics)
-        report["overall_output_throughput"] = df["completion_tokens"].sum() / (
-            end_time - start_time
-        )
-
-        logger.info("Metrics Summary:\n%s", json.dumps(report, indent=4, default=str))
         return report

python/mlc_llm/bench/prompts.py

@@ -58,8 +58,7 @@ def __init__(
                     assert "prompt" in json_line, "The prompt field is required in the JSONL file."
                     if "prompt_tokens" not in json_line:
                         json_line["prompt_tokens"] = self._count_tokens(json_line["prompt"])
-                    self.prompts.append(json.loads(line))
-                self.prompts = [json.loads(line) for line in file]
+                    self.prompts.append(json_line)
         else:
             if not prompts_path:
                 prompts_path = Path(__file__).parent / "prompts.txt"  # type: ignore

python/mlc_llm/bench/request.py

@@ -1,9 +1,9 @@
 """MLC LLM Bench Request"""
 import json
+import os
 import time
 from typing import Any, Dict, List, Optional
 
-import httpx
 from openai import AsyncOpenAI
 from pydantic import BaseModel
 from typing_extensions import Self
@@ -24,9 +24,10 @@ class RequestRecords(BaseModel):
     output: str
     end_to_end_latency_s: float
     ttft: Optional[float] = None
+    server_metrics: Optional[Dict] = None
 
 
-class OpenAIRequestSender:
+class OpenAIRequestSender:  # pylint: disable=too-many-instance-attributes
     """
     Manages the sending of requests to a specified API endpoint and gathers inference statistics.
 
@@ -40,20 +41,27 @@ class OpenAIRequestSender:
         Specifies if streaming should be enabled, default is True.
     timeout : Optional[float]
         The maximum duration in seconds for each request, default is 180.
+    client : Optional[Any]
+        The client to use for sending requests.
+    include_server_metrics : Optional[bool]
+        Specifies if server metrics should be included, default is False.
 
     Attributes
     ----------
     stats : dict
         Statistics about the performance.
     """
 
-    def __init__(
+    def __init__(  # pylint: disable=too-many-arguments
         self,
         host: Optional[str] = "127.0.0.1",
         port: Optional[int] = 8008,
         stream: Optional[bool] = None,
         timeout: Optional[float] = None,
+        client: Optional[Any] = None,
+        include_server_metrics: Optional[bool] = False,
     ) -> None:
+        import aiohttp  # pylint: disable=import-outside-toplevel,import-error
         from transformers import (  # pylint: disable=import-outside-toplevel,import-error
             LlamaTokenizerFast,
         )
@@ -63,75 +71,102 @@ def __init__(
         self.tokenizer = LlamaTokenizerFast.from_pretrained("hf-internal-testing/llama-tokenizer")
         self.prompt_generator = PromptsGenerator()
         self.request_records: List[RequestRecords] = []
-        self.client = AsyncOpenAI(
-            base_url=f"http://{host}:{port}/v1",
-            api_key="None",
-            http_client=httpx.AsyncClient(http2=True),
-        )
+        self.client = client if client else aiohttp.ClientSession()
+        self.include_server_metrics = include_server_metrics
+        self.url = f"http://{host}:{port}/v1/chat/completions"
+        self.headers = {"Content-Type": "application/json"}
+        if os.getenv("MLC_LLM_API_KEY"):
+            self.headers["Authorization"] = f"Bearer {os.getenv('MLC_LLM_API_KEY')}"
 
     async def __aenter__(self) -> Self:
         return self
 
     async def __aexit__(self, exc_type, exc_value, traceback) -> None:
         await self.client.close()
 
-    async def __call__(self, params: Dict[str, Any] = None) -> None:
-        """
-        Send a request to the deployed serving endpoint and collect request records.
-
-        Parameters
-        ----------
-        params : Dict[str, Any]
-            The parameters for the request.
-
-        Returns
-        -------
-        response : Union[Dict, None]
-            The JSON response from the server or None if an error occurs.
-        """
+    async def __call__(  # pylint: disable=too-many-locals, too-many-branches, too-many-statements
+        self, params: Dict[str, Any] = None
+    ) -> None:
         if "messages" not in params:
             prompt_tokens = 128
             if "prompt_tokens" in params:
                 prompt_tokens = params["prompt_tokens"]
             else:
                 logger.warning("A random prompt with %d tokens will be generated.", prompt_tokens)
-
             prompt = self.prompt_generator.generate_prompt(prompt_tokens)
             params["messages"] = [{"role": "system", "content": prompt}]
         else:
-            prompt = params["messages"][0]["content"]
+            prompt = params["messages"][-1]["content"]
         chat_params = self._get_chat_completion_params(params)
         if "stream" not in chat_params:
             chat_params["stream"] = self.stream
         if "timeout" not in chat_params:
             chat_params["timeout"] = self.timeout
+        if self.include_server_metrics:
+            if "stream_options" not in chat_params:
+                chat_params["stream_options"] = {"include_usage": True}
+            else:
+                chat_params["stream_options"]["include_usage"] = True
 
         total_request_time = 0
         generated_text = ""
         ttft = None
         start_time = time.monotonic()
-        # chat_params["stream_options"] = {"include_usage": True}
-        response = await self.client.chat.completions.create(**chat_params)
-
-        if chat_params["stream"]:
-            async for chunk in response:
-                if chunk.usage:
-                    logger.info(
-                        "Server Metrics:\n%s", json.dumps(chunk.usage.extra, indent=4, default=str)
-                    )
-                elif chunk.choices[0].delta.content is not None:
-                    if not ttft:
-                        ttft = time.monotonic() - start_time  # type: ignore
-                    generated_text += chunk.choices[0].delta.content
+        server_metrics = None
+
+        # AsyncOpenAI chat completion
+        if isinstance(self.client, AsyncOpenAI):
+            response = await self.client.chat.completions.create(**chat_params)
+            if chat_params["stream"]:
+                async for chunk in response:
+                    if chunk.usage:
+                        server_metrics = chunk.usage.extra
+                    elif chunk.choices[0].delta.content is not None:
+                        if not ttft:
+                            ttft = time.monotonic() - start_time  # type: ignore
+                        generated_text += chunk.choices[0].delta.content
+            else:
+                generated_text = response.choices[0].message.content
         else:
-            generated_text = response.choices[0].message.content
+            try:
+                async with self.client.post(
+                    self.url, json=chat_params, headers=self.headers
+                ) as response:
+                    if chat_params["stream"]:
+                        async for chunk in response.content:
+                            chunk = chunk.strip()
+                            if not chunk or chunk == b"\n":
+                                continue
+                            # Get rid of the prefix "data: " and suffix "\n"
+                            raw_data = chunk[6:].strip()
+                            if raw_data == b"[DONE]":
+                                continue
+                            data = json.loads(raw_data)
+                            if data["usage"] is not None:
+                                server_metrics = data["usage"]["extra"]
+                            if not data["choices"]:
+                                continue
+                            delta = data["choices"][0]["delta"]
+                            if delta.get("content", None):
+                                if not ttft:
+                                    ttft = time.monotonic() - start_time
+
+                            generated_text += delta["content"]
+                    else:
+                        data = await response.json()
+                        generated_text = data["choices"][0]["message"]["content"]
+            except Exception as e:  # pylint: disable=broad-except
+                logger.error("Error sending request: %s", str(e))
+                raise e
 
         total_request_time = time.monotonic() - start_time  # type: ignore
+
         req_rec = RequestRecords(
             input=prompt,
             output=generated_text,
             end_to_end_latency_s=total_request_time,
             ttft=ttft,
+            server_metrics=server_metrics,
         )
         self.request_records.append(req_rec)