Skip to content

Commit 7bb29b9

Browse files
fxdgearclaude
andcommitted
feat: env-var configurable httpx pool config in split_pdf_hook
Make httpx.AsyncClient pool config in run_tasks configurable via env vars: - UNSTRUCTURED_CLIENT_MAX_CONNECTIONS (default 100) - UNSTRUCTURED_CLIENT_MAX_KEEPALIVE_CONNECTIONS (default 20) - UNSTRUCTURED_CLIENT_KEEPALIVE_EXPIRY (default 5.0 seconds) Defaults match httpx's built-in defaults, so this is backward compatible. Also extends the split_pdf event=plan_created INFO log to include the resolved pool values so operators can verify the active config from production logs. Motivation: when the SDK runs inside a Kubernetes Deployment that load-balances only at TCP-connect time (kube-proxy ClusterIP, meshless), httpx's default keepalive pooling locks onto a small subset of backend pods. Allowing operators to lower MAX_KEEPALIVE_CONNECTIONS and KEEPALIVE_EXPIRY forces frequent reconnects and redistributes traffic. Encountered on INC-11422. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
1 parent 5e6f238 commit 7bb29b9

2 files changed

Lines changed: 72 additions & 3 deletions

File tree

‎_test_unstructured_client/unit/test_split_pdf_hook.py‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -482,6 +482,44 @@ async def test_remaining_tasks_cancelled_when_fails_disallowed():
482482
assert len(tasks) > cancelled_counter["cancelled"] > 0
483483

484484

485+
@pytest.mark.asyncio
486+
async def test_unit_run_tasks_pool_limits_configurable_via_env(
487+
monkeypatch: pytest.MonkeyPatch,
488+
):
489+
"""Env vars override the httpx.AsyncClient connection-pool limits.
490+
491+
Operators running the SDK in a Kubernetes Deployment with
492+
connect-time-only load balancing (kube-proxy ClusterIP, meshless)
493+
need to be able to shrink the keepalive pool so connections recycle
494+
frequently and redistribute across backend pods.
495+
"""
496+
monkeypatch.setenv("UNSTRUCTURED_CLIENT_MAX_CONNECTIONS", "7")
497+
monkeypatch.setenv("UNSTRUCTURED_CLIENT_MAX_KEEPALIVE_CONNECTIONS", "1")
498+
monkeypatch.setenv("UNSTRUCTURED_CLIENT_KEEPALIVE_EXPIRY", "30.0")
499+
500+
captured: dict[str, httpx.Limits] = {}
501+
real_async_client = httpx.AsyncClient
502+
503+
def _capturing_async_client(*args, **kwargs):
504+
captured["limits"] = kwargs.get("limits")
505+
return real_async_client(*args, **kwargs)
506+
507+
with patch(
508+
"unstructured_client._hooks.custom.split_pdf_hook.httpx.AsyncClient",
509+
side_effect=_capturing_async_client,
510+
):
511+
await run_tasks(
512+
[partial(_request_mock, fails=False, content="ok")],
513+
allow_failed=True,
514+
)
515+
516+
limits = captured["limits"]
517+
assert isinstance(limits, httpx.Limits)
518+
assert limits.max_connections == 7
519+
assert limits.max_keepalive_connections == 1
520+
assert limits.keepalive_expiry == 30.0
521+
522+
485523
@patch("unstructured_client._hooks.custom.form_utils.Path")
486524
def test_unit_get_split_pdf_cache_tmp_data_dir_uses_dir_from_form_data(mock_path: MagicMock):
487525
"""Test get_split_pdf_cache_tmp_data_dir uses the directory from the form data."""

‎src/unstructured_client/_hooks/custom/split_pdf_hook.py‎

