From d4820679f4012fb4a1883fec44916980d721207a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw?= Date: Wed, 25 Mar 2026 16:09:46 +0100 Subject: [PATCH 1/3] fix: respect max_parallel_requests in HTTP connection pool size MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pass a pre-configured HTTPTransport/AsyncHTTPTransport with the correct limits into RetryTransport instead of letting it create its own pool with httpx defaults (100 connections). Previously, the limits calculated from max_parallel_requests were passed to httpx.Client(limits=...) which silently ignores them when a custom transport is provided. Fixes #459 Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Przemysław --- .../clients/adapters/http_model_client.py | 12 ++++-- .../engine/models/clients/retry.py | 7 +++- .../clients/test_native_http_clients.py | 38 +++++++++++++++++++ .../tests/engine/models/clients/test_retry.py | 25 ++++++++++++ 4 files changed, 77 insertions(+), 5 deletions(-) diff --git a/packages/data-designer-engine/src/data_designer/engine/models/clients/adapters/http_model_client.py b/packages/data-designer-engine/src/data_designer/engine/models/clients/adapters/http_model_client.py index 8b030e0a7..082555855 100644 --- a/packages/data-designer-engine/src/data_designer/engine/models/clients/adapters/http_model_client.py +++ b/packages/data-designer-engine/src/data_designer/engine/models/clients/adapters/http_model_client.py @@ -97,10 +97,12 @@ def _get_sync_client(self) -> httpx.Client: raise RuntimeError("Model client is closed.") if self._client is None: if self._transport is None: - self._transport = create_retry_transport(self._retry_config, strip_rate_limit_codes=False) + inner = lazy.httpx.HTTPTransport(limits=self._limits) + self._transport = create_retry_transport( + self._retry_config, strip_rate_limit_codes=False, transport=inner + ) self._client = lazy.httpx.Client( transport=self._transport, - limits=self._limits, timeout=lazy.httpx.Timeout(self._timeout_s), ) return self._client @@ -113,10 +115,12 @@ def _get_async_client(self) -> httpx.AsyncClient: raise RuntimeError("Model client is closed.") if self._aclient is None: if self._transport is None: - self._transport = create_retry_transport(self._retry_config, strip_rate_limit_codes=True) + inner = lazy.httpx.AsyncHTTPTransport(limits=self._limits) + self._transport = create_retry_transport( + self._retry_config, strip_rate_limit_codes=True, transport=inner + ) self._aclient = lazy.httpx.AsyncClient( transport=self._transport, - limits=self._limits, timeout=lazy.httpx.Timeout(self._timeout_s), ) return self._aclient diff --git a/packages/data-designer-engine/src/data_designer/engine/models/clients/retry.py b/packages/data-designer-engine/src/data_designer/engine/models/clients/retry.py index c74fb5548..bef4f25c9 100644 --- a/packages/data-designer-engine/src/data_designer/engine/models/clients/retry.py +++ b/packages/data-designer-engine/src/data_designer/engine/models/clients/retry.py @@ -5,9 +5,13 @@ import logging from dataclasses import dataclass, field +from typing import TYPE_CHECKING from httpx_retries import Retry, RetryTransport +if TYPE_CHECKING: + import httpx + logger = logging.getLogger(__name__) # 429 must not be retried at the transport layer so that rate-limit signals @@ -37,6 +41,7 @@ def create_retry_transport( config: RetryConfig | None = None, *, strip_rate_limit_codes: bool = True, + transport: httpx.BaseTransport | httpx.AsyncBaseTransport | None = None, ) -> RetryTransport: """Build an httpx ``RetryTransport`` from a :class:`RetryConfig`. @@ -72,4 +77,4 @@ def create_retry_transport( respect_retry_after_header=True, allowed_methods=Retry.RETRYABLE_METHODS | frozenset(["POST"]), ) - return RetryTransport(retry=retry) + return RetryTransport(transport=transport, retry=retry) diff --git a/packages/data-designer-engine/tests/engine/models/clients/test_native_http_clients.py b/packages/data-designer-engine/tests/engine/models/clients/test_native_http_clients.py index b093fd7bf..14c2e3593 100644 --- a/packages/data-designer-engine/tests/engine/models/clients/test_native_http_clients.py +++ b/packages/data-designer-engine/tests/engine/models/clients/test_native_http_clients.py @@ -289,3 +289,41 @@ async def test_acompletion_lazy_initializes_async_client( mock_ctor.assert_called_once() assert result.message.content == "lazy result" + + +# --------------------------------------------------------------------------- +# Connection pool size regression tests (issue #459) +# --------------------------------------------------------------------------- + + +def test_sync_client_pool_size_respects_max_parallel_requests() -> None: + """Connection pool max_connections must be 2*max_parallel_requests, not the httpx default of 100.""" + client = OpenAICompatibleClient( + provider_name=_OPENAI_PROVIDER, + endpoint=_OPENAI_ENDPOINT, + api_key="sk-test", + max_parallel_requests=300, + concurrency_mode=ClientConcurrencyMode.SYNC, + ) + with patch(_SYNC_CLIENT_PATCH): + client._get_sync_client() + + # pool_max = max(32, 2 * 300) = 600 + assert client._transport._sync_transport._pool._max_connections == 600 + + +@pytest.mark.asyncio +async def test_async_client_pool_size_respects_max_parallel_requests() -> None: + """Async connection pool max_connections must be 2*max_parallel_requests, not the httpx default of 100.""" + client = OpenAICompatibleClient( + provider_name=_OPENAI_PROVIDER, + endpoint=_OPENAI_ENDPOINT, + api_key="sk-test", + max_parallel_requests=300, + concurrency_mode=ClientConcurrencyMode.ASYNC, + ) + with patch(_ASYNC_CLIENT_PATCH): + client._get_async_client() + + # pool_max = max(32, 2 * 300) = 600 + assert client._transport._async_transport._pool._max_connections == 600 diff --git a/packages/data-designer-engine/tests/engine/models/clients/test_retry.py b/packages/data-designer-engine/tests/engine/models/clients/test_retry.py index 9c86e5b7b..91a9d584a 100644 --- a/packages/data-designer-engine/tests/engine/models/clients/test_retry.py +++ b/packages/data-designer-engine/tests/engine/models/clients/test_retry.py @@ -82,3 +82,28 @@ def test_create_retry_transport_default_codes_exclude_429() -> None: """Default retryable_status_codes do not include 429 regardless of strip flag.""" transport = create_retry_transport(strip_rate_limit_codes=False) assert 429 not in transport.retry.status_forcelist + + +def test_create_retry_transport_forwards_sync_transport() -> None: + """Provided sync transport is used directly rather than replaced with a default pool.""" + import httpx + + inner = httpx.HTTPTransport(limits=httpx.Limits(max_connections=600)) + transport = create_retry_transport(transport=inner) + assert transport._sync_transport is inner + + +def test_create_retry_transport_forwards_async_transport() -> None: + """Provided async transport is used directly rather than replaced with a default pool.""" + import httpx + + inner = httpx.AsyncHTTPTransport(limits=httpx.Limits(max_connections=600)) + transport = create_retry_transport(transport=inner) + assert transport._async_transport is inner + + +def test_create_retry_transport_no_transport_creates_defaults() -> None: + """Without a transport argument both sync and async default pools are created.""" + transport = create_retry_transport() + assert transport._sync_transport is not None + assert transport._async_transport is not None From e558dd7b39ee93bc9ccd1c019c7d169b741176f0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw?= Date: Wed, 25 Mar 2026 16:23:10 +0100 Subject: [PATCH 2/3] fix: document transport param and annotate private attr chains in tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address Greptile P2 review comments on PR #460: - Add docstring entry for the new `transport` parameter in `create_retry_transport` explaining accepted types and None default - Add inline comments in pool-size regression tests explaining the private attribute chain (_sync/_async_transport → _pool → _max_connections) Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Przemysław --- .../src/data_designer/engine/models/clients/retry.py | 6 ++++++ .../engine/models/clients/test_native_http_clients.py | 8 ++++++++ 2 files changed, 14 insertions(+) diff --git a/packages/data-designer-engine/src/data_designer/engine/models/clients/retry.py b/packages/data-designer-engine/src/data_designer/engine/models/clients/retry.py index bef4f25c9..886ea6230 100644 --- a/packages/data-designer-engine/src/data_designer/engine/models/clients/retry.py +++ b/packages/data-designer-engine/src/data_designer/engine/models/clients/retry.py @@ -56,6 +56,12 @@ def create_retry_transport( AIMD feedback loop. When ``False`` (used by the sync engine, which has no salvage queue), 429 is kept in the retry list so the transport layer retries it transparently. + transport: Optional pre-configured inner transport to pass directly to + ``RetryTransport``. Pass ``httpx.HTTPTransport`` for sync clients or + ``httpx.AsyncHTTPTransport`` for async clients — typically with a custom + ``limits=`` — so that the connection pool is sized correctly. When + ``None`` (default), ``RetryTransport`` creates its own default pools for + both sync and async requests. """ cfg = config or RetryConfig() status_codes = cfg.retryable_status_codes diff --git a/packages/data-designer-engine/tests/engine/models/clients/test_native_http_clients.py b/packages/data-designer-engine/tests/engine/models/clients/test_native_http_clients.py index 14c2e3593..6962dfba4 100644 --- a/packages/data-designer-engine/tests/engine/models/clients/test_native_http_clients.py +++ b/packages/data-designer-engine/tests/engine/models/clients/test_native_http_clients.py @@ -308,6 +308,10 @@ def test_sync_client_pool_size_respects_max_parallel_requests() -> None: with patch(_SYNC_CLIENT_PATCH): client._get_sync_client() + # Attribute chain explained: + # RetryTransport._sync_transport → the httpx.HTTPTransport we injected + # HTTPTransport._pool → the underlying httpcore.ConnectionPool + # ConnectionPool._max_connections → the hard cap configured via Limits # pool_max = max(32, 2 * 300) = 600 assert client._transport._sync_transport._pool._max_connections == 600 @@ -325,5 +329,9 @@ async def test_async_client_pool_size_respects_max_parallel_requests() -> None: with patch(_ASYNC_CLIENT_PATCH): client._get_async_client() + # Attribute chain explained: + # RetryTransport._async_transport → the httpx.AsyncHTTPTransport we injected + # AsyncHTTPTransport._pool → the underlying httpcore.AsyncConnectionPool + # AsyncConnectionPool._max_connections → the hard cap configured via Limits # pool_max = max(32, 2 * 300) = 600 assert client._transport._async_transport._pool._max_connections == 600 From 088e0374462305a1ade976bb0592777ea160ac7b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw?= Date: Wed, 25 Mar 2026 21:14:22 +0100 Subject: [PATCH 3/3] =?UTF-8?q?refactor:=20address=20review=20feedback=20?= =?UTF-8?q?=E2=80=94=20public=20limits=20property,=20clean=20tests?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Drop three tests from test_retry.py that reached into private attributes of third-party RetryTransport (_sync_transport, _async_transport); the end-to-end contract is covered by the pool-size regression test - Expose a public `limits` property on HttpModelClient so tests and diagnostic code can assert the pool configuration without walking private attribute chains across three libraries - Replace two private-chain pool assertions with a single `client.limits.max_connections == 600` check against the new property - Trim "inner" from the transport docstring entry (nabinchha suggestion) Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Przemysław --- .../clients/adapters/http_model_client.py | 5 +++ .../engine/models/clients/retry.py | 2 +- .../clients/test_native_http_clients.py | 38 +++---------------- .../tests/engine/models/clients/test_retry.py | 25 ------------ 4 files changed, 12 insertions(+), 58 deletions(-) diff --git a/packages/data-designer-engine/src/data_designer/engine/models/clients/adapters/http_model_client.py b/packages/data-designer-engine/src/data_designer/engine/models/clients/adapters/http_model_client.py index 082555855..a07305f58 100644 --- a/packages/data-designer-engine/src/data_designer/engine/models/clients/adapters/http_model_client.py +++ b/packages/data-designer-engine/src/data_designer/engine/models/clients/adapters/http_model_client.py @@ -83,6 +83,11 @@ def __init__( def concurrency_mode(self) -> ClientConcurrencyMode: return self._mode + @property + def limits(self) -> httpx.Limits: + """Connection pool limits derived from ``max_parallel_requests`` at construction time.""" + return self._limits + @abstractmethod def _build_headers(self, extra_headers: dict[str, str]) -> dict[str, str]: """Build provider-specific request headers.""" diff --git a/packages/data-designer-engine/src/data_designer/engine/models/clients/retry.py b/packages/data-designer-engine/src/data_designer/engine/models/clients/retry.py index 886ea6230..56aa1eec4 100644 --- a/packages/data-designer-engine/src/data_designer/engine/models/clients/retry.py +++ b/packages/data-designer-engine/src/data_designer/engine/models/clients/retry.py @@ -56,7 +56,7 @@ def create_retry_transport( AIMD feedback loop. When ``False`` (used by the sync engine, which has no salvage queue), 429 is kept in the retry list so the transport layer retries it transparently. - transport: Optional pre-configured inner transport to pass directly to + transport: Optional pre-configured transport to pass directly to ``RetryTransport``. Pass ``httpx.HTTPTransport`` for sync clients or ``httpx.AsyncHTTPTransport`` for async clients — typically with a custom ``limits=`` — so that the connection pool is sized correctly. When diff --git a/packages/data-designer-engine/tests/engine/models/clients/test_native_http_clients.py b/packages/data-designer-engine/tests/engine/models/clients/test_native_http_clients.py index 6962dfba4..1e58f0c12 100644 --- a/packages/data-designer-engine/tests/engine/models/clients/test_native_http_clients.py +++ b/packages/data-designer-engine/tests/engine/models/clients/test_native_http_clients.py @@ -296,42 +296,16 @@ async def test_acompletion_lazy_initializes_async_client( # --------------------------------------------------------------------------- -def test_sync_client_pool_size_respects_max_parallel_requests() -> None: - """Connection pool max_connections must be 2*max_parallel_requests, not the httpx default of 100.""" - client = OpenAICompatibleClient( - provider_name=_OPENAI_PROVIDER, - endpoint=_OPENAI_ENDPOINT, - api_key="sk-test", - max_parallel_requests=300, - concurrency_mode=ClientConcurrencyMode.SYNC, - ) - with patch(_SYNC_CLIENT_PATCH): - client._get_sync_client() - - # Attribute chain explained: - # RetryTransport._sync_transport → the httpx.HTTPTransport we injected - # HTTPTransport._pool → the underlying httpcore.ConnectionPool - # ConnectionPool._max_connections → the hard cap configured via Limits - # pool_max = max(32, 2 * 300) = 600 - assert client._transport._sync_transport._pool._max_connections == 600 +def test_client_limits_respect_max_parallel_requests() -> None: + """Connection pool limits must reflect max_parallel_requests (regression for issue #459). - -@pytest.mark.asyncio -async def test_async_client_pool_size_respects_max_parallel_requests() -> None: - """Async connection pool max_connections must be 2*max_parallel_requests, not the httpx default of 100.""" + pool_max = max(32, 2 * max_parallel_requests) = max(32, 600) = 600 + """ client = OpenAICompatibleClient( provider_name=_OPENAI_PROVIDER, endpoint=_OPENAI_ENDPOINT, api_key="sk-test", max_parallel_requests=300, - concurrency_mode=ClientConcurrencyMode.ASYNC, + concurrency_mode=ClientConcurrencyMode.SYNC, ) - with patch(_ASYNC_CLIENT_PATCH): - client._get_async_client() - - # Attribute chain explained: - # RetryTransport._async_transport → the httpx.AsyncHTTPTransport we injected - # AsyncHTTPTransport._pool → the underlying httpcore.AsyncConnectionPool - # AsyncConnectionPool._max_connections → the hard cap configured via Limits - # pool_max = max(32, 2 * 300) = 600 - assert client._transport._async_transport._pool._max_connections == 600 + assert client.limits.max_connections == 600 diff --git a/packages/data-designer-engine/tests/engine/models/clients/test_retry.py b/packages/data-designer-engine/tests/engine/models/clients/test_retry.py index 91a9d584a..9c86e5b7b 100644 --- a/packages/data-designer-engine/tests/engine/models/clients/test_retry.py +++ b/packages/data-designer-engine/tests/engine/models/clients/test_retry.py @@ -82,28 +82,3 @@ def test_create_retry_transport_default_codes_exclude_429() -> None: """Default retryable_status_codes do not include 429 regardless of strip flag.""" transport = create_retry_transport(strip_rate_limit_codes=False) assert 429 not in transport.retry.status_forcelist - - -def test_create_retry_transport_forwards_sync_transport() -> None: - """Provided sync transport is used directly rather than replaced with a default pool.""" - import httpx - - inner = httpx.HTTPTransport(limits=httpx.Limits(max_connections=600)) - transport = create_retry_transport(transport=inner) - assert transport._sync_transport is inner - - -def test_create_retry_transport_forwards_async_transport() -> None: - """Provided async transport is used directly rather than replaced with a default pool.""" - import httpx - - inner = httpx.AsyncHTTPTransport(limits=httpx.Limits(max_connections=600)) - transport = create_retry_transport(transport=inner) - assert transport._async_transport is inner - - -def test_create_retry_transport_no_transport_creates_defaults() -> None: - """Without a transport argument both sync and async default pools are created.""" - transport = create_retry_transport() - assert transport._sync_transport is not None - assert transport._async_transport is not None