Merge pull request #27 from SachinVarghese/vllm

Adding vLLM Client to inference perf runner

Author: k8s-ci-robot

.github/workflows/format.yml

@@ -13,10 +13,8 @@ jobs:
     steps:
       - name: Checkout Code
         uses: actions/checkout@v4
-      - name: Set up Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: '3.13'
+      - name: Set up Python PDM
+        uses: pdm-project/setup-pdm@v4
       - name: Do Linting and Type Checks
         run: |
           make check

inference_perf/client/__init__.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 from .base import ModelServerClient
 from .mock_client import MockModelServerClient
+from .vllm_client import vLLMModelServerClient
 
 
-__all__ = ["ModelServerClient", "MockModelServerClient"]
+__all__ = ["ModelServerClient", "MockModelServerClient", "vLLMModelServerClient"]

inference_perf/client/base.py

@@ -27,5 +27,5 @@ def set_report_generator(self, reportgen: ReportGenerator) -> None:
         self.reportgen = reportgen
 
     @abstractmethod
-    def process_request(self, data: InferenceData) -> None:
+    async def process_request(self, data: InferenceData) -> None:
         raise NotImplementedError

inference_perf/client/mock_client.py

@@ -12,17 +12,25 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from inference_perf.datagen import InferenceData
-from inference_perf.reportgen import ReportGenerator, Metric
+from inference_perf.reportgen import ReportGenerator, RequestMetric
 from .base import ModelServerClient
+import asyncio
 
 
 class MockModelServerClient(ModelServerClient):
-    def __init__(self, uri: str) -> None:
-        self.uri = uri
+    def __init__(self) -> None:
+        pass
 
     def set_report_generator(self, reportgen: ReportGenerator) -> None:
         self.reportgen = reportgen
 
-    def process_request(self, data: InferenceData) -> None:
+    async def process_request(self, data: InferenceData) -> None:
         print("Processing request - " + data.system_prompt)
-        self.reportgen.collect_metrics(Metric(name=data.system_prompt))
+        await asyncio.sleep(3)
+        self.reportgen.collect_request_metrics(
+            RequestMetric(
+                prompt_tokens=0,
+                output_tokens=0,
+                time_per_request=3,
+            )
+        )

inference_perf/client/vllm_client.py

@@ -0,0 +1,56 @@
+# Copyright 2025 The Kubernetes Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from inference_perf.datagen import InferenceData
+from inference_perf.reportgen import ReportGenerator, RequestMetric
+from .base import ModelServerClient
+from typing import Any
+import aiohttp
+import json
+import time
+
+
+class vLLMModelServerClient(ModelServerClient):
+    def __init__(self, uri: str, model_name: str) -> None:
+        self.model_name = model_name
+        self.uri = uri + "/v1/completions"
+        self.max_completion_tokens = 30
+
+    def set_report_generator(self, reportgen: ReportGenerator) -> None:
+        self.reportgen = reportgen
+
+    def _createPayload(self, data: InferenceData) -> dict[str, Any]:
+        return {"model": self.model_name, "prompt": data.system_prompt, "max_tokens": self.max_completion_tokens}
+
+    async def process_request(self, data: InferenceData) -> None:
+        payload = self._createPayload(data)
+        headers = {"Content-Type": "application/json"}
+        async with aiohttp.ClientSession() as session:
+            start = time.monotonic()
+            try:
+                async with session.post(self.uri, headers=headers, data=json.dumps(payload)) as response:
+                    if response.status == 200:
+                        content = await response.json()
+                        end = time.monotonic()
+                        usage = content["usage"]
+                        self.reportgen.collect_request_metrics(
+                            RequestMetric(
+                                prompt_tokens=usage["prompt_tokens"],
+                                output_tokens=usage["completion_tokens"],
+                                time_per_request=end - start,
+                            )
+                        )
+                    else:
+                        print(await response.text())
+            except aiohttp.ClientConnectorError as e:
+                print("vLLM Server connection error:\n", str(e))

inference_perf/loadgen/load_generator.py

@@ -15,6 +15,7 @@
 from .load_timer import LoadTimer, ConstantLoadTimer, PoissonLoadTimer
 from inference_perf.datagen import DataGenerator
 from inference_perf.client import ModelServerClient
+from asyncio import TaskGroup, sleep
 import time
 
 