Lines changed: 34 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -177,6 +177,28 @@ async def _order_keeper(index: int, coro: Awaitable) -> Tuple[int, httpx.Respons
177177
return index, response
178178

179179

180+
def _resolve_pool_limits() -> httpx.Limits:
181+
"""Resolve httpx connection-pool limits from environment variables.
182+
183+
Defaults match httpx's built-in defaults (max_connections=100,
184+
max_keepalive_connections=20, keepalive_expiry=5.0) so behavior is
185+
unchanged for callers that do not set the env vars. Operators running
186+
the SDK inside a Kubernetes Deployment that load-balances only at
187+
TCP-connect time (e.g. kube-proxy ClusterIP, no service mesh) can
188+
lower these values to force frequent reconnects and redistribute
189+
traffic across backend pods.
190+
"""
191+
return httpx.Limits(
192+
max_connections=int(os.getenv("UNSTRUCTURED_CLIENT_MAX_CONNECTIONS", "100")),
193+
max_keepalive_connections=int(
194+
os.getenv("UNSTRUCTURED_CLIENT_MAX_KEEPALIVE_CONNECTIONS", "20")
195+
),
196+
keepalive_expiry=float(
197+
os.getenv("UNSTRUCTURED_CLIENT_KEEPALIVE_EXPIRY", "5.0")
198+
),
199+
)
200+
201+
180202
async def run_tasks(
181203
coroutines: list[partial[Coroutine[Any, Any, httpx.Response]]],
182204
allow_failed: bool = False,
@@ -205,16 +227,21 @@ async def run_tasks(
205227
client_timeout_minutes = int(timeout_var)
206228
client_timeout = httpx.Timeout(60 * client_timeout_minutes)
207229

230+
limits = _resolve_pool_limits()
231+
208232
logger.debug(
209-
"split_pdf event=batch_async_start operation_id=%s chunk_count=%d concurrency=%d client_timeout=%s allow_failed=%s",
233+
"split_pdf event=batch_async_start operation_id=%s chunk_count=%d concurrency=%d client_timeout=%s allow_failed=%s pool_max_connections=%s pool_max_keepalive=%s pool_keepalive_expiry=%s",
210234
operation_id,
211235
len(coroutines),
212236
concurrency_level,
213237
client_timeout,
214238
allow_failed,
239+
limits.max_connections,
240+
limits.max_keepalive_connections,
241+
limits.keepalive_expiry,
215242
)
216243

217-
async with httpx.AsyncClient(timeout=client_timeout) as client:
244+
async with httpx.AsyncClient(timeout=client_timeout, limits=limits) as client:
218245
armed_coroutines = [coro(async_client=client, limiter=limiter) for coro in coroutines] # type: ignore
219246
tasks = [
220247
asyncio.create_task(_order_keeper(index, coro))
@@ -770,8 +797,9 @@ def _before_request_unlocked(
770797
)
771798
self.coroutines_to_execute[operation_id].append(coroutine)
772799

800+
plan_limits = _resolve_pool_limits()
773801
logger.info(
774-
"split_pdf event=plan_created operation_id=%s filename=%s strategy=%s page_range=%s-%s page_count=%d split_size=%d chunk_count=%d concurrency=%d allow_failed=%s cache_mode=%s timeout_seconds=%s retry_config_mode=%s",
802+
"split_pdf event=plan_created operation_id=%s filename=%s strategy=%s page_range=%s-%s page_count=%d split_size=%d chunk_count=%d concurrency=%d allow_failed=%s cache_mode=%s timeout_seconds=%s retry_config_mode=%s pool_max_connections=%d pool_max_keepalive=%d pool_keepalive_expiry=%.1fs",
775803
operation_id,
776804
Path(pdf_file_meta["filename"]).name,
777805
form_data.get("strategy"),
@@ -790,6 +818,9 @@ def _before_request_unlocked(
790818
self._retry_config_observability_mode(
791819
self.operation_retry_configs.get(operation_id),
792820
),
821+
plan_limits.max_connections,
822+
plan_limits.max_keepalive_connections,
823+
plan_limits.keepalive_expiry,
793824
)
794825

795826
self.pending_operation_ids[operation_id] = operation_id

0 commit comments

Comments
 (0)