chore: import upstream snapshot with attribution

python/mlc_llm/bench/api_endpoint.py

@@ -0,0 +1,452 @@
+"""MLC LLM bench backends"""
+
+import argparse
+import json
+import os
+import time
+import traceback
+from typing import Optional
+
+from typing_extensions import Self
+
+from mlc_llm.bench.request_record import Metrics, RequestRecord, ServerMetrics
+from mlc_llm.support import logging
+
+logger = logging.getLogger(__name__)
+
+
+class APIEndPoint:
+    """Manages the sending of requests to a specified API endpoint and gathers
+    inference statistics.
+    """
+
+    def __init__(self, include_server_metrics: bool = False) -> None:
+        self.include_server_metrics = include_server_metrics
+
+    async def __aenter__(self) -> Self:
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, tb) -> None:
+        pass
+
+    async def __call__(self, request: RequestRecord) -> RequestRecord:
+        raise NotImplementedError()
+
+
+class OpenAIChatEndPoint(APIEndPoint):
+    """The backend of sending HTTP requests in OpenAI API through "v1/chat/completions"."""
+
+    def __init__(
+        self,
+        host: str,
+        port: int,
+        timeout: Optional[float] = None,
+        include_server_metrics: bool = False,
+    ) -> None:
+        super().__init__(include_server_metrics=include_server_metrics)
+
+        import aiohttp
+
+        self.timeout = timeout
+        self.client: aiohttp.ClientSession = None
+        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:
+        import aiohttp
+
+        self.client = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(self.timeout))
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, tb) -> None:
+        await self.client.close()
+
+    async def __call__(self, request_record: RequestRecord) -> RequestRecord:
+        payload = request_record.chat_cmpl.model_dump()
+        if self.timeout is not None and "timeout" not in payload:
+            payload["timeout"] = self.timeout
+        if self.include_server_metrics:
+            if "stream_options" not in payload or payload["stream_options"] is None:
+                payload["stream_options"] = {"include_usage": True}
+            else:
+                payload["stream_options"]["include_usage"] = True
+        if (
+            request_record.chat_cmpl.debug_config is not None
+            and request_record.chat_cmpl.debug_config.ignore_eos
+        ):
+            payload["ignore_eos"] = True
+
+        generated_text = ""
+        first_chunk_output_str = ""
+        time_to_first_token_s = None
+        start_time = time.monotonic()
+        server_metrics = None
+
+        try:
+            async with self.client.post(self.url, json=payload, headers=self.headers) as response:
+                assert response.status == 200, await response.text()
+                if payload["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 not data["choices"]:
+                            continue
+                        delta = data["choices"][0]["delta"]
+                        content = delta.get("content", None)
+                        if content is not None and not time_to_first_token_s:
+                            time_to_first_token_s = time.monotonic() - start_time
+                            first_chunk_output_str = content
+                        if self.include_server_metrics and data["usage"] is not None:
+                            # fmt: off
+                            server_metrics = ServerMetrics(
+                                input_tokens=data["usage"]["extra"]["prompt_tokens"],
+                                prefill_tokens=data["usage"]["extra"]["prefill_tokens"],
+                                output_tokens=data["usage"]["extra"]["completion_tokens"],
+                                end_to_end_latency_s=data["usage"]["extra"]["end_to_end_latency_s"],
+                                prefill_tokens_per_s=data["usage"]["extra"]["prefill_tokens_per_s"],
+                                inter_token_latency_s=data["usage"]["extra"]["inter_token_latency_s"],
+                                time_per_output_token_s=1 / data["usage"]["extra"]["decode_tokens_per_s"],  # noqa: E501
+                                time_to_first_token_s=data["usage"]["extra"]["ttft_s"],
+                            )
+                            # fmt: on
+
+                        if content is not None:
+                            generated_text += content
+                else:
+                    data = await response.json()
+                    generated_text = data["choices"][0]["message"]["content"]
+                    if self.include_server_metrics and data["usage"] is not None:
+                        # fmt: off
+                        server_metrics = ServerMetrics(
+                            input_tokens=data["usage"]["extra"]["prompt_tokens"],
+                            prefill_tokens=data["usage"]["extra"]["prefill_tokens"],
+                            output_tokens=data["usage"]["extra"]["completion_tokens"],
+                            end_to_end_latency_s=data["usage"]["extra"]["end_to_end_latency_s"],
+                            prefill_tokens_per_s=data["usage"]["extra"]["prefill_tokens_per_s"],
+                            inter_token_latency_s=data["usage"]["extra"]["inter_token_latency_s"],
+                            time_per_output_token_s=1 / data["usage"]["extra"]["decode_tokens_per_s"],  # noqa: E501
+                            time_to_first_token_s=data["usage"]["extra"]["ttft_s"],
+                        )
+                        # fmt: on
+        except Exception:
+            error_msg = "API endpoint errored when sending request: " + traceback.format_exc()
+            logger.info(error_msg)
+            finish_time = time.monotonic()
+            request_record.output_str = generated_text
+            request_record.first_chunk_output_str = first_chunk_output_str
+            request_record.metrics = Metrics(
+                success=False,
+                start_time=start_time,
+                finish_time=finish_time,
+                end_to_end_latency_s=finish_time - start_time,
+                input_tokens=request_record.metrics.input_tokens,
+                time_to_first_token_s=time_to_first_token_s,
+                server_metrics=server_metrics,
+                exec_feature=request_record.metrics.exec_feature,
+            )
+            request_record.error_msg = error_msg
+            return request_record
+
+        finish_time = time.monotonic()
+        request_record.output_str = generated_text
+        request_record.first_chunk_output_str = first_chunk_output_str
+        success = True
+        error_msg = None
+        if len(generated_text) == 0:
+            success = False
+            error_msg = "Empty generated text."
+        request_record.metrics = Metrics(
+            success=success,
+            start_time=start_time,
+            finish_time=finish_time,
+            end_to_end_latency_s=finish_time - start_time,
+            input_tokens=request_record.metrics.input_tokens,
+            time_to_first_token_s=time_to_first_token_s,
+            server_metrics=server_metrics,
+            exec_feature=request_record.metrics.exec_feature,
+        )
+        request_record.error_msg = error_msg
+        return request_record
+
+
+class OpenAIEndPoint(APIEndPoint):
+    """The backend of sending HTTP requests in OpenAI API through "v1/completions"."""
+
+    def __init__(
+        self,
+        host: str,
+        port: int,
+        timeout: Optional[float] = None,
+        include_server_metrics: bool = False,
+        no_debug_config: bool = False,
+    ) -> None:
+        super().__init__(include_server_metrics=include_server_metrics)
+
+        import aiohttp
+
+        self.timeout = timeout
+        self.client: aiohttp.ClientSession = None
+        self.url = f"http://{host}:{port}/v1/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')}"
+        assert not include_server_metrics, (
+            '"include_server_metrics" only works for "openai-chat" endpoint for now'
+        )
+        self.no_debug_config = no_debug_config
+
+    async def __aenter__(self) -> Self:
+        import aiohttp
+
+        self.client = aiohttp.ClientSession()
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, tb) -> None:
+        await self.client.close()
+
+    async def __call__(self, request_record: RequestRecord) -> RequestRecord:
+        assert len(request_record.chat_cmpl.messages) == 1, (
+            'Endpoint "openai" does not support system prompt and multi-round conversation.'
+        )
+        assert isinstance(request_record.chat_cmpl.messages[0].content, str)
+        payload = {
+            "model": request_record.chat_cmpl.model,
+            "prompt": request_record.chat_cmpl.messages[0].content,
+            "temperature": request_record.chat_cmpl.temperature,
+            "top_p": request_record.chat_cmpl.top_p,
+            "max_tokens": request_record.chat_cmpl.max_tokens,
+            "stream": True,
+        }
+        if self.timeout is not None and "timeout" not in payload:
+            payload["timeout"] = self.timeout
+        if (
+            request_record.chat_cmpl.debug_config is not None
+            and request_record.chat_cmpl.debug_config.ignore_eos
+        ):
+            payload["ignore_eos"] = True
+            if not self.no_debug_config:
+                payload["debug_config"] = {"ignore_eos": True}
+
+        generated_text = ""
+        first_chunk_output_str = ""
+        time_to_first_token_s = None
+        start_time = time.monotonic()
+
+        try:
+            async with self.client.post(
+                self.url, json=payload, headers=self.headers, timeout=3600
+            ) as response:
+                assert response.status == 200, await response.text()
+                if payload["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 not data["choices"]:
+                            continue
+                        content = data["choices"][0]["text"]
+                        if content is not None and not time_to_first_token_s:
+                            time_to_first_token_s = time.monotonic() - start_time
+                            first_chunk_output_str = content
+                        if content is not None:
+                            generated_text += content
+                else:
+                    data = await response.json()
+                    generated_text = data["choices"][0]["message"]["content"]
+        except Exception:
+            error_msg = "API endpoint errored when sending request: " + traceback.format_exc()
+            logger.info(error_msg)
+            finish_time = time.monotonic()
+            request_record.output_str = generated_text
+            request_record.first_chunk_output_str = first_chunk_output_str
+            request_record.metrics = Metrics(
+                success=False,
+                start_time=start_time,
+                finish_time=finish_time,
+                end_to_end_latency_s=finish_time - start_time,
+                input_tokens=request_record.metrics.input_tokens,
+                time_to_first_token_s=time_to_first_token_s,
+                server_metrics=None,
+                exec_feature=request_record.metrics.exec_feature,
+            )
+            request_record.error_msg = error_msg
+            return request_record
+
+        finish_time = time.monotonic()
+        request_record.output_str = generated_text
+        request_record.first_chunk_output_str = first_chunk_output_str
+        success = True
+        error_msg = None
+        if len(generated_text) == 0:
+            success = False
+            error_msg = "Empty generated text."
+        request_record.metrics = Metrics(
+            success=success,
+            start_time=start_time,
+            finish_time=finish_time,
+            end_to_end_latency_s=finish_time - start_time,
+            input_tokens=request_record.metrics.input_tokens,
+            time_to_first_token_s=time_to_first_token_s,
+            server_metrics=None,
+            exec_feature=request_record.metrics.exec_feature,
+        )
+        request_record.error_msg = error_msg
+        return request_record
+
+
+class TensorRTLLMEndPoint(APIEndPoint):
+    """The backend of sending HTTP requests in TensorRT-LLM API."""
+
+    def __init__(self, host: str, port: int, timeout: Optional[float] = None) -> None:
+        super().__init__(include_server_metrics=False)
+
+        import aiohttp
+
+        self.timeout = timeout
+        self.client: aiohttp.ClientSession = None
+        self.url_stream = f"http://{host}:{port}/v2/models/ensemble/generate_stream"
+        self.url_no_stream = f"http://{host}:{port}/v2/models/ensemble/generate"
+
+    async def __aenter__(self) -> Self:
+        import aiohttp
+
+        self.client = aiohttp.ClientSession()
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, tb) -> None:
+        await self.client.close()
+
+    async def __call__(self, request_record: RequestRecord) -> RequestRecord:
+        assert len(request_record.chat_cmpl.messages) == 1
+        assert isinstance(request_record.chat_cmpl.messages[0].content, str)
+        payload = {
+            "accumulate_tokens": True,
+            "text_input": request_record.chat_cmpl.messages[0].content,
+            "temperature": (
+                max(request_record.chat_cmpl.temperature, 1e-5)
+                if request_record.chat_cmpl.temperature
+                else 1
+            ),
+            "top_p": request_record.chat_cmpl.top_p if request_record.chat_cmpl.top_p else 1,
+            "max_tokens": request_record.chat_cmpl.max_tokens,
+            "stream": request_record.chat_cmpl.stream,
+        }
+        if (
+            request_record.chat_cmpl.debug_config is not None
+            and request_record.chat_cmpl.debug_config.ignore_eos
+        ):
+            payload["min_length"] = payload["max_tokens"]
+        if self.timeout is not None and "timeout" not in payload:
+            payload["timeout"] = self.timeout
+
+        generated_text = ""
+        first_chunk_output_str = ""
+        url = self.url_stream if request_record.chat_cmpl.stream else self.url_no_stream
+        time_to_first_token_s = None
+        start_time = time.monotonic()
+
+        try:
+            async with self.client.post(url, json=payload) as response:
+                assert response.status == 200, await response.text()
+                if payload["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[5:].strip()
+                        data = json.loads(raw_data)
+                        delta = data["text_output"]
+                        if delta is None:
+                            continue
+
+                        if not time_to_first_token_s:
+                            time_to_first_token_s = time.monotonic() - start_time
+                            first_chunk_output_str = delta
+                        generated_text += delta
+                else:
+                    data = await response.json()
+                    generated_text = data["text_output"]
+        except Exception:
+            error_msg = "API endpoint errored when sending request: " + traceback.format_exc()
+            logger.info(error_msg)
+            finish_time = time.monotonic()
+            request_record.output_str = generated_text
+            request_record.first_chunk_output_str = first_chunk_output_str
+            request_record.metrics = Metrics(
+                success=False,
+                start_time=start_time,
+                finish_time=finish_time,
+                end_to_end_latency_s=finish_time - start_time,
+                input_tokens=request_record.metrics.input_tokens,
+                time_to_first_token_s=time_to_first_token_s,
+                exec_feature=request_record.metrics.exec_feature,
+            )
+            request_record.error_msg = error_msg
+            return request_record
+
+        finish_time = time.monotonic()
+        request_record.output_str = generated_text
+        request_record.first_chunk_output_str = first_chunk_output_str
+        success = True
+        error_msg = None
+        if len(generated_text) == 0:
+            success = False
+            error_msg = "Empty generated text."
+        request_record.metrics = Metrics(
+            success=success,
+            start_time=start_time,
+            finish_time=finish_time,
+            end_to_end_latency_s=finish_time - start_time,
+            input_tokens=request_record.metrics.input_tokens,
+            time_to_first_token_s=time_to_first_token_s,
+            exec_feature=request_record.metrics.exec_feature,
+        )
+        request_record.error_msg = error_msg
+        return request_record
+
+
+# Todo: APIEndPoint with AsyncOpenAI Python interface
+# class OpenAIPythonEndPoint(APIEndPoint):
+#     pass
+
+SUPPORTED_BACKENDS = [
+    "openai",
+    "openai-chat",
+    "mlc",
+    "sglang",
+    "tensorrt-llm",
+    "vllm",
+]
+
+
+def create_api_endpoint(args: argparse.Namespace) -> APIEndPoint:
+    """Create an API endpoint instance with regard to the specified endpoint kind."""
+    if args.api_endpoint in ["openai", "mlc", "sglang"]:
+        return OpenAIEndPoint(args.host, args.port, args.timeout, args.include_server_metrics)
+    if args.api_endpoint == "vllm":
+        return OpenAIEndPoint(
+            args.host,
+            args.port,
+            args.timeout,
+            include_server_metrics=False,
+            no_debug_config=True,
+        )
+    if args.api_endpoint == "openai-chat":
+        return OpenAIChatEndPoint(args.host, args.port, args.timeout, args.include_server_metrics)
+    if args.api_endpoint == "tensorrt-llm":
+        return TensorRTLLMEndPoint(args.host, args.port, args.timeout)
+    raise ValueError(f'Unrecognized endpoint "{args.api_endpoint}"')