@@ -35,13 +36,20 @@ def __init__(self, datagen: DataGenerator, load_type: LoadType, rate: float, dur
         else:
             raise
 
-    def run(self, client: ModelServerClient) -> None:
-        print("Run started")
+    async def run(self, client: ModelServerClient) -> None:
         start_time = time.time()
         end_time = start_time + self.duration
-        for _, (data, time_index) in enumerate(zip(self.datagen.get_data(), self.timer.start_timer(start_time), strict=True)):
-            if time_index < end_time:
-                client.process_request(data)
-            else:
-                print("Run complete")
-                break
+        print("Run started")
+        async with TaskGroup() as tg:
+            for _, (data, time_index) in enumerate(
+                zip(self.datagen.get_data(), self.timer.start_timer(start_time), strict=True)
+            ):
+                now = time.time()
+                if time_index < end_time and now < end_time:
+                    if time_index > now:
+                        await sleep(time_index - time.time())
+                    tg.create_task(client.process_request(data))
+                    continue
+                else:
+                    break
+        print("Run completed")

inference_perf/loadgen/load_timer.py

@@ -48,7 +48,7 @@ def start_timer(self, initial: Optional[float] = None) -> Generator[float, None,
 
         # Given a rate, yield a time to wait before the next request
         while True:
-            next_time += self._rand.exponential(1 / self._rate)
+            next_time += self._rand.uniform(0, 1 / self._rate)
             yield next_time
 
 
@@ -73,7 +73,6 @@ def start_timer(self, initial: Optional[float] = None) -> Generator[float, None,
 
             # Schedule the requests over the next second
             timer = ConstantLoadTimer(req_count)
-            times = timer.start_timer(next_time)
             for _ in range(req_count):
-                next_time = next(times)
+                next_time = next(timer.start_timer(next_time))
                 yield next_time

inference_perf/main.py

@@ -13,8 +13,9 @@
 # limitations under the License.
 from inference_perf.loadgen import LoadGenerator, LoadType
 from inference_perf.datagen import MockDataGenerator
-from inference_perf.client import ModelServerClient, MockModelServerClient
+from inference_perf.client import ModelServerClient, vLLMModelServerClient
 from inference_perf.reportgen import ReportGenerator, MockReportGenerator
+import asyncio
 
 
 class InferencePerfRunner:
@@ -25,15 +26,15 @@ def __init__(self, client: ModelServerClient, loadgen: LoadGenerator, reportgen:
         self.client.set_report_generator(self.reportgen)
 
     def run(self) -> None:
-        self.loadgen.run(self.client)
+        asyncio.run(self.loadgen.run(self.client))
 
     def generate_report(self) -> None:
-        self.reportgen.generate_report()
+        asyncio.run(self.reportgen.generate_report())
 
 
 def main_cli() -> None:
     # Define Model Server Client
-    client = MockModelServerClient(uri="0.0.0.0:0")
+    client = vLLMModelServerClient(uri="http://0.0.0.0:8000", model_name="openai-community/gpt2")
 
     # Define LoadGenerator
     loadgen = LoadGenerator(MockDataGenerator(), LoadType.CONSTANT, rate=2, duration=5)

inference_perf/reportgen/__init__.py

@@ -11,8 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from .base import ReportGenerator, Metric
+from .base import ReportGenerator, RequestMetric
 from .mock_reportgen import MockReportGenerator
 
 
-__all__ = ["ReportGenerator", "Metric", "MockReportGenerator"]
+__all__ = ["ReportGenerator", "RequestMetric", "MockReportGenerator"]

inference_perf/reportgen/base.py

@@ -16,8 +16,17 @@
 from typing import Tuple
 
 
-class Metric(BaseModel):
-    name: str
+class MetricsSummary(BaseModel):
+    total_requests: int
+    avg_prompt_tokens: float
+    avg_output_tokens: float
+    avg_time_per_request: float
+
+
+class RequestMetric(BaseModel):
+    prompt_tokens: int
+    output_tokens: int
+    time_per_request: float
 
 
 class ReportGenerator(ABC):
@@ -26,9 +35,9 @@ def __init__(self, *args: Tuple[int, ...]) -> None:
         pass
 
     @abstractmethod
-    def collect_metrics(self, metric: Metric) -> None:
+    def collect_request_metrics(self, metric: RequestMetric) -> None:
         raise NotImplementedError
 
     @abstractmethod
-    def generate_report(self) -> None:
+    async def generate_report(self) -> None:
         raise NotImplementedError