Share aiohttp.ClientSessions per worker

Slightly refactor `openAIModelServerClient` to add a new method,
`process_request_with_session`, that accepts a custom
`ReusableHTTPClientSession` per request, which allows the caller to
reuse an HTTP client session per worker.

The previous method, `process_request`, is made to create a fresh HTTP
client session then call `process_request_with_session`, preserving the
previous behavior.

Prior to this commit, a new `aiohttp.ClientSession` is created for each
request. Not only is this inefficient and lowers throughput, on certain
environments, it also leads to inotify watch issues:

    aiodns - WARNING - Failed to create DNS resolver channel with
    automatic monitoring of resolver configuration changes. This usually
    means the system ran out of inotify watches. Falling back to socket
    state callback. Consider increasing the system inotify watch limit:
    Failed to initialize c-ares channel

Indeed, because each DNS resolver is created for a new `ClientSession`,
creating tons of new `ClientSession`s causes eventual inotify watch
exhaustion. Sharing `ClientSession`s solves this issue.

Relevant links:

- https://docs.aiohttp.org/en/stable/http_request_lifecycle.html
- https://stackoverflow.com/questions/62707369/one-aiohttp-clientsession-per-thread
- home-assistant/core#144457 (comment)

Relevant PR: kubernetes-sigs#247
(doesn't address the issue of worker sharing).

Author: diamondburned

inference_perf/client/modelserver/__init__.py

@@ -11,10 +11,16 @@
 # 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 ModelServerClient
+from .base import ModelServerClient, ReusableHTTPClientSession
 from .mock_client import MockModelServerClient
 from .vllm_client import vLLMModelServerClient
 from .sglang_client import SGlangModelServerClient
 
 
-__all__ = ["ModelServerClient", "MockModelServerClient", "vLLMModelServerClient", "SGlangModelServerClient"]
+__all__ = [
+    "ModelServerClient",
+    "ReusableHTTPClientSession",
+    "MockModelServerClient",
+    "vLLMModelServerClient",
+    "SGlangModelServerClient",
+]

inference_perf/client/modelserver/base.py

@@ -12,11 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from abc import ABC, abstractmethod
-from typing import List, Optional, Tuple
+from typing import List, Optional, Tuple, Any
 from inference_perf.client.metricsclient.base import MetricsMetadata
 from inference_perf.config import APIConfig, APIType
-
 from inference_perf.apis import InferenceAPIData
+import aiohttp
 
 
 class ModelServerPrometheusMetric:
@@ -87,10 +87,38 @@ def get_supported_apis(self) -> List[APIType]:
         raise NotImplementedError
 
     @abstractmethod
-    async def process_request(self, data: InferenceAPIData, stage_id: int, scheduled_time: float) -> None:
+    async def process_request(
+        self, data: InferenceAPIData, stage_id: int, scheduled_time: float, *args: Any, **kwargs: Any
+    ) -> None:
         raise NotImplementedError
 
     @abstractmethod
     def get_prometheus_metric_metadata(self) -> PrometheusMetricMetadata:
         # assumption: all metrics clients have metrics exported in Prometheus format
         raise NotImplementedError
+
+
+class ReusableHTTPClientSession:
+    """
+    A wrapper for aiohttp.ClientSession to allow for reusable sessions.
+    This is useful for sharing among many HTTP clients.
+    """
+
+    def __init__(self, session: aiohttp.ClientSession, dont_close: bool = False) -> None:
+        self.session = session
+        self.dont_close = dont_close
+
+    def make_dont_close(self) -> "ReusableHTTPClientSession":
+        return ReusableHTTPClientSession(session=self.session, dont_close=True)
+
+    async def close(self) -> None:
+        if self.dont_close:
+            self.dont_close = False
+            return
+        await self.session.close()
+
+    async def __aenter__(self) -> None:
+        pass
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:  # type: ignore[no-untyped-def]
+        await self.close()

inference_perf/client/modelserver/openai_client.py

@@ -17,7 +17,7 @@
 from inference_perf.config import APIConfig, APIType, CustomTokenizerConfig
 from inference_perf.apis import InferenceAPIData, InferenceInfo, RequestLifecycleMetric, ErrorResponseInfo
 from inference_perf.utils import CustomTokenizer
-from .base import ModelServerClient, PrometheusMetricMetadata
+from .base import ModelServerClient, PrometheusMetricMetadata, ReusableHTTPClientSession
 from typing import List, Optional
 import aiohttp
 import asyncio
@@ -30,6 +30,8 @@
 
 
 class openAIModelServerClient(ModelServerClient):
+    _session: aiohttp.ClientSession | None = None
+
     def __init__(
         self,
         metrics_collector: RequestDataCollector,
@@ -71,6 +73,17 @@ def __init__(
         self.tokenizer = CustomTokenizer(tokenizer_config)
 
     async def process_request(self, data: InferenceAPIData, stage_id: int, scheduled_time: float) -> None:
+        session = self.new_reusable_session()
+        async with session:
+            await self.process_request_with_session(data, stage_id, scheduled_time, session)
+
+    async def process_request_with_session(
+        self,
+        data: InferenceAPIData,
+        stage_id: int,
+        scheduled_time: float,
+        session: ReusableHTTPClientSession,
+    ) -> None:
         payload = data.to_payload(
             model_name=self.model_name,
             max_tokens=self.max_completion_tokens,
@@ -87,14 +100,10 @@ async def process_request(self, data: InferenceAPIData, stage_id: int, scheduled
 
         request_data = json.dumps(payload)
 
-        timeout = aiohttp.ClientTimeout(total=self.timeout) if self.timeout else aiohttp.helpers.sentinel
-
-        async with aiohttp.ClientSession(
-            connector=aiohttp.TCPConnector(limit=self.max_tcp_connections), timeout=timeout
-        ) as session:
+        async with session.make_dont_close():
             start = time.perf_counter()
             try:
-                async with session.post(self.uri + data.get_route(), headers=headers, data=request_data) as response:
+                async with session.session.post(self.uri + data.get_route(), headers=headers, data=request_data) as response:
                     response_info = await data.process_response(
                         response=response, config=self.api_config, tokenizer=self.tokenizer
                     )
@@ -138,6 +147,14 @@ async def process_request(self, data: InferenceAPIData, stage_id: int, scheduled
                     )
                 )
 
+    def new_reusable_session(self) -> ReusableHTTPClientSession:
+        return ReusableHTTPClientSession(
+            aiohttp.ClientSession(
+                timeout=aiohttp.ClientTimeout(total=self.timeout) if self.timeout else aiohttp.helpers.sentinel,
+                connector=aiohttp.TCPConnector(limit=self.max_tcp_connections),
+            )
+        )
+
     def get_supported_apis(self) -> List[APIType]:
         return []
 

inference_perf/loadgen/load_generator.py

@@ -17,7 +17,8 @@
 from .load_timer import LoadTimer, ConstantLoadTimer, PoissonLoadTimer, TraceReplayLoadTimer
 from inference_perf.datagen import DataGenerator
 from inference_perf.apis import InferenceAPIData
-from inference_perf.client.modelserver import ModelServerClient
+from inference_perf.client.modelserver import ModelServerClient, ReusableHTTPClientSession
+from inference_perf.client.modelserver.openai_client import openAIModelServerClient
 from inference_perf.circuit_breaker import get_circuit_breaker
 from inference_perf.config import LoadConfig, LoadStage, LoadType, StageGenType, TraceFormat
 from asyncio import (
@@ -83,6 +84,10 @@ async def loop(self) -> None:
         item = None
         timeout = 0.5
 
+        session: ReusableHTTPClientSession | None = None
+        if isinstance(self.client, openAIModelServerClient):
+            session = self.client.new_reusable_session()
+
         while not self.stop_signal.is_set():
             while self.request_phase.is_set() and not self.cancel_signal.is_set():
                 await semaphore.acquire()
@@ -120,7 +125,12 @@ async def schedule_client(
                         with self.active_requests_counter.get_lock():
                             self.active_requests_counter.value += 1
                             inflight = True
-                        await self.client.process_request(request_data, stage_id, request_time)
+
+                        if isinstance(self.client, openAIModelServerClient):
+                            assert session
+                            await self.client.process_request_with_session(request_data, stage_id, request_time, session)
+                        else:
+                            await self.client.process_request(request_data, stage_id, request_time)
                     except CancelledError:
                         pass
                     finally:
@@ -151,6 +161,8 @@ async def schedule_client(
                 logger.debug(f"[Worker {self.id}] waiting for next phase")
                 self.request_phase.wait()
 
+        if session:
+            await session.close()
         logger.debug(f"[Worker {self.id}] stopped")
 
     def run(self) -> None: