Share aiohttp.ClientSessions per worker

Slightly refactor `openAIModelServerClient` to accept a custom
`aiohttp.ClientSession` per request, which allows us to use exactly 1
client session per worker.

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/base.py

@@ -12,11 +12,12 @@
 # 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
+import copy
 
 
 class ModelServerPrometheusMetric:
@@ -87,10 +88,35 @@ 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 dont_close_if(self, dont_close: bool = True) -> "ReusableHTTPClientSession":
+        return ReusableHTTPClientSession(session=self.session, dont_close=dont_close)
+
+    async def __aenter__(self) -> None:
+        pass
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:  # type: ignore[no-untyped-def]
+        if self.dont_close:
+            self.dont_close = False
+            return
+        await self.session.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,
@@ -70,7 +72,24 @@ def __init__(
             tokenizer_config = CustomTokenizerConfig(pretrained_model_name_or_path=self.model_name)
         self.tokenizer = CustomTokenizer(tokenizer_config)
 
-    async def process_request(self, data: InferenceAPIData, stage_id: int, scheduled_time: float) -> None:
+    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),
+            )
+        )
+
+    async def process_request(
+        self,
+        data: InferenceAPIData,
+        stage_id: int,
+        scheduled_time: float,
+        session: Optional[ReusableHTTPClientSession] = None,
+    ) -> None:
+        reusing_session = session is not None
+        session = session or self.new_reusable_session()
+
         payload = data.to_payload(
             model_name=self.model_name,
             max_tokens=self.max_completion_tokens,
@@ -87,14 +106,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.dont_close_if(reusing_session):
             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
                     )

inference_perf/loadgen/load_generator.py

@@ -16,6 +16,7 @@
 from inference_perf.datagen import DataGenerator
 from inference_perf.apis import InferenceAPIData
 from inference_perf.client.modelserver import ModelServerClient
+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
 from asyncio import (
@@ -29,7 +30,7 @@
     set_event_loop_policy,
     get_event_loop,
 )
-from typing import List, Tuple, TypeAlias, Optional
+from typing import List, Tuple, TypeAlias, Optional, Any
 from types import FrameType
 import time
 import multiprocessing as mp
@@ -81,6 +82,10 @@ async def loop(self) -> None:
         item = None
         timeout = 0.5
 
+        session: Any = None
+        if issubclass(type(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()
@@ -118,7 +123,7 @@ 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)
+                        await self.client.process_request(request_data, stage_id, request_time, session=session)
                     except CancelledError:
                         pass
                     finally:
@@ -149,6 +154,7 @@ async def schedule_client(
                 logger.debug(f"[Worker {self.id}] waiting for next phase")
                 self.request_phase.wait()
 
+        await session.close()
         logger.debug(f"[Worker {self.id}] stopped")
 
     def run(self) -> None: