Share aiohttp.ClientSessions per worker

diamondburned · diamondburned · commit 39dbaea053de · 2025-11-19T16:26:03.000-08:00
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).
diff --git a/inference_perf/client/modelserver/__init__.py b/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, ModelServerClientSession
 from .mock_client import MockModelServerClient
 from .vllm_client import vLLMModelServerClient
 from .sglang_client import SGlangModelServerClient
 
 
-__all__ = ["ModelServerClient", "MockModelServerClient", "vLLMModelServerClient", "SGlangModelServerClient"]
+__all__ = [
+    "ModelServerClient",
+    "ModelServerClientSession",
+    "MockModelServerClient",
+    "vLLMModelServerClient",
+    "SGlangModelServerClient",
+]
diff --git a/inference_perf/client/modelserver/base.py b/inference_perf/client/modelserver/base.py
@@ -15,7 +15,6 @@
 from typing import List, Optional, Tuple
 from inference_perf.client.metricsclient.base import MetricsMetadata
 from inference_perf.config import APIConfig, APIType
-
 from inference_perf.apis import InferenceAPIData
 
 
@@ -82,6 +81,9 @@ def __init__(self, api_config: APIConfig, timeout: Optional[float] = None, *args
         self.api_config = api_config
         self.timeout = timeout
 
+    def new_session(self) -> "ModelServerClientSession":
+        return ModelServerClientSession(self)
+
     @abstractmethod
     def get_supported_apis(self) -> List[APIType]:
         raise NotImplementedError
@@ -94,3 +96,14 @@ async def process_request(self, data: InferenceAPIData, stage_id: int, scheduled
     def get_prometheus_metric_metadata(self) -> PrometheusMetricMetadata:
         # assumption: all metrics clients have metrics exported in Prometheus format
         raise NotImplementedError
+
+
+class ModelServerClientSession:
+    def __init__(self, client: ModelServerClient):
+        self.client = client
+
+    async def process_request(self, data: InferenceAPIData, stage_id: int, scheduled_time: float) -> None:
+        await self.client.process_request(data, stage_id, scheduled_time)
+
+    async def close(self) -> None:  # noqa - subclasses optionally override this
+        pass
diff --git a/inference_perf/client/modelserver/openai_client.py b/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, ModelServerClientSession, PrometheusMetricMetadata
 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,73 +72,11 @@ 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:
-        payload = data.to_payload(
-            model_name=self.model_name,
-            max_tokens=self.max_completion_tokens,
-            ignore_eos=self.ignore_eos,
-            streaming=self.api_config.streaming,
-        )
-        headers = {"Content-Type": "application/json"}
-
-        if self.api_key:
-            headers["Authorization"] = f"Bearer {self.api_key}"
+    def new_session(self) -> "ModelServerClientSession":
+        return openAIModelServerClientSession(self)
 
-        if self.api_config.headers:
-            headers.update(self.api_config.headers)
-
-        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:
-            start = time.perf_counter()
-            try:
-                async with 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
-                    )
-                    response_content = await response.text()
-
-                    end_time = time.perf_counter()
-                    error = None
-                    if response.status != 200:
-                        error = ErrorResponseInfo(error_msg=response_content, error_type="Error response")
-
-                    self.metrics_collector.record_metric(
-                        RequestLifecycleMetric(
-                            stage_id=stage_id,
-                            request_data=request_data,
-                            response_data=response_content,
-                            info=response_info,
-                            error=error,
-                            start_time=start,
-                            end_time=end_time,
-                            scheduled_time=scheduled_time,
-                        )
-                    )
-            except Exception as e:
-                if isinstance(e, asyncio.exceptions.TimeoutError):
-                    logger.error("request timed out:", exc_info=True)
-                else:
-                    logger.error("error occured during request processing:", exc_info=True)
-                self.metrics_collector.record_metric(
-                    RequestLifecycleMetric(
-                        stage_id=stage_id,
-                        request_data=request_data,
-                        response_data=response_content if "response_content" in locals() else "",
-                        info=response_info if "response_info" in locals() else InferenceInfo(),
-                        error=ErrorResponseInfo(
-                            error_msg=str(e),
-                            error_type=type(e).__name__,
-                        ),
-                        start_time=start,
-                        end_time=time.perf_counter(),
-                        scheduled_time=scheduled_time,
-                    )
-                )
+    async def process_request(self, data: InferenceAPIData, stage_id: int, scheduled_time: float) -> None:
+        await self.new_session().process_request(data, stage_id, scheduled_time)
 
     def get_supported_apis(self) -> List[APIType]:
         return []
@@ -157,3 +97,83 @@ def get_supported_models(self) -> List[str]:
         except Exception as e:
             logger.error(f"Got exception retrieving supported models {e}")
             return []
+
+
+class openAIModelServerClientSession(ModelServerClientSession):
+    def __init__(self, client: openAIModelServerClient):
+        self.client = client
+        self.session = aiohttp.ClientSession(
+            timeout=aiohttp.ClientTimeout(total=client.timeout) if client.timeout else aiohttp.helpers.sentinel,
+            connector=aiohttp.TCPConnector(limit=client.max_tcp_connections),
+        )
+
+    async def process_request(self, data: InferenceAPIData, stage_id: int, scheduled_time: float) -> None:
+        payload = data.to_payload(
+            model_name=self.client.model_name,
+            max_tokens=self.client.max_completion_tokens,
+            ignore_eos=self.client.ignore_eos,
+            streaming=self.client.api_config.streaming,
+        )
+        headers = {"Content-Type": "application/json"}
+
+        if self.client.api_key:
+            headers["Authorization"] = f"Bearer {self.client.api_key}"
+
+        if self.client.api_config.headers:
+            headers.update(self.client.api_config.headers)
+
+        request_data = json.dumps(payload)
+        response_info: InferenceInfo | None = None
+        response_content: str | None = None
+
+        start = time.perf_counter()
+        try:
+            async with self.session.post(self.client.uri + data.get_route(), headers=headers, data=request_data) as response:
+                response_info = await data.process_response(
+                    response=response,
+                    config=self.client.api_config,
+                    tokenizer=self.client.tokenizer,
+                )
+                response_content = await response.text()
+
+                end_time = time.perf_counter()
+                error = None
+                if response.status != 200:
+                    error = ErrorResponseInfo(error_msg=response_content, error_type="Error response")
+
+                self.client.metrics_collector.record_metric(
+                    RequestLifecycleMetric(
+                        stage_id=stage_id,
+                        request_data=request_data,
+                        response_data=response_content,
+                        info=response_info,
+                        error=error,
+                        start_time=start,
+                        end_time=end_time,
+                        scheduled_time=scheduled_time,
+                    )
+                )
+        except Exception as e:
+            if isinstance(e, asyncio.exceptions.TimeoutError):
+                logger.error("request timed out:", exc_info=True)
+            else:
+                logger.error("error occured during request processing:", exc_info=True)
+
+            self.client.metrics_collector.record_metric(
+                RequestLifecycleMetric(
+                    stage_id=stage_id,
+                    request_data=request_data,
+                    response_data=response_content or "",
+                    info=response_info or InferenceInfo(),
+                    error=ErrorResponseInfo(
+                        error_msg=str(e),
+                        error_type=type(e).__name__,
+                    ),
+                    start_time=start,
+                    end_time=time.perf_counter(),
+                    scheduled_time=scheduled_time,
+                )
+            )
+
+    async def close(self) -> None:
+        await self.session.close()
diff --git a/inference_perf/loadgen/load_generator.py b/inference_perf/loadgen/load_generator.py
@@ -94,6 +94,7 @@ async def loop(self) -> None:
         event_loop = get_event_loop()
         item = None
         timeout = 0.5
+        session = self.client.new_session()
 
         while not self.stop_signal.is_set():
             # Check if max_concurrency has been updated and recreate semaphore if needed (concurrent load type)
@@ -154,7 +155,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 session.process_request(request_data, stage_id, request_time)
                     except CancelledError:
                         pass
                     finally:
@@ -188,6 +189,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:
