From 7fc11ba538258a1f6f39f95ccd6ae61d2c0a835f Mon Sep 17 00:00:00 2001 From: Reflex Date: Mon, 27 Jul 2026 21:25:59 +0000 Subject: [PATCH 01/11] fix: open extra H2 conns under stream pressure; isolate file transfers on HTTP/1.1 httpcore blocks on a per-connection stream semaphore once slots fill instead of opening another connection. Patch is_available + non-blocking acquire so the pool spreads load. Route upload_file/download_file through a dedicated HTTP/1.1 pool so bulk bodies cannot starve latency-sensitive RPCs on the shared H2 session. Co-authored-by: Cursor --- src/runloop_api_client/__init__.py | 5 + src/runloop_api_client/_base_client.py | 110 ++++++++++++- src/runloop_api_client/_http2_pool_fix.py | 185 ++++++++++++++++++++++ tests/test_http2_pool_fix.py | 72 +++++++++ tests/test_shared_pool.py | 14 +- tests/test_transfer_client.py | 133 ++++++++++++++++ 6 files changed, 512 insertions(+), 7 deletions(-) create mode 100644 src/runloop_api_client/_http2_pool_fix.py create mode 100644 tests/test_http2_pool_fix.py create mode 100644 tests/test_transfer_client.py diff --git a/src/runloop_api_client/__init__.py b/src/runloop_api_client/__init__.py index c74afc2c1..e0baca008 100644 --- a/src/runloop_api_client/__init__.py +++ b/src/runloop_api_client/__init__.py @@ -28,8 +28,13 @@ APIResponseValidationError, ) from ._base_client import DefaultHttpxClient, DefaultAioHttpClient, DefaultAsyncHttpxClient +from ._http2_pool_fix import install as _install_http2_pool_fix from ._utils._logs import setup_logging as _setup_logging +# Open additional HTTP/2 connections when a connection's stream slots are full +# instead of blocking on httpcore's per-connection stream semaphore. +_install_http2_pool_fix() + __all__ = [ "types", "__version__", diff --git a/src/runloop_api_client/_base_client.py b/src/runloop_api_client/_base_client.py index 88e0bbb3b..bc636112d 100644 --- a/src/runloop_api_client/_base_client.py +++ b/src/runloop_api_client/_base_client.py @@ -174,6 +174,21 @@ async def aclose(self) -> None: weakref.WeakKeyDictionary() ) +# Separate HTTP/1.1 pool for bulk file upload/download. Keeping these off the +# main HTTP/2 connection avoids H2 session-window / stream-slot contention with +# latency-sensitive RPCs (create, wait_for_status, execute, …). +_shared_sync_transfer_transport: _SharedTransport | None = None +_shared_async_transfer_transports: weakref.WeakKeyDictionary[ + asyncio.AbstractEventLoop, _SharedAsyncTransport +] = weakref.WeakKeyDictionary() + +# Paths that carry large request/response bodies and should use the transfer pool. +_FILE_TRANSFER_PATH_SUFFIXES = ("/upload_file", "/download_file") + + +def _is_file_transfer_path(path: str) -> bool: + return path.endswith(_FILE_TRANSFER_PATH_SUFFIXES) + # TODO: make base page type vars covariant SyncPageT = TypeVar("SyncPageT", bound="BaseSyncPage[Any]") AsyncPageT = TypeVar("AsyncPageT", bound="BaseAsyncPage[Any]") @@ -929,8 +944,10 @@ def __del__(self) -> None: class SyncAPIClient(BaseClient[httpx.Client, Stream[Any]]): _client: httpx.Client + _transfer_client: httpx.Client | None _default_stream_cls: type[Stream[Any]] | None = None _uses_shared_pool: bool + _isolate_file_transfers: bool _closed: bool def __init__( @@ -976,6 +993,9 @@ def __init__( ) self._closed = False + self._transfer_client = None + # Custom http_client owns the full transport stack; don't invent a sibling pool. + self._isolate_file_transfers = http_client is None if http_client is not None: self._client = http_client @@ -1000,6 +1020,38 @@ def __init__( ) self._uses_shared_pool = False + def _ensure_transfer_client(self) -> httpx.Client: + """Lazy HTTP/1.1 client for upload_file / download_file.""" + if self._transfer_client is not None: + return self._transfer_client + + timeout = cast(Timeout, self.timeout) + if self._uses_shared_pool: + global _shared_sync_transfer_transport + with _pool_lock: + if _shared_sync_transfer_transport is None or not _shared_sync_transfer_transport.acquire(): + _shared_sync_transfer_transport = _SharedTransport( + httpx.HTTPTransport(limits=DEFAULT_CONNECTION_LIMITS, http2=False), + ) + self._transfer_client = SyncHttpxClientWrapper( + base_url=self._base_url, + timeout=timeout, + transport=_shared_sync_transfer_transport, + http2=False, + ) + else: + self._transfer_client = SyncHttpxClientWrapper( + base_url=self._base_url, + timeout=timeout, + http2=False, + ) + return self._transfer_client + + def _send_client_for_request(self, request: httpx.Request) -> httpx.Client: + if self._isolate_file_transfers and _is_file_transfer_path(request.url.path): + return self._ensure_transfer_client() + return self._client + def is_closed(self) -> bool: return self._closed or self._client.is_closed @@ -1014,6 +1066,10 @@ def close(self) -> None: return self._closed = True self._client.close() + transfer = self._transfer_client + self._transfer_client = None + if transfer is not None: + transfer.close() def __enter__(self: _T) -> _T: return self @@ -1114,7 +1170,7 @@ def request( response = None try: - response = self._client.send( + response = self._send_client_for_request(request).send( request, stream=stream or self._should_stream_response_body(request=request), **kwargs, @@ -1561,8 +1617,10 @@ def __del__(self) -> None: class AsyncAPIClient(BaseClient[httpx.AsyncClient, AsyncStream[Any]]): _client: httpx.AsyncClient + _transfer_client: httpx.AsyncClient | None _default_stream_cls: type[AsyncStream[Any]] | None = None _uses_shared_pool: bool + _isolate_file_transfers: bool _closed: bool def __init__( @@ -1608,6 +1666,9 @@ def __init__( ) self._closed = False + self._transfer_client = None + # Custom http_client owns the full transport stack; don't invent a sibling pool. + self._isolate_file_transfers = http_client is None if http_client is not None: self._client = http_client @@ -1646,6 +1707,47 @@ def __init__( ) self._uses_shared_pool = False + def _ensure_transfer_client(self) -> httpx.AsyncClient: + """Lazy HTTP/1.1 client for upload_file / download_file.""" + if self._transfer_client is not None: + return self._transfer_client + + timeout = cast(Timeout, self.timeout) + if self._uses_shared_pool: + try: + loop: asyncio.AbstractEventLoop | None = asyncio.get_running_loop() + except RuntimeError: + loop = None + if loop is not None: + with _pool_lock: + existing = _shared_async_transfer_transports.get(loop) + if existing is not None and existing.acquire(): + transport: _SharedAsyncTransport = existing + else: + transport = _SharedAsyncTransport( + httpx.AsyncHTTPTransport(limits=DEFAULT_CONNECTION_LIMITS, http2=False), + ) + _shared_async_transfer_transports[loop] = transport + self._transfer_client = AsyncHttpxClientWrapper( + base_url=self._base_url, + timeout=timeout, + transport=transport, + http2=False, + ) + return self._transfer_client + + self._transfer_client = AsyncHttpxClientWrapper( + base_url=self._base_url, + timeout=timeout, + http2=False, + ) + return self._transfer_client + + def _send_client_for_request(self, request: httpx.Request) -> httpx.AsyncClient: + if self._isolate_file_transfers and _is_file_transfer_path(request.url.path): + return self._ensure_transfer_client() + return self._client + def is_closed(self) -> bool: return self._closed or self._client.is_closed @@ -1660,6 +1762,10 @@ async def close(self) -> None: return self._closed = True await self._client.aclose() + transfer = self._transfer_client + self._transfer_client = None + if transfer is not None: + await transfer.aclose() async def __aenter__(self: _T) -> _T: return self @@ -1765,7 +1871,7 @@ async def request( response = None try: - response = await self._client.send( + response = await self._send_client_for_request(request).send( request, stream=stream or self._should_stream_response_body(request=request), **kwargs, diff --git a/src/runloop_api_client/_http2_pool_fix.py b/src/runloop_api_client/_http2_pool_fix.py new file mode 100644 index 000000000..4c8c24493 --- /dev/null +++ b/src/runloop_api_client/_http2_pool_fix.py @@ -0,0 +1,185 @@ +"""Make httpcore open additional HTTP/2 connections when stream slots run low. + +httpcore multiplexes HTTP/2 on a single connection and gates new streams with a +semaphore (default MAX_CONCURRENT_STREAMS=100). Its connection pool's +``is_available()`` ignores that limit, so once stream slots are exhausted, +further requests *block* on the semaphore instead of the pool opening another +connection — even when ``max_connections`` still has headroom. + +Additionally, the pool may assign a *burst* of requests to a connection that +only has a few free slots left (``is_available()`` is checked once per request +but slots are not reserved). Winners then share an overloaded connection while +the rest raise ``ConnectionNotAvailable`` and retry. + +This module patches sync + async HTTP/2 connections so that: + +1. ``is_available()`` is False when free stream slots are below a headroom + threshold, so the pool prefers opening/reusing another connection. +2. Stream-slot acquire is non-blocking; if a race still over-assigns, we raise + ``ConnectionNotAvailable`` and the pool retries on another connection. + +Idempotent: safe to call ``install()`` more than once. +""" + +from __future__ import annotations + +import logging +from typing import Any, Callable + +logger = logging.getLogger("runloop_api_client._http2_pool_fix") + +# If a connection has fewer free stream slots than this, treat it as unavailable +# so the pool opens another connection instead of stampeding the remainder. +# Must be >1: the pool assigns many queued requests to one "available" conn +# without reserving slots, so small remaining capacity still over-assigns. +_MIN_FREE_STREAM_SLOTS = 16 + +_installed = False + + +def _free_stream_slots(connection: Any) -> int | None: + """Return free H2 stream slots, or None if unknown.""" + max_streams = getattr(connection, "_max_streams", None) + events = getattr(connection, "_events", None) + if not isinstance(max_streams, int) or max_streams <= 0 or events is None: + return None + return max_streams - len(events) + + +def _stream_slots_saturated(connection: Any) -> bool: + """True when the connection should not accept more streams.""" + free = _free_stream_slots(connection) + if free is None: + return False + return free < _MIN_FREE_STREAM_SLOTS + + +def _patch_is_available(cls: type) -> None: + if getattr(cls.is_available, "_runloop_stream_overflow_patched", False): + return + + original: Callable[[Any], bool] = cls.is_available + + def is_available(self: Any) -> bool: + if not original(self): + return False + return not _stream_slots_saturated(self) + + is_available._runloop_stream_overflow_patched = True # type: ignore[attr-defined] + cls.is_available = is_available # type: ignore[method-assign] + + +def _make_async_nonblocking_acquire(sem: Any, connection_not_available: type) -> Callable[[], Any]: + original_acquire = sem.acquire + + async def acquire() -> None: + if not getattr(sem, "_backend", ""): + sem.setup() + + if sem._backend == "asyncio": + inner = sem._anyio_semaphore + try: + inner.acquire_nowait() + return + except Exception: + raise connection_not_available( + "HTTP/2 connection has no free stream slots" + ) from None + + if sem._backend == "trio": + inner = sem._trio_semaphore + acquire_nowait = getattr(inner, "acquire_nowait", None) + if acquire_nowait is not None: + try: + acquire_nowait() + return + except Exception: + raise connection_not_available( + "HTTP/2 connection has no free stream slots" + ) from None + + await original_acquire() + + return acquire + + +def _make_sync_nonblocking_acquire(sem: Any, connection_not_available: type) -> Callable[[], None]: + inner = getattr(sem, "_semaphore", sem) + + def acquire() -> None: + ok = inner.acquire(False) + if not ok: + raise connection_not_available("HTTP/2 connection has no free stream slots") + + return acquire + + +def _patch_async_connection(cls: type) -> None: + from httpcore import ConnectionNotAvailable + + if getattr(cls.handle_async_request, "_runloop_stream_overflow_patched", False): + return + + original = cls.handle_async_request + + async def handle_async_request(self: Any, request: Any) -> Any: + sem = getattr(self, "_max_streams_semaphore", None) + if sem is None: + return await original(self, request) + + real_acquire = sem.acquire + sem.acquire = _make_async_nonblocking_acquire(sem, ConnectionNotAvailable) + try: + return await original(self, request) + finally: + sem.acquire = real_acquire + + handle_async_request._runloop_stream_overflow_patched = True # type: ignore[attr-defined] + cls.handle_async_request = handle_async_request # type: ignore[method-assign] + + +def _patch_sync_connection(cls: type) -> None: + from httpcore import ConnectionNotAvailable + + if getattr(cls.handle_request, "_runloop_stream_overflow_patched", False): + return + + original = cls.handle_request + + def handle_request(self: Any, request: Any) -> Any: + sem = getattr(self, "_max_streams_semaphore", None) + if sem is None: + return original(self, request) + + real_acquire = sem.acquire + sem.acquire = _make_sync_nonblocking_acquire(sem, ConnectionNotAvailable) + try: + return original(self, request) + finally: + sem.acquire = real_acquire + + handle_request._runloop_stream_overflow_patched = True # type: ignore[attr-defined] + cls.handle_request = handle_request # type: ignore[method-assign] + + +def install() -> bool: + """Patch httpcore HTTP/2 connections. Returns True if a patch was applied.""" + global _installed + if _installed: + return False + + try: + from httpcore._async.http2 import AsyncHTTP2Connection + from httpcore._sync.http2 import HTTP2Connection + except ImportError as exc: # pragma: no cover + logger.warning("http2 stream-overflow patch skipped: %s", exc) + return False + + _patch_is_available(AsyncHTTP2Connection) + _patch_is_available(HTTP2Connection) + _patch_async_connection(AsyncHTTP2Connection) + _patch_sync_connection(HTTP2Connection) + + _installed = True + logger.debug("installed httpcore HTTP/2 stream-overflow connection patch") + return True diff --git a/tests/test_http2_pool_fix.py b/tests/test_http2_pool_fix.py new file mode 100644 index 000000000..aa5135ab9 --- /dev/null +++ b/tests/test_http2_pool_fix.py @@ -0,0 +1,72 @@ +"""Tests for HTTP/2 stream-overflow → new connection behavior.""" + +from __future__ import annotations + +from typing import Any +from unittest.mock import MagicMock + +import pytest + +from runloop_api_client._http2_pool_fix import ( + _stream_slots_saturated, + install, +) + + +def test_install_is_idempotent() -> None: + # Package import already installs the patch. + assert install() is False + + +def test_stream_slots_saturated_helper() -> None: + conn: Any = MagicMock() + conn._max_streams = 100 + conn._events = {i: [] for i in range(1, 90)} # 89 in flight → 11 free < 16 + assert _stream_slots_saturated(conn) is True + + conn._events = {i: [] for i in range(1, 80)} # 79 in flight → 21 free + assert _stream_slots_saturated(conn) is False + + conn._max_streams = 0 + assert _stream_slots_saturated(conn) is False + + +def _bare_async_h2_connection() -> Any: + from httpcore._async.http2 import AsyncHTTP2Connection, HTTPConnectionState + + conn = object.__new__(AsyncHTTP2Connection) + conn._state = HTTPConnectionState.ACTIVE + conn._connection_error = False + conn._used_all_stream_ids = False + conn._h2_state = MagicMock() + # Anything other than ConnectionState.CLOSED. + conn._h2_state.state_machine.state = object() + return conn + + +def test_is_available_false_when_streams_nearly_full() -> None: + conn = _bare_async_h2_connection() + conn._max_streams = 100 + conn._events = {i: [] for i in range(1, 95)} # 6 free + assert conn.is_available() is False + + +def test_is_available_true_when_stream_slots_remain() -> None: + conn = _bare_async_h2_connection() + conn._max_streams = 100 + conn._events = {i: [] for i in range(1, 50)} # 51 free + assert conn.is_available() is True + + +def test_sync_is_available_false_when_streams_nearly_full() -> None: + from httpcore._sync.http2 import HTTP2Connection, HTTPConnectionState + + conn = object.__new__(HTTP2Connection) + conn._state = HTTPConnectionState.ACTIVE + conn._connection_error = False + conn._used_all_stream_ids = False + conn._h2_state = MagicMock() + conn._h2_state.state_machine.state = object() + conn._max_streams = 100 + conn._events = {i: [] for i in range(1, 95)} + assert conn.is_available() is False diff --git a/tests/test_shared_pool.py b/tests/test_shared_pool.py index 4220f8ba9..982350f07 100644 --- a/tests/test_shared_pool.py +++ b/tests/test_shared_pool.py @@ -30,13 +30,17 @@ def _reset_shared_pool() -> Iterator[None]: # pyright: ignore[reportUnusedFunct def _clear_pool_state() -> None: with _base_mod._pool_lock: old_sync = _base_mod._shared_sync_transport + old_sync_transfer = _base_mod._shared_sync_transfer_transport _base_mod._shared_sync_transport = None + _base_mod._shared_sync_transfer_transport = None _base_mod._shared_async_transports.clear() - if old_sync is not None: - try: - old_sync._transport.close() - except Exception: - pass + _base_mod._shared_async_transfer_transports.clear() + for transport in (old_sync, old_sync_transfer): + if transport is not None: + try: + transport._transport.close() + except Exception: + pass def _make_client(**kwargs: Any) -> Runloop: diff --git a/tests/test_transfer_client.py b/tests/test_transfer_client.py new file mode 100644 index 000000000..2ca84ec6e --- /dev/null +++ b/tests/test_transfer_client.py @@ -0,0 +1,133 @@ +"""Tests for dedicated HTTP/1.1 transfer pool used by upload/download.""" + +from __future__ import annotations + +import os +from typing import Any, Iterator + +import httpx +import pytest + +import runloop_api_client._base_client as _base_mod +from runloop_api_client import Runloop, AsyncRunloop +from runloop_api_client._base_client import _is_file_transfer_path + +base_url = os.environ.get("TEST_API_BASE_URL", "http://127.0.0.1:4010") +bearer_token = "My Bearer Token" + + +@pytest.fixture(autouse=True) +def _reset_shared_pool() -> Iterator[None]: # pyright: ignore[reportUnusedFunction] + _clear_pool_state() + yield + _clear_pool_state() + + +def _clear_pool_state() -> None: + with _base_mod._pool_lock: + old_sync = _base_mod._shared_sync_transport + old_sync_transfer = _base_mod._shared_sync_transfer_transport + _base_mod._shared_sync_transport = None + _base_mod._shared_sync_transfer_transport = None + _base_mod._shared_async_transports.clear() + _base_mod._shared_async_transfer_transports.clear() + for transport in (old_sync, old_sync_transfer): + if transport is not None: + try: + transport._transport.close() + except Exception: + pass + + +def _make_client(**kwargs: Any) -> Runloop: + kwargs.setdefault("base_url", base_url) + kwargs.setdefault("bearer_token", bearer_token) + return Runloop(**kwargs) + + +def test_is_file_transfer_path() -> None: + assert _is_file_transfer_path("/v1/devboxes/dbx_1/upload_file") is True + assert _is_file_transfer_path("/v1/devboxes/dbx_1/download_file") is True + assert _is_file_transfer_path("/v1/devboxes") is False + assert _is_file_transfer_path("/v1/objects/obj_1/download") is False + assert _is_file_transfer_path("/v1/scenarios/runs/run_1/download_logs") is False + + +def test_transfer_client_is_http1_and_separate_from_api_pool() -> None: + client = _make_client(shared_http_pool=True) + try: + assert client._isolate_file_transfers is True + assert client._transfer_client is None + + api_request = httpx.Request("POST", f"{base_url}/v1/devboxes") + transfer_request = httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_1/upload_file") + + assert client._send_client_for_request(api_request) is client._client + + transfer = client._send_client_for_request(transfer_request) + assert transfer is not client._client + assert transfer is client._transfer_client + + api_transport = client._client._transport # type: ignore[attr-defined] + transfer_transport = transfer._transport # type: ignore[attr-defined] + assert api_transport is not transfer_transport + assert _base_mod._shared_sync_transfer_transport is transfer_transport + # Real httpx transport under the shared wrapper is HTTP/1.1-only. + inner = transfer_transport._transport + assert getattr(inner, "_http2", False) is False + finally: + client.close() + + +def test_transfer_pool_is_shared_across_sdk_clients() -> None: + c1 = _make_client(shared_http_pool=True) + c2 = _make_client(shared_http_pool=True) + try: + req = httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_1/download_file") + t1 = c1._send_client_for_request(req) + t2 = c2._send_client_for_request(req) + assert t1 is not t2 + assert t1._transport is t2._transport # type: ignore[attr-defined] + assert _base_mod._shared_sync_transfer_transport is not None + assert _base_mod._shared_sync_transfer_transport.refcount == 2 + finally: + c1.close() + c2.close() + assert _base_mod._shared_sync_transfer_transport is not None + assert _base_mod._shared_sync_transfer_transport.refcount == 0 + + +def test_custom_http_client_skips_transfer_isolation() -> None: + custom = httpx.Client() + client = _make_client(http_client=custom) + try: + assert client._isolate_file_transfers is False + req = httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_1/upload_file") + assert client._send_client_for_request(req) is custom + assert client._transfer_client is None + finally: + client.close() + custom.close() + + +def test_private_pool_still_isolates_transfers() -> None: + client = _make_client(shared_http_pool=False) + try: + req = httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_1/upload_file") + transfer = client._send_client_for_request(req) + assert transfer is not client._client + assert client._uses_shared_pool is False + finally: + client.close() + + +@pytest.mark.asyncio +async def test_async_transfer_client_is_separate() -> None: + client = AsyncRunloop(base_url=base_url, bearer_token=bearer_token, shared_http_pool=True) + try: + req = httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_1/upload_file") + transfer = client._send_client_for_request(req) + assert transfer is not client._client + assert client._transfer_client is transfer + finally: + await client.close() From bfc71dfc4e2da6606aa9192645fa646bed17d4f4 Mon Sep 17 00:00:00 2001 From: Reflex Date: Mon, 27 Jul 2026 23:07:16 +0000 Subject: [PATCH 02/11] fix: isolate waits and file transfers on sharded H2 pools Drop the httpcore monkeypatch. Route long-polls (/wait_for_status) and upload/download onto separate shared H2 bulkhead pools with configurable shards (default 2), hashed by resource id, so control-plane RPCs keep their own connection without HTTP/1.1 explosion. Co-authored-by: Cursor --- src/runloop_api_client/__init__.py | 5 - src/runloop_api_client/_base_client.py | 293 ++++++++++++++++------ src/runloop_api_client/_client.py | 33 +++ src/runloop_api_client/_constants.py | 6 + src/runloop_api_client/_http2_pool_fix.py | 185 -------------- src/runloop_api_client/sdk/async_.py | 9 + src/runloop_api_client/sdk/sync.py | 9 + tests/test_http2_pool_fix.py | 72 ------ tests/test_shared_pool.py | 9 +- tests/test_transfer_client.py | 158 +++++++----- 10 files changed, 376 insertions(+), 403 deletions(-) delete mode 100644 src/runloop_api_client/_http2_pool_fix.py delete mode 100644 tests/test_http2_pool_fix.py diff --git a/src/runloop_api_client/__init__.py b/src/runloop_api_client/__init__.py index e0baca008..c74afc2c1 100644 --- a/src/runloop_api_client/__init__.py +++ b/src/runloop_api_client/__init__.py @@ -28,13 +28,8 @@ APIResponseValidationError, ) from ._base_client import DefaultHttpxClient, DefaultAioHttpClient, DefaultAsyncHttpxClient -from ._http2_pool_fix import install as _install_http2_pool_fix from ._utils._logs import setup_logging as _setup_logging -# Open additional HTTP/2 connections when a connection's stream slots are full -# instead of blocking on httpcore's per-connection stream semaphore. -_install_http2_pool_fix() - __all__ = [ "types", "__version__", diff --git a/src/runloop_api_client/_base_client.py b/src/runloop_api_client/_base_client.py index bc636112d..cf03ac1c1 100644 --- a/src/runloop_api_client/_base_client.py +++ b/src/runloop_api_client/_base_client.py @@ -4,6 +4,7 @@ import json import time import uuid +import zlib import email import asyncio import inspect @@ -80,6 +81,8 @@ RAW_RESPONSE_HEADER, OVERRIDE_CAST_TO_HEADER, DEFAULT_CONNECTION_LIMITS, + DEFAULT_TRANSFER_POOL_SHARDS, + DEFAULT_BACKGROUND_POOL_SHARDS, ) from ._streaming import Stream, SSEDecoder, AsyncStream, SSEBytesDecoder from ._exceptions import ( @@ -174,20 +177,85 @@ async def aclose(self) -> None: weakref.WeakKeyDictionary() ) -# Separate HTTP/1.1 pool for bulk file upload/download. Keeping these off the -# main HTTP/2 connection avoids H2 session-window / stream-slot contention with -# latency-sensitive RPCs (create, wait_for_status, execute, …). -_shared_sync_transfer_transport: _SharedTransport | None = None +# Sharded H2 bulkheads: long-polls and file transfers stay off the API control-plane +# connection. Each shard index maps to its own shared transport (≈ one H2 connection). +# Removable once httpcore respects stream capacity when opening connections. +_shared_sync_background_transports: dict[int, _SharedTransport] = {} +_shared_sync_transfer_transports: dict[int, _SharedTransport] = {} +_shared_async_background_transports: weakref.WeakKeyDictionary[ + asyncio.AbstractEventLoop, dict[int, _SharedAsyncTransport] +] = weakref.WeakKeyDictionary() _shared_async_transfer_transports: weakref.WeakKeyDictionary[ - asyncio.AbstractEventLoop, _SharedAsyncTransport + asyncio.AbstractEventLoop, dict[int, _SharedAsyncTransport] ] = weakref.WeakKeyDictionary() -# Paths that carry large request/response bodies and should use the transfer pool. -_FILE_TRANSFER_PATH_SUFFIXES = ("/upload_file", "/download_file") +_BACKGROUND_PATH_SUFFIXES = ("/wait_for_status",) +_TRANSFER_PATH_SUFFIXES = ("/upload_file", "/download_file") + + +def _is_background_path(path: str) -> bool: + return path.endswith(_BACKGROUND_PATH_SUFFIXES) + +def _is_transfer_path(path: str) -> bool: + return path.endswith(_TRANSFER_PATH_SUFFIXES) + + +def _pool_affinity_key(path: str) -> str: + """Pick a stable resource id from the URL for shard routing.""" + parts = [p for p in path.split("/") if p] + try: + if "executions" in parts: + idx = parts.index("executions") + if idx + 1 < len(parts): + return parts[idx + 1] + if "devboxes" in parts: + idx = parts.index("devboxes") + if idx + 1 < len(parts): + return parts[idx + 1] + except ValueError: + pass + return path + + +def _shard_index(key: str, shards: int) -> int: + if shards <= 1: + return 0 + # crc32 is stable across processes (unlike PYTHONHASHSEED-randomized hash()). + return zlib.crc32(key.encode("utf-8")) % shards + + +def _acquire_shared_sync_transport(bucket: dict[int, _SharedTransport], shard: int) -> _SharedTransport: + with _pool_lock: + existing = bucket.get(shard) + if existing is not None and existing.acquire(): + return existing + transport = _SharedTransport( + httpx.HTTPTransport(limits=DEFAULT_CONNECTION_LIMITS, http2=True), + ) + bucket[shard] = transport + return transport + + +def _acquire_shared_async_transport( + by_loop: weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, dict[int, _SharedAsyncTransport]], + loop: asyncio.AbstractEventLoop, + shard: int, +) -> _SharedAsyncTransport: + with _pool_lock: + bucket = by_loop.get(loop) + if bucket is None: + bucket = {} + by_loop[loop] = bucket + existing = bucket.get(shard) + if existing is not None and existing.acquire(): + return existing + transport = _SharedAsyncTransport( + httpx.AsyncHTTPTransport(limits=DEFAULT_CONNECTION_LIMITS, http2=True), + ) + bucket[shard] = transport + return transport -def _is_file_transfer_path(path: str) -> bool: - return path.endswith(_FILE_TRANSFER_PATH_SUFFIXES) # TODO: make base page type vars covariant SyncPageT = TypeVar("SyncPageT", bound="BaseSyncPage[Any]") @@ -944,10 +1012,13 @@ def __del__(self) -> None: class SyncAPIClient(BaseClient[httpx.Client, Stream[Any]]): _client: httpx.Client - _transfer_client: httpx.Client | None + _background_clients: dict[int, httpx.Client] + _transfer_clients: dict[int, httpx.Client] _default_stream_cls: type[Stream[Any]] | None = None _uses_shared_pool: bool - _isolate_file_transfers: bool + _isolate_workload_pools: bool + _background_pool_shards: int + _transfer_pool_shards: int _closed: bool def __init__( @@ -962,6 +1033,8 @@ def __init__( custom_query: Mapping[str, object] | None = None, _strict_response_validation: bool, shared_http_pool: bool = True, + background_pool_shards: int = DEFAULT_BACKGROUND_POOL_SHARDS, + transfer_pool_shards: int = DEFAULT_TRANSFER_POOL_SHARDS, ) -> None: if not is_given(timeout): # if the user passed in a custom http client with a non-default @@ -981,6 +1054,11 @@ def __init__( f"Invalid `http_client` argument; Expected an instance of `httpx.Client` but got {type(http_client)}" ) + if background_pool_shards < 1: + raise ValueError("background_pool_shards must be >= 1") + if transfer_pool_shards < 1: + raise ValueError("transfer_pool_shards must be >= 1") + super().__init__( version=version, # cast to a valid type because mypy doesn't understand our type narrowing @@ -993,9 +1071,12 @@ def __init__( ) self._closed = False - self._transfer_client = None - # Custom http_client owns the full transport stack; don't invent a sibling pool. - self._isolate_file_transfers = http_client is None + self._background_clients = {} + self._transfer_clients = {} + self._background_pool_shards = background_pool_shards + self._transfer_pool_shards = transfer_pool_shards + # Custom http_client owns the full transport stack; don't invent sibling pools. + self._isolate_workload_pools = http_client is None if http_client is not None: self._client = http_client @@ -1020,36 +1101,56 @@ def __init__( ) self._uses_shared_pool = False - def _ensure_transfer_client(self) -> httpx.Client: - """Lazy HTTP/1.1 client for upload_file / download_file.""" - if self._transfer_client is not None: - return self._transfer_client - + def _make_bulkhead_client(self, *, transport: httpx.BaseTransport | None) -> httpx.Client: timeout = cast(Timeout, self.timeout) - if self._uses_shared_pool: - global _shared_sync_transfer_transport - with _pool_lock: - if _shared_sync_transfer_transport is None or not _shared_sync_transfer_transport.acquire(): - _shared_sync_transfer_transport = _SharedTransport( - httpx.HTTPTransport(limits=DEFAULT_CONNECTION_LIMITS, http2=False), - ) - self._transfer_client = SyncHttpxClientWrapper( + if transport is not None: + return SyncHttpxClientWrapper( base_url=self._base_url, timeout=timeout, - transport=_shared_sync_transfer_transport, - http2=False, + transport=transport, + ) + return SyncHttpxClientWrapper( + base_url=self._base_url, + timeout=timeout, + ) + + def _ensure_background_client(self, shard: int) -> httpx.Client: + existing = self._background_clients.get(shard) + if existing is not None: + return existing + if self._uses_shared_pool: + transport: httpx.BaseTransport | None = _acquire_shared_sync_transport( + _shared_sync_background_transports, shard ) else: - self._transfer_client = SyncHttpxClientWrapper( - base_url=self._base_url, - timeout=timeout, - http2=False, + transport = None + client = self._make_bulkhead_client(transport=transport) + self._background_clients[shard] = client + return client + + def _ensure_transfer_client(self, shard: int) -> httpx.Client: + existing = self._transfer_clients.get(shard) + if existing is not None: + return existing + if self._uses_shared_pool: + transport: httpx.BaseTransport | None = _acquire_shared_sync_transport( + _shared_sync_transfer_transports, shard ) - return self._transfer_client + else: + transport = None + client = self._make_bulkhead_client(transport=transport) + self._transfer_clients[shard] = client + return client def _send_client_for_request(self, request: httpx.Request) -> httpx.Client: - if self._isolate_file_transfers and _is_file_transfer_path(request.url.path): - return self._ensure_transfer_client() + if not self._isolate_workload_pools: + return self._client + path = request.url.path + key = _pool_affinity_key(path) + if _is_background_path(path): + return self._ensure_background_client(_shard_index(key, self._background_pool_shards)) + if _is_transfer_path(path): + return self._ensure_transfer_client(_shard_index(key, self._transfer_pool_shards)) return self._client def is_closed(self) -> bool: @@ -1066,10 +1167,12 @@ def close(self) -> None: return self._closed = True self._client.close() - transfer = self._transfer_client - self._transfer_client = None - if transfer is not None: - transfer.close() + for client in self._background_clients.values(): + client.close() + self._background_clients.clear() + for client in self._transfer_clients.values(): + client.close() + self._transfer_clients.clear() def __enter__(self: _T) -> _T: return self @@ -1617,10 +1720,13 @@ def __del__(self) -> None: class AsyncAPIClient(BaseClient[httpx.AsyncClient, AsyncStream[Any]]): _client: httpx.AsyncClient - _transfer_client: httpx.AsyncClient | None + _background_clients: dict[int, httpx.AsyncClient] + _transfer_clients: dict[int, httpx.AsyncClient] _default_stream_cls: type[AsyncStream[Any]] | None = None _uses_shared_pool: bool - _isolate_file_transfers: bool + _isolate_workload_pools: bool + _background_pool_shards: int + _transfer_pool_shards: int _closed: bool def __init__( @@ -1635,6 +1741,8 @@ def __init__( custom_headers: Mapping[str, str] | None = None, custom_query: Mapping[str, object] | None = None, shared_http_pool: bool = True, + background_pool_shards: int = DEFAULT_BACKGROUND_POOL_SHARDS, + transfer_pool_shards: int = DEFAULT_TRANSFER_POOL_SHARDS, ) -> None: if not is_given(timeout): # if the user passed in a custom http client with a non-default @@ -1654,6 +1762,11 @@ def __init__( f"Invalid `http_client` argument; Expected an instance of `httpx.AsyncClient` but got {type(http_client)}" ) + if background_pool_shards < 1: + raise ValueError("background_pool_shards must be >= 1") + if transfer_pool_shards < 1: + raise ValueError("transfer_pool_shards must be >= 1") + super().__init__( version=version, base_url=base_url, @@ -1666,9 +1779,12 @@ def __init__( ) self._closed = False - self._transfer_client = None - # Custom http_client owns the full transport stack; don't invent a sibling pool. - self._isolate_file_transfers = http_client is None + self._background_clients = {} + self._transfer_clients = {} + self._background_pool_shards = background_pool_shards + self._transfer_pool_shards = transfer_pool_shards + # Custom http_client owns the full transport stack; don't invent sibling pools. + self._isolate_workload_pools = http_client is None if http_client is not None: self._client = http_client @@ -1707,45 +1823,64 @@ def __init__( ) self._uses_shared_pool = False - def _ensure_transfer_client(self) -> httpx.AsyncClient: - """Lazy HTTP/1.1 client for upload_file / download_file.""" - if self._transfer_client is not None: - return self._transfer_client - + def _make_bulkhead_client(self, *, transport: httpx.AsyncBaseTransport | None) -> httpx.AsyncClient: timeout = cast(Timeout, self.timeout) + if transport is not None: + return AsyncHttpxClientWrapper( + base_url=self._base_url, + timeout=timeout, + transport=transport, + ) + return AsyncHttpxClientWrapper( + base_url=self._base_url, + timeout=timeout, + ) + + def _ensure_background_client(self, shard: int) -> httpx.AsyncClient: + existing = self._background_clients.get(shard) + if existing is not None: + return existing + transport: httpx.AsyncBaseTransport | None = None if self._uses_shared_pool: try: loop: asyncio.AbstractEventLoop | None = asyncio.get_running_loop() except RuntimeError: loop = None if loop is not None: - with _pool_lock: - existing = _shared_async_transfer_transports.get(loop) - if existing is not None and existing.acquire(): - transport: _SharedAsyncTransport = existing - else: - transport = _SharedAsyncTransport( - httpx.AsyncHTTPTransport(limits=DEFAULT_CONNECTION_LIMITS, http2=False), - ) - _shared_async_transfer_transports[loop] = transport - self._transfer_client = AsyncHttpxClientWrapper( - base_url=self._base_url, - timeout=timeout, - transport=transport, - http2=False, + transport = _acquire_shared_async_transport( + _shared_async_background_transports, loop, shard ) - return self._transfer_client - - self._transfer_client = AsyncHttpxClientWrapper( - base_url=self._base_url, - timeout=timeout, - http2=False, - ) - return self._transfer_client + client = self._make_bulkhead_client(transport=transport) + self._background_clients[shard] = client + return client + + def _ensure_transfer_client(self, shard: int) -> httpx.AsyncClient: + existing = self._transfer_clients.get(shard) + if existing is not None: + return existing + transport: httpx.AsyncBaseTransport | None = None + if self._uses_shared_pool: + try: + loop: asyncio.AbstractEventLoop | None = asyncio.get_running_loop() + except RuntimeError: + loop = None + if loop is not None: + transport = _acquire_shared_async_transport( + _shared_async_transfer_transports, loop, shard + ) + client = self._make_bulkhead_client(transport=transport) + self._transfer_clients[shard] = client + return client def _send_client_for_request(self, request: httpx.Request) -> httpx.AsyncClient: - if self._isolate_file_transfers and _is_file_transfer_path(request.url.path): - return self._ensure_transfer_client() + if not self._isolate_workload_pools: + return self._client + path = request.url.path + key = _pool_affinity_key(path) + if _is_background_path(path): + return self._ensure_background_client(_shard_index(key, self._background_pool_shards)) + if _is_transfer_path(path): + return self._ensure_transfer_client(_shard_index(key, self._transfer_pool_shards)) return self._client def is_closed(self) -> bool: @@ -1762,10 +1897,12 @@ async def close(self) -> None: return self._closed = True await self._client.aclose() - transfer = self._transfer_client - self._transfer_client = None - if transfer is not None: - await transfer.aclose() + for client in self._background_clients.values(): + await client.aclose() + self._background_clients.clear() + for client in self._transfer_clients.values(): + await client.aclose() + self._transfer_clients.clear() async def __aenter__(self: _T) -> _T: return self diff --git a/src/runloop_api_client/_client.py b/src/runloop_api_client/_client.py index ed052ec87..0526fadde 100644 --- a/src/runloop_api_client/_client.py +++ b/src/runloop_api_client/_client.py @@ -33,6 +33,7 @@ SyncAPIClient, AsyncAPIClient, ) +from ._constants import DEFAULT_BACKGROUND_POOL_SHARDS, DEFAULT_TRANSFER_POOL_SHARDS if TYPE_CHECKING: from .resources import ( @@ -96,6 +97,10 @@ def __init__( # Enables HTTP/2 multiplexing and avoids ConnectTimeout storms under high concurrency. # Set to False to create a private connection pool (old behavior). shared_http_pool: bool = True, + # Separate H2 connection pools for long-polls (/wait_for_status) and file + # transfers. Each shard ≈ one H2 connection, selected by hash(resource_id). + background_pool_shards: int = DEFAULT_BACKGROUND_POOL_SHARDS, + transfer_pool_shards: int = DEFAULT_TRANSFER_POOL_SHARDS, # Enable or disable schema validation for data returned by the API. # When enabled an error APIResponseValidationError is raised # if the API responds with invalid data for the expected schema. @@ -142,6 +147,8 @@ def __init__( custom_query=default_query, _strict_response_validation=_strict_response_validation, shared_http_pool=shared_http_pool, + background_pool_shards=background_pool_shards, + transfer_pool_shards=transfer_pool_shards, ) self._idempotency_header = "x-request-id" @@ -284,6 +291,8 @@ def copy( timeout: float | Timeout | None | NotGiven = not_given, http_client: httpx.Client | None = None, shared_http_pool: bool | None = None, + background_pool_shards: int | None = None, + transfer_pool_shards: int | None = None, max_retries: int | NotGiven = not_given, default_headers: Mapping[str, str] | None = None, set_default_headers: Mapping[str, str] | None = None, @@ -325,6 +334,14 @@ def copy( timeout=self.timeout if isinstance(timeout, NotGiven) else timeout, http_client=http_client, shared_http_pool=resolved_shared, + background_pool_shards=( + background_pool_shards + if background_pool_shards is not None + else self._background_pool_shards + ), + transfer_pool_shards=( + transfer_pool_shards if transfer_pool_shards is not None else self._transfer_pool_shards + ), max_retries=max_retries if is_given(max_retries) else self.max_retries, default_headers=headers, default_query=params, @@ -390,6 +407,10 @@ def __init__( # Enables HTTP/2 multiplexing and avoids ConnectTimeout storms under high concurrency. # Set to False to create a private connection pool (old behavior). shared_http_pool: bool = True, + # Separate H2 connection pools for long-polls (/wait_for_status) and file + # transfers. Each shard ≈ one H2 connection, selected by hash(resource_id). + background_pool_shards: int = DEFAULT_BACKGROUND_POOL_SHARDS, + transfer_pool_shards: int = DEFAULT_TRANSFER_POOL_SHARDS, # Enable or disable schema validation for data returned by the API. # When enabled an error APIResponseValidationError is raised # if the API responds with invalid data for the expected schema. @@ -436,6 +457,8 @@ def __init__( custom_query=default_query, _strict_response_validation=_strict_response_validation, shared_http_pool=shared_http_pool, + background_pool_shards=background_pool_shards, + transfer_pool_shards=transfer_pool_shards, ) self._idempotency_header = "x-request-id" @@ -578,6 +601,8 @@ def copy( timeout: float | Timeout | None | NotGiven = not_given, http_client: httpx.AsyncClient | None = None, shared_http_pool: bool | None = None, + background_pool_shards: int | None = None, + transfer_pool_shards: int | None = None, max_retries: int | NotGiven = not_given, default_headers: Mapping[str, str] | None = None, set_default_headers: Mapping[str, str] | None = None, @@ -619,6 +644,14 @@ def copy( timeout=self.timeout if isinstance(timeout, NotGiven) else timeout, http_client=http_client, shared_http_pool=resolved_shared, + background_pool_shards=( + background_pool_shards + if background_pool_shards is not None + else self._background_pool_shards + ), + transfer_pool_shards=( + transfer_pool_shards if transfer_pool_shards is not None else self._transfer_pool_shards + ), max_retries=max_retries if is_given(max_retries) else self.max_retries, default_headers=headers, default_query=params, diff --git a/src/runloop_api_client/_constants.py b/src/runloop_api_client/_constants.py index 83506afe3..b91e63d90 100644 --- a/src/runloop_api_client/_constants.py +++ b/src/runloop_api_client/_constants.py @@ -10,6 +10,12 @@ DEFAULT_MAX_RETRIES = 5 DEFAULT_CONNECTION_LIMITS = httpx.Limits(max_connections=20, max_keepalive_connections=10) +# Separate H2 connection pools for long-polls and file transfers. Each "shard" is +# its own httpx client / transport (≈ one H2 connection). Default 2 shards ≈ two +# H2 connections so waits/uploads do not share the API control-plane connection. +DEFAULT_BACKGROUND_POOL_SHARDS = 2 +DEFAULT_TRANSFER_POOL_SHARDS = 2 + INITIAL_RETRY_DELAY = 1.0 MAX_RETRY_DELAY = 60.0 diff --git a/src/runloop_api_client/_http2_pool_fix.py b/src/runloop_api_client/_http2_pool_fix.py deleted file mode 100644 index 4c8c24493..000000000 --- a/src/runloop_api_client/_http2_pool_fix.py +++ /dev/null @@ -1,185 +0,0 @@ -"""Make httpcore open additional HTTP/2 connections when stream slots run low. - -httpcore multiplexes HTTP/2 on a single connection and gates new streams with a -semaphore (default MAX_CONCURRENT_STREAMS=100). Its connection pool's -``is_available()`` ignores that limit, so once stream slots are exhausted, -further requests *block* on the semaphore instead of the pool opening another -connection — even when ``max_connections`` still has headroom. - -Additionally, the pool may assign a *burst* of requests to a connection that -only has a few free slots left (``is_available()`` is checked once per request -but slots are not reserved). Winners then share an overloaded connection while -the rest raise ``ConnectionNotAvailable`` and retry. - -This module patches sync + async HTTP/2 connections so that: - -1. ``is_available()`` is False when free stream slots are below a headroom - threshold, so the pool prefers opening/reusing another connection. -2. Stream-slot acquire is non-blocking; if a race still over-assigns, we raise - ``ConnectionNotAvailable`` and the pool retries on another connection. - -Idempotent: safe to call ``install()`` more than once. -""" - -from __future__ import annotations - -import logging -from typing import Any, Callable - -logger = logging.getLogger("runloop_api_client._http2_pool_fix") - -# If a connection has fewer free stream slots than this, treat it as unavailable -# so the pool opens another connection instead of stampeding the remainder. -# Must be >1: the pool assigns many queued requests to one "available" conn -# without reserving slots, so small remaining capacity still over-assigns. -_MIN_FREE_STREAM_SLOTS = 16 - -_installed = False - - -def _free_stream_slots(connection: Any) -> int | None: - """Return free H2 stream slots, or None if unknown.""" - max_streams = getattr(connection, "_max_streams", None) - events = getattr(connection, "_events", None) - if not isinstance(max_streams, int) or max_streams <= 0 or events is None: - return None - return max_streams - len(events) - - -def _stream_slots_saturated(connection: Any) -> bool: - """True when the connection should not accept more streams.""" - free = _free_stream_slots(connection) - if free is None: - return False - return free < _MIN_FREE_STREAM_SLOTS - - -def _patch_is_available(cls: type) -> None: - if getattr(cls.is_available, "_runloop_stream_overflow_patched", False): - return - - original: Callable[[Any], bool] = cls.is_available - - def is_available(self: Any) -> bool: - if not original(self): - return False - return not _stream_slots_saturated(self) - - is_available._runloop_stream_overflow_patched = True # type: ignore[attr-defined] - cls.is_available = is_available # type: ignore[method-assign] - - -def _make_async_nonblocking_acquire(sem: Any, connection_not_available: type) -> Callable[[], Any]: - original_acquire = sem.acquire - - async def acquire() -> None: - if not getattr(sem, "_backend", ""): - sem.setup() - - if sem._backend == "asyncio": - inner = sem._anyio_semaphore - try: - inner.acquire_nowait() - return - except Exception: - raise connection_not_available( - "HTTP/2 connection has no free stream slots" - ) from None - - if sem._backend == "trio": - inner = sem._trio_semaphore - acquire_nowait = getattr(inner, "acquire_nowait", None) - if acquire_nowait is not None: - try: - acquire_nowait() - return - except Exception: - raise connection_not_available( - "HTTP/2 connection has no free stream slots" - ) from None - - await original_acquire() - - return acquire - - -def _make_sync_nonblocking_acquire(sem: Any, connection_not_available: type) -> Callable[[], None]: - inner = getattr(sem, "_semaphore", sem) - - def acquire() -> None: - ok = inner.acquire(False) - if not ok: - raise connection_not_available("HTTP/2 connection has no free stream slots") - - return acquire - - -def _patch_async_connection(cls: type) -> None: - from httpcore import ConnectionNotAvailable - - if getattr(cls.handle_async_request, "_runloop_stream_overflow_patched", False): - return - - original = cls.handle_async_request - - async def handle_async_request(self: Any, request: Any) -> Any: - sem = getattr(self, "_max_streams_semaphore", None) - if sem is None: - return await original(self, request) - - real_acquire = sem.acquire - sem.acquire = _make_async_nonblocking_acquire(sem, ConnectionNotAvailable) - try: - return await original(self, request) - finally: - sem.acquire = real_acquire - - handle_async_request._runloop_stream_overflow_patched = True # type: ignore[attr-defined] - cls.handle_async_request = handle_async_request # type: ignore[method-assign] - - -def _patch_sync_connection(cls: type) -> None: - from httpcore import ConnectionNotAvailable - - if getattr(cls.handle_request, "_runloop_stream_overflow_patched", False): - return - - original = cls.handle_request - - def handle_request(self: Any, request: Any) -> Any: - sem = getattr(self, "_max_streams_semaphore", None) - if sem is None: - return original(self, request) - - real_acquire = sem.acquire - sem.acquire = _make_sync_nonblocking_acquire(sem, ConnectionNotAvailable) - try: - return original(self, request) - finally: - sem.acquire = real_acquire - - handle_request._runloop_stream_overflow_patched = True # type: ignore[attr-defined] - cls.handle_request = handle_request # type: ignore[method-assign] - - -def install() -> bool: - """Patch httpcore HTTP/2 connections. Returns True if a patch was applied.""" - global _installed - if _installed: - return False - - try: - from httpcore._async.http2 import AsyncHTTP2Connection - from httpcore._sync.http2 import HTTP2Connection - except ImportError as exc: # pragma: no cover - logger.warning("http2 stream-overflow patch skipped: %s", exc) - return False - - _patch_is_available(AsyncHTTP2Connection) - _patch_is_available(HTTP2Connection) - _patch_async_connection(AsyncHTTP2Connection) - _patch_sync_connection(HTTP2Connection) - - _installed = True - logger.debug("installed httpcore HTTP/2 stream-overflow connection patch") - return True diff --git a/src/runloop_api_client/sdk/async_.py b/src/runloop_api_client/sdk/async_.py index 6319e02bd..41f10cd1f 100644 --- a/src/runloop_api_client/sdk/async_.py +++ b/src/runloop_api_client/sdk/async_.py @@ -40,6 +40,7 @@ ) from .._types import Timeout, NotGiven, not_given from .._client import DEFAULT_MAX_RETRIES, AsyncRunloop +from .._constants import DEFAULT_BACKGROUND_POOL_SHARDS, DEFAULT_TRANSFER_POOL_SHARDS from ._helpers import detect_content_type from .async_axon import AsyncAxon from .async_agent import AsyncAgent @@ -1329,6 +1330,8 @@ def __init__( default_headers: Mapping[str, str] | None = None, default_query: Mapping[str, object] | None = None, http_client: httpx.AsyncClient | None = None, + background_pool_shards: int = DEFAULT_BACKGROUND_POOL_SHARDS, + transfer_pool_shards: int = DEFAULT_TRANSFER_POOL_SHARDS, ) -> None: """Configure the asynchronous SDK wrapper. @@ -1346,6 +1349,10 @@ def __init__( :type default_query: Mapping[str, object] | None, optional :param http_client: Custom ``httpx.AsyncClient`` instance to reuse, defaults to None :type http_client: httpx.AsyncClient | None, optional + :param background_pool_shards: H2 shards for long-polls, defaults to 2 + :type background_pool_shards: int, optional + :param transfer_pool_shards: H2 shards for upload/download, defaults to 2 + :type transfer_pool_shards: int, optional """ self.api = AsyncRunloop( bearer_token=bearer_token, @@ -1355,6 +1362,8 @@ def __init__( default_headers=default_headers, default_query=default_query, http_client=http_client, + background_pool_shards=background_pool_shards, + transfer_pool_shards=transfer_pool_shards, ) self.agent = AsyncAgentOps(self.api) diff --git a/src/runloop_api_client/sdk/sync.py b/src/runloop_api_client/sdk/sync.py index 7d3afcb24..db45f5379 100644 --- a/src/runloop_api_client/sdk/sync.py +++ b/src/runloop_api_client/sdk/sync.py @@ -44,6 +44,7 @@ from .secret import Secret from .._types import Timeout, NotGiven, not_given from .._client import DEFAULT_MAX_RETRIES, Runloop +from .._constants import DEFAULT_BACKGROUND_POOL_SHARDS, DEFAULT_TRANSFER_POOL_SHARDS from ._helpers import detect_content_type from .scenario import Scenario from .snapshot import Snapshot @@ -1354,6 +1355,8 @@ def __init__( default_headers: Mapping[str, str] | None = None, default_query: Mapping[str, object] | None = None, http_client: httpx.Client | None = None, + background_pool_shards: int = DEFAULT_BACKGROUND_POOL_SHARDS, + transfer_pool_shards: int = DEFAULT_TRANSFER_POOL_SHARDS, ) -> None: """Configure the synchronous SDK wrapper. @@ -1371,6 +1374,10 @@ def __init__( :type default_query: Mapping[str, object] | None, optional :param http_client: Custom ``httpx.Client`` instance to reuse, defaults to None :type http_client: httpx.Client | None, optional + :param background_pool_shards: H2 shards for long-polls, defaults to 2 + :type background_pool_shards: int, optional + :param transfer_pool_shards: H2 shards for upload/download, defaults to 2 + :type transfer_pool_shards: int, optional """ self.api = Runloop( bearer_token=bearer_token, @@ -1380,6 +1387,8 @@ def __init__( default_headers=default_headers, default_query=default_query, http_client=http_client, + background_pool_shards=background_pool_shards, + transfer_pool_shards=transfer_pool_shards, ) self.agent = AgentOps(self.api) diff --git a/tests/test_http2_pool_fix.py b/tests/test_http2_pool_fix.py deleted file mode 100644 index aa5135ab9..000000000 --- a/tests/test_http2_pool_fix.py +++ /dev/null @@ -1,72 +0,0 @@ -"""Tests for HTTP/2 stream-overflow → new connection behavior.""" - -from __future__ import annotations - -from typing import Any -from unittest.mock import MagicMock - -import pytest - -from runloop_api_client._http2_pool_fix import ( - _stream_slots_saturated, - install, -) - - -def test_install_is_idempotent() -> None: - # Package import already installs the patch. - assert install() is False - - -def test_stream_slots_saturated_helper() -> None: - conn: Any = MagicMock() - conn._max_streams = 100 - conn._events = {i: [] for i in range(1, 90)} # 89 in flight → 11 free < 16 - assert _stream_slots_saturated(conn) is True - - conn._events = {i: [] for i in range(1, 80)} # 79 in flight → 21 free - assert _stream_slots_saturated(conn) is False - - conn._max_streams = 0 - assert _stream_slots_saturated(conn) is False - - -def _bare_async_h2_connection() -> Any: - from httpcore._async.http2 import AsyncHTTP2Connection, HTTPConnectionState - - conn = object.__new__(AsyncHTTP2Connection) - conn._state = HTTPConnectionState.ACTIVE - conn._connection_error = False - conn._used_all_stream_ids = False - conn._h2_state = MagicMock() - # Anything other than ConnectionState.CLOSED. - conn._h2_state.state_machine.state = object() - return conn - - -def test_is_available_false_when_streams_nearly_full() -> None: - conn = _bare_async_h2_connection() - conn._max_streams = 100 - conn._events = {i: [] for i in range(1, 95)} # 6 free - assert conn.is_available() is False - - -def test_is_available_true_when_stream_slots_remain() -> None: - conn = _bare_async_h2_connection() - conn._max_streams = 100 - conn._events = {i: [] for i in range(1, 50)} # 51 free - assert conn.is_available() is True - - -def test_sync_is_available_false_when_streams_nearly_full() -> None: - from httpcore._sync.http2 import HTTP2Connection, HTTPConnectionState - - conn = object.__new__(HTTP2Connection) - conn._state = HTTPConnectionState.ACTIVE - conn._connection_error = False - conn._used_all_stream_ids = False - conn._h2_state = MagicMock() - conn._h2_state.state_machine.state = object() - conn._max_streams = 100 - conn._events = {i: [] for i in range(1, 95)} - assert conn.is_available() is False diff --git a/tests/test_shared_pool.py b/tests/test_shared_pool.py index 982350f07..b798b5baf 100644 --- a/tests/test_shared_pool.py +++ b/tests/test_shared_pool.py @@ -30,12 +30,15 @@ def _reset_shared_pool() -> Iterator[None]: # pyright: ignore[reportUnusedFunct def _clear_pool_state() -> None: with _base_mod._pool_lock: old_sync = _base_mod._shared_sync_transport - old_sync_transfer = _base_mod._shared_sync_transfer_transport + old_bg = list(_base_mod._shared_sync_background_transports.values()) + old_xfer = list(_base_mod._shared_sync_transfer_transports.values()) _base_mod._shared_sync_transport = None - _base_mod._shared_sync_transfer_transport = None + _base_mod._shared_sync_background_transports.clear() + _base_mod._shared_sync_transfer_transports.clear() _base_mod._shared_async_transports.clear() + _base_mod._shared_async_background_transports.clear() _base_mod._shared_async_transfer_transports.clear() - for transport in (old_sync, old_sync_transfer): + for transport in [old_sync, *old_bg, *old_xfer]: if transport is not None: try: transport._transport.close() diff --git a/tests/test_transfer_client.py b/tests/test_transfer_client.py index 2ca84ec6e..6d371ff6a 100644 --- a/tests/test_transfer_client.py +++ b/tests/test_transfer_client.py @@ -1,4 +1,4 @@ -"""Tests for dedicated HTTP/1.1 transfer pool used by upload/download.""" +"""Tests for sharded H2 background + transfer bulkhead pools.""" from __future__ import annotations @@ -10,7 +10,12 @@ import runloop_api_client._base_client as _base_mod from runloop_api_client import Runloop, AsyncRunloop -from runloop_api_client._base_client import _is_file_transfer_path +from runloop_api_client._base_client import ( + _is_background_path, + _is_transfer_path, + _pool_affinity_key, + _shard_index, +) base_url = os.environ.get("TEST_API_BASE_URL", "http://127.0.0.1:4010") bearer_token = "My Bearer Token" @@ -26,12 +31,15 @@ def _reset_shared_pool() -> Iterator[None]: # pyright: ignore[reportUnusedFunct def _clear_pool_state() -> None: with _base_mod._pool_lock: old_sync = _base_mod._shared_sync_transport - old_sync_transfer = _base_mod._shared_sync_transfer_transport + old_bg = list(_base_mod._shared_sync_background_transports.values()) + old_xfer = list(_base_mod._shared_sync_transfer_transports.values()) _base_mod._shared_sync_transport = None - _base_mod._shared_sync_transfer_transport = None + _base_mod._shared_sync_background_transports.clear() + _base_mod._shared_sync_transfer_transports.clear() _base_mod._shared_async_transports.clear() + _base_mod._shared_async_background_transports.clear() _base_mod._shared_async_transfer_transports.clear() - for transport in (old_sync, old_sync_transfer): + for transport in [old_sync, *old_bg, *old_xfer]: if transport is not None: try: transport._transport.close() @@ -45,89 +53,119 @@ def _make_client(**kwargs: Any) -> Runloop: return Runloop(**kwargs) -def test_is_file_transfer_path() -> None: - assert _is_file_transfer_path("/v1/devboxes/dbx_1/upload_file") is True - assert _is_file_transfer_path("/v1/devboxes/dbx_1/download_file") is True - assert _is_file_transfer_path("/v1/devboxes") is False - assert _is_file_transfer_path("/v1/objects/obj_1/download") is False - assert _is_file_transfer_path("/v1/scenarios/runs/run_1/download_logs") is False +def test_path_classification() -> None: + assert _is_background_path("/v1/devboxes/dbx_1/wait_for_status") is True + assert _is_background_path("/v1/devboxes/dbx_1/executions/ex_1/wait_for_status") is True + assert _is_transfer_path("/v1/devboxes/dbx_1/upload_file") is True + assert _is_transfer_path("/v1/devboxes/dbx_1/download_file") is True + assert _is_background_path("/v1/devboxes") is False + assert _is_transfer_path("/v1/objects/obj_1/download") is False -def test_transfer_client_is_http1_and_separate_from_api_pool() -> None: - client = _make_client(shared_http_pool=True) - try: - assert client._isolate_file_transfers is True - assert client._transfer_client is None +def test_affinity_key_and_shard() -> None: + assert _pool_affinity_key("/v1/devboxes/dbx_1/upload_file") == "dbx_1" + assert _pool_affinity_key("/v1/devboxes/dbx_1/wait_for_status") == "dbx_1" + assert _pool_affinity_key("/v1/devboxes/dbx_1/executions/ex_9/wait_for_status") == "ex_9" + assert _shard_index("dbx_1", 1) == 0 + assert _shard_index("dbx_1", 2) in (0, 1) + assert _shard_index("dbx_1", 2) == _shard_index("dbx_1", 2) - api_request = httpx.Request("POST", f"{base_url}/v1/devboxes") - transfer_request = httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_1/upload_file") - assert client._send_client_for_request(api_request) is client._client +def test_api_background_transfer_use_distinct_transports() -> None: + client = _make_client(shared_http_pool=True, background_pool_shards=2, transfer_pool_shards=2) + try: + api_req = httpx.Request("POST", f"{base_url}/v1/devboxes") + wait_req = httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_a/wait_for_status") + upload_req = httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_a/upload_file") + + api = client._send_client_for_request(api_req) + wait = client._send_client_for_request(wait_req) + upload = client._send_client_for_request(upload_req) + + assert api is client._client + assert wait is not api + assert upload is not api + assert wait is not upload + assert wait._transport is not api._transport # type: ignore[attr-defined] + assert upload._transport is not api._transport # type: ignore[attr-defined] + assert wait._transport is not upload._transport # type: ignore[attr-defined] + finally: + client.close() - transfer = client._send_client_for_request(transfer_request) - assert transfer is not client._client - assert transfer is client._transfer_client - api_transport = client._client._transport # type: ignore[attr-defined] - transfer_transport = transfer._transport # type: ignore[attr-defined] - assert api_transport is not transfer_transport - assert _base_mod._shared_sync_transfer_transport is transfer_transport - # Real httpx transport under the shared wrapper is HTTP/1.1-only. - inner = transfer_transport._transport - assert getattr(inner, "_http2", False) is False +def test_same_resource_affinity_uses_same_shard() -> None: + client = _make_client(background_pool_shards=2, transfer_pool_shards=2) + try: + w1 = client._send_client_for_request( + httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_same/wait_for_status") + ) + w2 = client._send_client_for_request( + httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_same/wait_for_status") + ) + assert w1 is w2 finally: client.close() -def test_transfer_pool_is_shared_across_sdk_clients() -> None: - c1 = _make_client(shared_http_pool=True) - c2 = _make_client(shared_http_pool=True) +def test_different_resources_can_land_on_different_shards() -> None: + client = _make_client(background_pool_shards=2) try: - req = httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_1/download_file") - t1 = c1._send_client_for_request(req) - t2 = c2._send_client_for_request(req) - assert t1 is not t2 - assert t1._transport is t2._transport # type: ignore[attr-defined] - assert _base_mod._shared_sync_transfer_transport is not None - assert _base_mod._shared_sync_transfer_transport.refcount == 2 + seen: set[int] = set() + for i in range(40): + req = httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_{i}/wait_for_status") + c = client._send_client_for_request(req) + seen.add(id(c._transport)) # type: ignore[attr-defined] + assert len(seen) == 2 + assert len(_base_mod._shared_sync_background_transports) == 2 finally: - c1.close() - c2.close() - assert _base_mod._shared_sync_transfer_transport is not None - assert _base_mod._shared_sync_transfer_transport.refcount == 0 + client.close() -def test_custom_http_client_skips_transfer_isolation() -> None: +def test_custom_http_client_skips_isolation() -> None: custom = httpx.Client() client = _make_client(http_client=custom) try: - assert client._isolate_file_transfers is False - req = httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_1/upload_file") + assert client._isolate_workload_pools is False + req = httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_1/wait_for_status") assert client._send_client_for_request(req) is custom - assert client._transfer_client is None + assert client._background_clients == {} finally: client.close() custom.close() -def test_private_pool_still_isolates_transfers() -> None: - client = _make_client(shared_http_pool=False) +def test_shards_shared_across_sdk_clients() -> None: + c1 = _make_client(background_pool_shards=2) + c2 = _make_client(background_pool_shards=2) try: - req = httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_1/upload_file") - transfer = client._send_client_for_request(req) - assert transfer is not client._client - assert client._uses_shared_pool is False + req = httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_shared/wait_for_status") + t1 = c1._send_client_for_request(req) + t2 = c2._send_client_for_request(req) + assert t1 is not t2 + assert t1._transport is t2._transport # type: ignore[attr-defined] finally: - client.close() + c1.close() + c2.close() @pytest.mark.asyncio -async def test_async_transfer_client_is_separate() -> None: - client = AsyncRunloop(base_url=base_url, bearer_token=bearer_token, shared_http_pool=True) +async def test_async_bulkheads() -> None: + client = AsyncRunloop( + base_url=base_url, + bearer_token=bearer_token, + shared_http_pool=True, + background_pool_shards=2, + transfer_pool_shards=2, + ) try: - req = httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_1/upload_file") - transfer = client._send_client_for_request(req) - assert transfer is not client._client - assert client._transfer_client is transfer + wait = client._send_client_for_request( + httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_1/executions/ex_1/wait_for_status") + ) + upload = client._send_client_for_request( + httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_1/upload_file") + ) + assert wait is not client._client + assert upload is not client._client + assert wait is not upload finally: await client.close() From 6783ef62858fac9676ae115528b0d61d40714083 Mon Sep 17 00:00:00 2001 From: Reflex Date: Mon, 27 Jul 2026 23:11:12 +0000 Subject: [PATCH 03/11] chore: fix ruff import sorting for pool bulkhead changes Co-authored-by: Cursor --- src/runloop_api_client/_client.py | 2 +- src/runloop_api_client/sdk/async_.py | 2 +- src/runloop_api_client/sdk/sync.py | 2 +- tests/test_transfer_client.py | 4 ++-- 4 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/runloop_api_client/_client.py b/src/runloop_api_client/_client.py index 0526fadde..c06c66a1d 100644 --- a/src/runloop_api_client/_client.py +++ b/src/runloop_api_client/_client.py @@ -26,6 +26,7 @@ ) from ._compat import cached_property from ._version import __version__ +from ._constants import DEFAULT_TRANSFER_POOL_SHARDS, DEFAULT_BACKGROUND_POOL_SHARDS from ._streaming import Stream as Stream, AsyncStream as AsyncStream from ._exceptions import RunloopError, APIStatusError from ._base_client import ( @@ -33,7 +34,6 @@ SyncAPIClient, AsyncAPIClient, ) -from ._constants import DEFAULT_BACKGROUND_POOL_SHARDS, DEFAULT_TRANSFER_POOL_SHARDS if TYPE_CHECKING: from .resources import ( diff --git a/src/runloop_api_client/sdk/async_.py b/src/runloop_api_client/sdk/async_.py index 41f10cd1f..2a28b86e4 100644 --- a/src/runloop_api_client/sdk/async_.py +++ b/src/runloop_api_client/sdk/async_.py @@ -40,9 +40,9 @@ ) from .._types import Timeout, NotGiven, not_given from .._client import DEFAULT_MAX_RETRIES, AsyncRunloop -from .._constants import DEFAULT_BACKGROUND_POOL_SHARDS, DEFAULT_TRANSFER_POOL_SHARDS from ._helpers import detect_content_type from .async_axon import AsyncAxon +from .._constants import DEFAULT_TRANSFER_POOL_SHARDS, DEFAULT_BACKGROUND_POOL_SHARDS from .async_agent import AsyncAgent from .async_devbox import AsyncDevbox from .async_scorer import AsyncScorer diff --git a/src/runloop_api_client/sdk/sync.py b/src/runloop_api_client/sdk/sync.py index db45f5379..0f23b105f 100644 --- a/src/runloop_api_client/sdk/sync.py +++ b/src/runloop_api_client/sdk/sync.py @@ -44,13 +44,13 @@ from .secret import Secret from .._types import Timeout, NotGiven, not_given from .._client import DEFAULT_MAX_RETRIES, Runloop -from .._constants import DEFAULT_BACKGROUND_POOL_SHARDS, DEFAULT_TRANSFER_POOL_SHARDS from ._helpers import detect_content_type from .scenario import Scenario from .snapshot import Snapshot from .benchmark import Benchmark from .blueprint import Blueprint from .mcp_config import McpConfig +from .._constants import DEFAULT_TRANSFER_POOL_SHARDS, DEFAULT_BACKGROUND_POOL_SHARDS from .gateway_config import GatewayConfig from .network_policy import NetworkPolicy from .storage_object import StorageObject diff --git a/tests/test_transfer_client.py b/tests/test_transfer_client.py index 6d371ff6a..6d265eb7f 100644 --- a/tests/test_transfer_client.py +++ b/tests/test_transfer_client.py @@ -11,10 +11,10 @@ import runloop_api_client._base_client as _base_mod from runloop_api_client import Runloop, AsyncRunloop from runloop_api_client._base_client import ( - _is_background_path, + _shard_index, _is_transfer_path, _pool_affinity_key, - _shard_index, + _is_background_path, ) base_url = os.environ.get("TEST_API_BASE_URL", "http://127.0.0.1:4010") From bb093914033b97f207af0d6bf3ca61de5efd02c8 Mon Sep 17 00:00:00 2001 From: Reflex Date: Mon, 27 Jul 2026 23:17:29 +0000 Subject: [PATCH 04/11] chore: apply ruff format to bulkhead pool changes Co-authored-by: Cursor --- src/runloop_api_client/_base_client.py | 8 ++------ src/runloop_api_client/_client.py | 8 ++------ tests/test_transfer_client.py | 12 +++--------- 3 files changed, 7 insertions(+), 21 deletions(-) diff --git a/src/runloop_api_client/_base_client.py b/src/runloop_api_client/_base_client.py index cf03ac1c1..266e2844f 100644 --- a/src/runloop_api_client/_base_client.py +++ b/src/runloop_api_client/_base_client.py @@ -1847,9 +1847,7 @@ def _ensure_background_client(self, shard: int) -> httpx.AsyncClient: except RuntimeError: loop = None if loop is not None: - transport = _acquire_shared_async_transport( - _shared_async_background_transports, loop, shard - ) + transport = _acquire_shared_async_transport(_shared_async_background_transports, loop, shard) client = self._make_bulkhead_client(transport=transport) self._background_clients[shard] = client return client @@ -1865,9 +1863,7 @@ def _ensure_transfer_client(self, shard: int) -> httpx.AsyncClient: except RuntimeError: loop = None if loop is not None: - transport = _acquire_shared_async_transport( - _shared_async_transfer_transports, loop, shard - ) + transport = _acquire_shared_async_transport(_shared_async_transfer_transports, loop, shard) client = self._make_bulkhead_client(transport=transport) self._transfer_clients[shard] = client return client diff --git a/src/runloop_api_client/_client.py b/src/runloop_api_client/_client.py index c06c66a1d..7d6fd924d 100644 --- a/src/runloop_api_client/_client.py +++ b/src/runloop_api_client/_client.py @@ -335,9 +335,7 @@ def copy( http_client=http_client, shared_http_pool=resolved_shared, background_pool_shards=( - background_pool_shards - if background_pool_shards is not None - else self._background_pool_shards + background_pool_shards if background_pool_shards is not None else self._background_pool_shards ), transfer_pool_shards=( transfer_pool_shards if transfer_pool_shards is not None else self._transfer_pool_shards @@ -645,9 +643,7 @@ def copy( http_client=http_client, shared_http_pool=resolved_shared, background_pool_shards=( - background_pool_shards - if background_pool_shards is not None - else self._background_pool_shards + background_pool_shards if background_pool_shards is not None else self._background_pool_shards ), transfer_pool_shards=( transfer_pool_shards if transfer_pool_shards is not None else self._transfer_pool_shards diff --git a/tests/test_transfer_client.py b/tests/test_transfer_client.py index 6d265eb7f..8a4faad5a 100644 --- a/tests/test_transfer_client.py +++ b/tests/test_transfer_client.py @@ -96,12 +96,8 @@ def test_api_background_transfer_use_distinct_transports() -> None: def test_same_resource_affinity_uses_same_shard() -> None: client = _make_client(background_pool_shards=2, transfer_pool_shards=2) try: - w1 = client._send_client_for_request( - httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_same/wait_for_status") - ) - w2 = client._send_client_for_request( - httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_same/wait_for_status") - ) + w1 = client._send_client_for_request(httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_same/wait_for_status")) + w2 = client._send_client_for_request(httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_same/wait_for_status")) assert w1 is w2 finally: client.close() @@ -161,9 +157,7 @@ async def test_async_bulkheads() -> None: wait = client._send_client_for_request( httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_1/executions/ex_1/wait_for_status") ) - upload = client._send_client_for_request( - httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_1/upload_file") - ) + upload = client._send_client_for_request(httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_1/upload_file")) assert wait is not client._client assert upload is not client._client assert wait is not upload From 76a9ecbe225fa58e84f9a52185f1e5b7b67e2dcc Mon Sep 17 00:00:00 2001 From: Reflex Date: Mon, 27 Jul 2026 23:26:27 +0000 Subject: [PATCH 05/11] test: prove H2 bulkheads on the wire; lock sync lazy init Add TLS+ALPN integration tests that saturate waits / stall upload bodies and assert create lands on a different connection and stays fast. Guard sync bulkhead client creation with a lock and enumerate long-lived resource paths so new endpoints cannot silently join the API pool. Co-authored-by: Cursor --- src/runloop_api_client/_base_client.py | 45 +-- tests/test_h2_bulkhead_integration.py | 416 +++++++++++++++++++++++++ tests/test_transfer_client.py | 72 +++++ 3 files changed, 515 insertions(+), 18 deletions(-) create mode 100644 tests/test_h2_bulkhead_integration.py diff --git a/src/runloop_api_client/_base_client.py b/src/runloop_api_client/_base_client.py index 266e2844f..25b73e816 100644 --- a/src/runloop_api_client/_base_client.py +++ b/src/runloop_api_client/_base_client.py @@ -1073,6 +1073,7 @@ def __init__( self._closed = False self._background_clients = {} self._transfer_clients = {} + self._bulkhead_lock = threading.Lock() self._background_pool_shards = background_pool_shards self._transfer_pool_shards = transfer_pool_shards # Custom http_client owns the full transport stack; don't invent sibling pools. @@ -1118,29 +1119,37 @@ def _ensure_background_client(self, shard: int) -> httpx.Client: existing = self._background_clients.get(shard) if existing is not None: return existing - if self._uses_shared_pool: - transport: httpx.BaseTransport | None = _acquire_shared_sync_transport( - _shared_sync_background_transports, shard - ) - else: - transport = None - client = self._make_bulkhead_client(transport=transport) - self._background_clients[shard] = client - return client + with self._bulkhead_lock: + existing = self._background_clients.get(shard) + if existing is not None: + return existing + if self._uses_shared_pool: + transport: httpx.BaseTransport | None = _acquire_shared_sync_transport( + _shared_sync_background_transports, shard + ) + else: + transport = None + client = self._make_bulkhead_client(transport=transport) + self._background_clients[shard] = client + return client def _ensure_transfer_client(self, shard: int) -> httpx.Client: existing = self._transfer_clients.get(shard) if existing is not None: return existing - if self._uses_shared_pool: - transport: httpx.BaseTransport | None = _acquire_shared_sync_transport( - _shared_sync_transfer_transports, shard - ) - else: - transport = None - client = self._make_bulkhead_client(transport=transport) - self._transfer_clients[shard] = client - return client + with self._bulkhead_lock: + existing = self._transfer_clients.get(shard) + if existing is not None: + return existing + if self._uses_shared_pool: + transport: httpx.BaseTransport | None = _acquire_shared_sync_transport( + _shared_sync_transfer_transports, shard + ) + else: + transport = None + client = self._make_bulkhead_client(transport=transport) + self._transfer_clients[shard] = client + return client def _send_client_for_request(self, request: httpx.Request) -> httpx.Client: if not self._isolate_workload_pools: diff --git a/tests/test_h2_bulkhead_integration.py b/tests/test_h2_bulkhead_integration.py new file mode 100644 index 000000000..c7f3896c7 --- /dev/null +++ b/tests/test_h2_bulkhead_integration.py @@ -0,0 +1,416 @@ +"""TLS+ALPN HTTP/2 integration tests for workload bulkhead pools. + +Proves wire behavior (not just Python object identity): +- Blocking waits and creates land on different server connections. +- Create completes promptly even when background streams are saturated. +- Upload body stall does not delay create on the API pool. +""" + +from __future__ import annotations + +import ssl +import json +import time +import asyncio +import tempfile +import threading +import subprocess +from typing import Any, Iterator +from pathlib import Path + +import httpx +import pytest + +h2 = pytest.importorskip("h2") +import h2.config # noqa: E402 +import h2.events # noqa: E402 +import h2.settings # noqa: E402 +import h2.connection # noqa: E402 + +import runloop_api_client._base_client as _base_mod +from runloop_api_client import AsyncRunloop + +pytestmark = pytest.mark.timeout(30) + + +def _clear_pool_state() -> None: + with _base_mod._pool_lock: + old_sync = _base_mod._shared_sync_transport + old_bg = list(_base_mod._shared_sync_background_transports.values()) + old_xfer = list(_base_mod._shared_sync_transfer_transports.values()) + _base_mod._shared_sync_transport = None + _base_mod._shared_sync_background_transports.clear() + _base_mod._shared_sync_transfer_transports.clear() + _base_mod._shared_async_transports.clear() + _base_mod._shared_async_background_transports.clear() + _base_mod._shared_async_transfer_transports.clear() + for transport in [old_sync, *old_bg, *old_xfer]: + if transport is not None: + try: + transport._transport.close() + except Exception: + pass + + +@pytest.fixture(autouse=True) +def _reset_pools() -> Iterator[None]: + _clear_pool_state() + yield + _clear_pool_state() + + +def _ensure_certs(dir_path: Path) -> tuple[Path, Path]: + cert = dir_path / "cert.pem" + key = dir_path / "key.pem" + if not cert.exists(): + subprocess.run( + [ + "openssl", + "req", + "-x509", + "-newkey", + "rsa:2048", + "-keyout", + str(key), + "-out", + str(cert), + "-days", + "1", + "-nodes", + "-subj", + "/CN=localhost", + ], + check=True, + capture_output=True, + ) + return cert, key + + +class _H2BulkheadServer: + """Minimal TLS+ALPN h2 server that records per-connection request timing.""" + + def __init__(self, *, max_concurrent_streams: int = 2, upload_stall_s: float = 0.0) -> None: + self.max_concurrent_streams = max_concurrent_streams + self.upload_stall_s = upload_stall_s + self.observations: list[dict[str, Any]] = [] + self.active_waits = 0 + self.active_uploads = 0 + self.release_upload_credit = asyncio.Event() + self._conn_seq = 0 + self._conn_lock = threading.Lock() + self._server: asyncio.Server | None = None + self._cert_dir = Path(tempfile.mkdtemp(prefix="h2-bulkhead-")) + self.cert, self.key = _ensure_certs(self._cert_dir) + self.host = "127.0.0.1" + self.port = 0 + + def _next_conn_id(self) -> int: + with self._conn_lock: + self._conn_seq += 1 + return self._conn_seq + + async def start(self) -> str: + ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + ctx.load_cert_chain(str(self.cert), str(self.key)) + ctx.set_alpn_protocols(["h2"]) + self._server = await asyncio.start_server(self._handle, self.host, 0, ssl=ctx) + sockets = self._server.sockets + assert sockets + self.port = int(sockets[0].getsockname()[1]) + return f"https://{self.host}:{self.port}" + + async def stop(self) -> None: + if self._server is not None: + self._server.close() + await self._server.wait_closed() + self._server = None + + async def _handle(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + conn_id = self._next_conn_id() + config = h2.config.H2Configuration(client_side=False, header_encoding="utf-8") + conn = h2.connection.H2Connection(config=config) + conn.initiate_connection() + conn.update_settings( + { + h2.settings.SettingCodes.MAX_CONCURRENT_STREAMS: self.max_concurrent_streams, + } + ) + writer.write(conn.data_to_send()) + await writer.drain() + + streams: dict[int, dict[str, Any]] = {} + write_lock = asyncio.Lock() + + async def flush_withheld_credit() -> None: + """When the test releases the stall, return WINDOW_UPDATE so the upload can finish.""" + await self.release_upload_credit.wait() + for _ in range(200): + async with write_lock: + for sid, st in list(streams.items()): + owed = int(st.get("flow_controlled") or 0) + if owed: + conn.acknowledge_received_data(owed, sid) + st["flow_controlled"] = 0 + to_send = conn.data_to_send() + if to_send: + writer.write(to_send) + await writer.drain() + if self.active_uploads == 0 and not any(int(st.get("flow_controlled") or 0) for st in streams.values()): + return + await asyncio.sleep(0.05) + + credit_task = asyncio.create_task(flush_withheld_credit()) + + async def respond(stream_id: int, st: dict[str, Any]) -> None: + path = st["path"] + method = st["method"] + started = st["started"] + hold_s = 0.0 + if path.endswith("/wait_for_status"): + hold_s = float(st.get("hold_s") or 2.0) + if hold_s > 0: + await asyncio.sleep(hold_s) + + if st.get("flow_controlled"): + async with write_lock: + conn.acknowledge_received_data(st["flow_controlled"], stream_id) + writer.write(conn.data_to_send()) + await writer.drain() + st["flow_controlled"] = 0 + + finished = time.perf_counter() + self.observations.append( + { + "conn_id": conn_id, + "path": path, + "method": method, + "started": started, + "finished": finished, + "duration_s": finished - started, + "body_bytes": len(st.get("body") or b""), + } + ) + if path.endswith("/wait_for_status"): + self.active_waits = max(0, self.active_waits - 1) + if path.endswith("/upload_file"): + self.active_uploads = max(0, self.active_uploads - 1) + + if path == "/v1/devboxes" and method == "POST": + payload = {"id": "dbx_test", "status": "running"} + else: + payload = {"ok": True} + raw = json.dumps(payload).encode() + headers = [ + (":status", "200"), + ("content-type", "application/json"), + ("content-length", str(len(raw))), + ] + async with write_lock: + conn.send_headers(stream_id, headers) + conn.send_data(stream_id, raw, end_stream=True) + writer.write(conn.data_to_send()) + await writer.drain() + streams.pop(stream_id, None) + + try: + while True: + data = await reader.read(65535) + if not data: + break + try: + events = conn.receive_data(data) + except h2.exceptions.ProtocolError: + break + for event in events: + if isinstance(event, h2.events.RequestReceived): + st = streams.setdefault( + event.stream_id, + { + "headers": [], + "body": bytearray(), + "method": "GET", + "path": "/", + "started": time.perf_counter(), + "flow_controlled": 0, + "hold_s": 2.0, + }, + ) + st["headers"] = list(event.headers) + for k, v in event.headers: + if k == ":method": + st["method"] = v + elif k == ":path": + st["path"] = v.split("?", 1)[0] + elif k.lower() == "x-hold-seconds": + try: + st["hold_s"] = float(v) + except ValueError: + pass + if str(st["path"]).endswith("/wait_for_status"): + self.active_waits += 1 + if str(st["path"]).endswith("/upload_file"): + self.active_uploads += 1 + elif isinstance(event, h2.events.DataReceived): + st = streams.setdefault(event.stream_id, {"body": bytearray(), "flow_controlled": 0}) + st["body"].extend(event.data) + stall_upload = self.upload_stall_s > 0 and str(st.get("path", "")).endswith("/upload_file") + if stall_upload and not self.release_upload_credit.is_set(): + st["flow_controlled"] = st.get("flow_controlled", 0) + event.flow_controlled_length + else: + conn.acknowledge_received_data(event.flow_controlled_length, event.stream_id) + elif isinstance(event, h2.events.StreamEnded): + st = streams.get(event.stream_id) + if st is not None: + asyncio.create_task(respond(event.stream_id, st)) + elif isinstance(event, h2.events.StreamReset): + streams.pop(event.stream_id, None) + async with write_lock: + to_send = conn.data_to_send() + if to_send: + writer.write(to_send) + await writer.drain() + except (ConnectionResetError, BrokenPipeError, asyncio.CancelledError): + pass + finally: + credit_task.cancel() + try: + writer.close() + await writer.wait_closed() + except Exception: + pass + + +async def _wait_until(predicate: Any, *, timeout_s: float = 5.0) -> None: + deadline = time.perf_counter() + timeout_s + while time.perf_counter() < deadline: + if predicate(): + return + await asyncio.sleep(0.02) + raise AssertionError("condition not met before timeout") + + +@pytest.mark.asyncio +async def test_waits_and_create_use_different_h2_connections(monkeypatch: pytest.MonkeyPatch) -> None: + # Force trust of the ephemeral self-signed cert used by the local h2 server. + orig_transport_init = httpx.AsyncHTTPTransport.__init__ + + def _transport_init(self: Any, *args: Any, **kwargs: Any) -> None: + kwargs["verify"] = False + orig_transport_init(self, *args, **kwargs) + + monkeypatch.setattr(httpx.AsyncHTTPTransport, "__init__", _transport_init) + + server = _H2BulkheadServer(max_concurrent_streams=2) + base = await server.start() + try: + client = AsyncRunloop( + bearer_token="test-token", + base_url=base, + shared_http_pool=False, + background_pool_shards=1, + transfer_pool_shards=1, + max_retries=0, + ) + try: + wait_headers = {"X-Hold-Seconds": "2.5"} + + async def wait_call() -> object: + return await client.post( + "/v1/devboxes/dbx_shared/wait_for_status", + body={"statuses": ["running"], "timeout_seconds": 5}, + options={"extra_headers": wait_headers}, + cast_to=object, + ) + + wait_tasks = [asyncio.create_task(wait_call()) for _ in range(2)] + await _wait_until(lambda: server.active_waits >= 2, timeout_s=5.0) + + t0 = time.perf_counter() + create_result = await client.post( + "/v1/devboxes", + body={"name": "bulkhead-create"}, + cast_to=object, + ) + create_s = time.perf_counter() - t0 + + assert isinstance(create_result, dict) + assert create_result["id"] == "dbx_test" + assert create_s < 0.75, f"create blocked by waits: {create_s:.3f}s" + + await asyncio.gather(*wait_tasks) + + wait_obs = [o for o in server.observations if o["path"].endswith("/wait_for_status")] + create_obs = [o for o in server.observations if o["path"] == "/v1/devboxes"] + assert len(wait_obs) == 2 + assert len(create_obs) == 1 + wait_conns = {o["conn_id"] for o in wait_obs} + create_conn = create_obs[0]["conn_id"] + assert create_conn not in wait_conns, f"create conn {create_conn} overlapped wait conns {wait_conns}" + finally: + await client.close() + finally: + await server.stop() + + +@pytest.mark.asyncio +async def test_upload_stall_does_not_block_create(monkeypatch: pytest.MonkeyPatch) -> None: + orig_transport_init = httpx.AsyncHTTPTransport.__init__ + + def _transport_init(self: Any, *args: Any, **kwargs: Any) -> None: + kwargs["verify"] = False + orig_transport_init(self, *args, **kwargs) + + monkeypatch.setattr(httpx.AsyncHTTPTransport, "__init__", _transport_init) + + server = _H2BulkheadServer(max_concurrent_streams=100, upload_stall_s=2.5) + base = await server.start() + try: + client = AsyncRunloop( + bearer_token="test-token", + base_url=base, + shared_http_pool=False, + background_pool_shards=1, + transfer_pool_shards=1, + max_retries=0, + ) + try: + payload = b"x" * (256 * 1024) + + async def upload() -> object: + return await client.post( + "/v1/devboxes/dbx_up/upload_file", + content=payload, + cast_to=object, + options={"extra_headers": {"content-type": "application/octet-stream"}}, + ) + + upload_task = asyncio.create_task(upload()) + await _wait_until(lambda: server.active_uploads >= 1, timeout_s=5.0) + + t0 = time.perf_counter() + create_result = await client.post( + "/v1/devboxes", + body={"name": "during-upload"}, + cast_to=object, + ) + create_s = time.perf_counter() - t0 + + assert isinstance(create_result, dict) + assert create_result["id"] == "dbx_test" + assert create_s < 0.75, f"create blocked by upload stall: {create_s:.3f}s" + + # Unblock the stalled upload body so the test can finish cleanly. + server.release_upload_credit.set() + # Nudge the connection by allowing the client to finish sending. + await asyncio.wait_for(upload_task, timeout=10.0) + + upload_obs = [o for o in server.observations if o["path"].endswith("/upload_file")] + create_obs = [o for o in server.observations if o["path"] == "/v1/devboxes"] + assert len(upload_obs) == 1 + assert len(create_obs) == 1 + assert create_obs[0]["conn_id"] != upload_obs[0]["conn_id"] + assert create_obs[0]["finished"] < upload_obs[0]["finished"] + finally: + await client.close() + finally: + await server.stop() diff --git a/tests/test_transfer_client.py b/tests/test_transfer_client.py index 8a4faad5a..4a57c3bb3 100644 --- a/tests/test_transfer_client.py +++ b/tests/test_transfer_client.py @@ -3,7 +3,10 @@ from __future__ import annotations import os +import re +import threading from typing import Any, Iterator +from pathlib import Path import httpx import pytest @@ -20,6 +23,18 @@ base_url = os.environ.get("TEST_API_BASE_URL", "http://127.0.0.1:4010") bearer_token = "My Bearer Token" +# Invariant: every long-lived / bulk-body RPC in the resources tree must match +# these suffixes or it silently shares the API H2 pool. Update both the +# classification helpers and this list when adding endpoints. +_KNOWN_BACKGROUND_PATHS = ( + "/v1/devboxes/dbx_1/wait_for_status", + "/v1/devboxes/dbx_1/executions/ex_1/wait_for_status", +) +_KNOWN_TRANSFER_PATHS = ( + "/v1/devboxes/dbx_1/upload_file", + "/v1/devboxes/dbx_1/download_file", +) + @pytest.fixture(autouse=True) def _reset_shared_pool() -> Iterator[None]: # pyright: ignore[reportUnusedFunction] @@ -163,3 +178,60 @@ async def test_async_bulkheads() -> None: assert wait is not upload finally: await client.close() + + +def test_concurrent_background_client_init_is_singleton() -> None: + client = _make_client(shared_http_pool=False, background_pool_shards=1) + barrier = threading.Barrier(16) + results: list[httpx.Client] = [] + errors: list[BaseException] = [] + + def worker() -> None: + try: + barrier.wait() + results.append(client._ensure_background_client(0)) + except BaseException as exc: # noqa: BLE001 + errors.append(exc) + + threads = [threading.Thread(target=worker) for _ in range(16)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert not errors + assert len(results) == 16 + assert len({id(c) for c in results}) == 1 + client.close() + + +def test_known_long_lived_paths_are_classified() -> None: + for path in _KNOWN_BACKGROUND_PATHS: + assert _is_background_path(path), path + assert not _is_transfer_path(path), path + for path in _KNOWN_TRANSFER_PATHS: + assert _is_transfer_path(path), path + assert not _is_background_path(path), path + + +def test_resources_tree_long_lived_ops_match_classifier() -> None: + """Fail if a new wait_for_status / upload_file / download_file path is added + without being covered by the bulkhead suffixes. + """ + root = Path(__file__).resolve().parents[1] / "src" / "runloop_api_client" / "resources" + text = "\n".join(p.read_text() for p in root.rglob("*.py")) + # Paths appear as f-strings or path_template("/v1/.../wait_for_status", ...) + found_wait = set(re.findall(r'(/v1/[^"\s]*wait_for_status)', text)) + found_upload = set(re.findall(r'(/v1/[^"\s]*upload_file)', text)) + found_download = set(re.findall(r'(/v1/[^"\s]*download_file)', text)) + + assert found_wait, "expected wait_for_status paths in resources/" + assert found_upload, "expected upload_file paths in resources/" + assert found_download, "expected download_file paths in resources/" + + for path in found_wait: + concrete = path.replace("{id}", "dbx_1").replace("{devbox_id}", "dbx_1").replace("{execution_id}", "ex_1") + assert _is_background_path(concrete), path + for path in found_upload | found_download: + concrete = path.replace("{id}", "dbx_1") + assert _is_transfer_path(concrete), path From 8bdcd53e8a7a24e312222ecb66af19c83772e1d8 Mon Sep 17 00:00:00 2001 From: Reflex Date: Mon, 27 Jul 2026 23:30:37 +0000 Subject: [PATCH 06/11] chore: fix pyright on H2 bulkhead integration test Co-authored-by: Cursor --- tests/test_h2_bulkhead_integration.py | 24 +++++++++++++++--------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/tests/test_h2_bulkhead_integration.py b/tests/test_h2_bulkhead_integration.py index c7f3896c7..2af683e97 100644 --- a/tests/test_h2_bulkhead_integration.py +++ b/tests/test_h2_bulkhead_integration.py @@ -15,7 +15,7 @@ import tempfile import threading import subprocess -from typing import Any, Iterator +from typing import Any, Iterator, cast from pathlib import Path import httpx @@ -26,9 +26,11 @@ import h2.events # noqa: E402 import h2.settings # noqa: E402 import h2.connection # noqa: E402 +import h2.exceptions # noqa: E402 import runloop_api_client._base_client as _base_mod from runloop_api_client import AsyncRunloop +from runloop_api_client._base_client import make_request_options pytestmark = pytest.mark.timeout(30) @@ -53,7 +55,7 @@ def _clear_pool_state() -> None: @pytest.fixture(autouse=True) -def _reset_pools() -> Iterator[None]: +def _reset_pools() -> Iterator[None]: # pyright: ignore[reportUnusedFunction] _clear_pool_state() yield _clear_pool_state() @@ -221,8 +223,10 @@ async def respond(stream_id: int, st: dict[str, Any]) -> None: events = conn.receive_data(data) except h2.exceptions.ProtocolError: break - for event in events: - if isinstance(event, h2.events.RequestReceived): + for raw_event in events: + # h2's Event hierarchy is poorly typed under pyright strict. + event = cast(Any, raw_event) + if isinstance(raw_event, h2.events.RequestReceived): st = streams.setdefault( event.stream_id, { @@ -250,7 +254,7 @@ async def respond(stream_id: int, st: dict[str, Any]) -> None: self.active_waits += 1 if str(st["path"]).endswith("/upload_file"): self.active_uploads += 1 - elif isinstance(event, h2.events.DataReceived): + elif isinstance(raw_event, h2.events.DataReceived): st = streams.setdefault(event.stream_id, {"body": bytearray(), "flow_controlled": 0}) st["body"].extend(event.data) stall_upload = self.upload_stall_s > 0 and str(st.get("path", "")).endswith("/upload_file") @@ -258,11 +262,11 @@ async def respond(stream_id: int, st: dict[str, Any]) -> None: st["flow_controlled"] = st.get("flow_controlled", 0) + event.flow_controlled_length else: conn.acknowledge_received_data(event.flow_controlled_length, event.stream_id) - elif isinstance(event, h2.events.StreamEnded): + elif isinstance(raw_event, h2.events.StreamEnded): st = streams.get(event.stream_id) if st is not None: asyncio.create_task(respond(event.stream_id, st)) - elif isinstance(event, h2.events.StreamReset): + elif isinstance(raw_event, h2.events.StreamReset): streams.pop(event.stream_id, None) async with write_lock: to_send = conn.data_to_send() @@ -318,7 +322,7 @@ async def wait_call() -> object: return await client.post( "/v1/devboxes/dbx_shared/wait_for_status", body={"statuses": ["running"], "timeout_seconds": 5}, - options={"extra_headers": wait_headers}, + options=make_request_options(extra_headers=wait_headers), cast_to=object, ) @@ -381,7 +385,9 @@ async def upload() -> object: "/v1/devboxes/dbx_up/upload_file", content=payload, cast_to=object, - options={"extra_headers": {"content-type": "application/octet-stream"}}, + options=make_request_options( + extra_headers={"content-type": "application/octet-stream"}, + ), ) upload_task = asyncio.create_task(upload()) From b6451eea652393146ef967e1927b63a6d3d0541d Mon Sep 17 00:00:00 2001 From: Reflex Date: Tue, 28 Jul 2026 00:13:23 +0000 Subject: [PATCH 07/11] fix: round-robin H2 bulkhead shards per SDK client Replace CRC32 resource affinity with per-client counters so concurrent waits/uploads spread across shards instead of pinning a busy resource to one connection. Co-authored-by: Cursor --- src/runloop_api_client/_base_client.py | 70 ++++++++++++++------------ src/runloop_api_client/_client.py | 4 +- src/runloop_api_client/sdk/async_.py | 4 +- src/runloop_api_client/sdk/sync.py | 4 +- tests/test_transfer_client.py | 54 +++++++++----------- 5 files changed, 69 insertions(+), 67 deletions(-) diff --git a/src/runloop_api_client/_base_client.py b/src/runloop_api_client/_base_client.py index 25b73e816..33af8100b 100644 --- a/src/runloop_api_client/_base_client.py +++ b/src/runloop_api_client/_base_client.py @@ -4,7 +4,6 @@ import json import time import uuid -import zlib import email import asyncio import inspect @@ -179,7 +178,8 @@ async def aclose(self) -> None: # Sharded H2 bulkheads: long-polls and file transfers stay off the API control-plane # connection. Each shard index maps to its own shared transport (≈ one H2 connection). -# Removable once httpcore respects stream capacity when opening connections. +# Per-client round-robin spreads concurrent requests across shards. Removable once +# httpcore respects stream capacity when opening connections. _shared_sync_background_transports: dict[int, _SharedTransport] = {} _shared_sync_transfer_transports: dict[int, _SharedTransport] = {} _shared_async_background_transports: weakref.WeakKeyDictionary[ @@ -201,30 +201,6 @@ def _is_transfer_path(path: str) -> bool: return path.endswith(_TRANSFER_PATH_SUFFIXES) -def _pool_affinity_key(path: str) -> str: - """Pick a stable resource id from the URL for shard routing.""" - parts = [p for p in path.split("/") if p] - try: - if "executions" in parts: - idx = parts.index("executions") - if idx + 1 < len(parts): - return parts[idx + 1] - if "devboxes" in parts: - idx = parts.index("devboxes") - if idx + 1 < len(parts): - return parts[idx + 1] - except ValueError: - pass - return path - - -def _shard_index(key: str, shards: int) -> int: - if shards <= 1: - return 0 - # crc32 is stable across processes (unlike PYTHONHASHSEED-randomized hash()). - return zlib.crc32(key.encode("utf-8")) % shards - - def _acquire_shared_sync_transport(bucket: dict[int, _SharedTransport], shard: int) -> _SharedTransport: with _pool_lock: existing = bucket.get(shard) @@ -1019,6 +995,8 @@ class SyncAPIClient(BaseClient[httpx.Client, Stream[Any]]): _isolate_workload_pools: bool _background_pool_shards: int _transfer_pool_shards: int + _background_next: int + _transfer_next: int _closed: bool def __init__( @@ -1076,6 +1054,8 @@ def __init__( self._bulkhead_lock = threading.Lock() self._background_pool_shards = background_pool_shards self._transfer_pool_shards = transfer_pool_shards + self._background_next = 0 + self._transfer_next = 0 # Custom http_client owns the full transport stack; don't invent sibling pools. self._isolate_workload_pools = http_client is None @@ -1151,15 +1131,27 @@ def _ensure_transfer_client(self, shard: int) -> httpx.Client: self._transfer_clients[shard] = client return client + def _next_background_client(self) -> httpx.Client: + # Select under the lock; ensure afterward so _ensure_* can take the same lock. + with self._bulkhead_lock: + shard = self._background_next % self._background_pool_shards + self._background_next += 1 + return self._ensure_background_client(shard) + + def _next_transfer_client(self) -> httpx.Client: + with self._bulkhead_lock: + shard = self._transfer_next % self._transfer_pool_shards + self._transfer_next += 1 + return self._ensure_transfer_client(shard) + def _send_client_for_request(self, request: httpx.Request) -> httpx.Client: if not self._isolate_workload_pools: return self._client path = request.url.path - key = _pool_affinity_key(path) if _is_background_path(path): - return self._ensure_background_client(_shard_index(key, self._background_pool_shards)) + return self._next_background_client() if _is_transfer_path(path): - return self._ensure_transfer_client(_shard_index(key, self._transfer_pool_shards)) + return self._next_transfer_client() return self._client def is_closed(self) -> bool: @@ -1736,6 +1728,8 @@ class AsyncAPIClient(BaseClient[httpx.AsyncClient, AsyncStream[Any]]): _isolate_workload_pools: bool _background_pool_shards: int _transfer_pool_shards: int + _background_next: int + _transfer_next: int _closed: bool def __init__( @@ -1792,6 +1786,8 @@ def __init__( self._transfer_clients = {} self._background_pool_shards = background_pool_shards self._transfer_pool_shards = transfer_pool_shards + self._background_next = 0 + self._transfer_next = 0 # Custom http_client owns the full transport stack; don't invent sibling pools. self._isolate_workload_pools = http_client is None @@ -1877,15 +1873,25 @@ def _ensure_transfer_client(self, shard: int) -> httpx.AsyncClient: self._transfer_clients[shard] = client return client + def _next_background_client(self) -> httpx.AsyncClient: + # Single-threaded event loop: counter bump needs no lock when there is no await. + shard = self._background_next % self._background_pool_shards + self._background_next += 1 + return self._ensure_background_client(shard) + + def _next_transfer_client(self) -> httpx.AsyncClient: + shard = self._transfer_next % self._transfer_pool_shards + self._transfer_next += 1 + return self._ensure_transfer_client(shard) + def _send_client_for_request(self, request: httpx.Request) -> httpx.AsyncClient: if not self._isolate_workload_pools: return self._client path = request.url.path - key = _pool_affinity_key(path) if _is_background_path(path): - return self._ensure_background_client(_shard_index(key, self._background_pool_shards)) + return self._next_background_client() if _is_transfer_path(path): - return self._ensure_transfer_client(_shard_index(key, self._transfer_pool_shards)) + return self._next_transfer_client() return self._client def is_closed(self) -> bool: diff --git a/src/runloop_api_client/_client.py b/src/runloop_api_client/_client.py index 7d6fd924d..5b5e555b9 100644 --- a/src/runloop_api_client/_client.py +++ b/src/runloop_api_client/_client.py @@ -98,7 +98,7 @@ def __init__( # Set to False to create a private connection pool (old behavior). shared_http_pool: bool = True, # Separate H2 connection pools for long-polls (/wait_for_status) and file - # transfers. Each shard ≈ one H2 connection, selected by hash(resource_id). + # transfers. Each shard ≈ one H2 connection; requests round-robin per client. background_pool_shards: int = DEFAULT_BACKGROUND_POOL_SHARDS, transfer_pool_shards: int = DEFAULT_TRANSFER_POOL_SHARDS, # Enable or disable schema validation for data returned by the API. @@ -406,7 +406,7 @@ def __init__( # Set to False to create a private connection pool (old behavior). shared_http_pool: bool = True, # Separate H2 connection pools for long-polls (/wait_for_status) and file - # transfers. Each shard ≈ one H2 connection, selected by hash(resource_id). + # transfers. Each shard ≈ one H2 connection; requests round-robin per client. background_pool_shards: int = DEFAULT_BACKGROUND_POOL_SHARDS, transfer_pool_shards: int = DEFAULT_TRANSFER_POOL_SHARDS, # Enable or disable schema validation for data returned by the API. diff --git a/src/runloop_api_client/sdk/async_.py b/src/runloop_api_client/sdk/async_.py index 2a28b86e4..8639771a4 100644 --- a/src/runloop_api_client/sdk/async_.py +++ b/src/runloop_api_client/sdk/async_.py @@ -1349,9 +1349,9 @@ def __init__( :type default_query: Mapping[str, object] | None, optional :param http_client: Custom ``httpx.AsyncClient`` instance to reuse, defaults to None :type http_client: httpx.AsyncClient | None, optional - :param background_pool_shards: H2 shards for long-polls, defaults to 2 + :param background_pool_shards: H2 shards for long-polls (round-robin), defaults to 2 :type background_pool_shards: int, optional - :param transfer_pool_shards: H2 shards for upload/download, defaults to 2 + :param transfer_pool_shards: H2 shards for upload/download (round-robin), defaults to 2 :type transfer_pool_shards: int, optional """ self.api = AsyncRunloop( diff --git a/src/runloop_api_client/sdk/sync.py b/src/runloop_api_client/sdk/sync.py index 0f23b105f..29a91303a 100644 --- a/src/runloop_api_client/sdk/sync.py +++ b/src/runloop_api_client/sdk/sync.py @@ -1374,9 +1374,9 @@ def __init__( :type default_query: Mapping[str, object] | None, optional :param http_client: Custom ``httpx.Client`` instance to reuse, defaults to None :type http_client: httpx.Client | None, optional - :param background_pool_shards: H2 shards for long-polls, defaults to 2 + :param background_pool_shards: H2 shards for long-polls (round-robin), defaults to 2 :type background_pool_shards: int, optional - :param transfer_pool_shards: H2 shards for upload/download, defaults to 2 + :param transfer_pool_shards: H2 shards for upload/download (round-robin), defaults to 2 :type transfer_pool_shards: int, optional """ self.api = Runloop( diff --git a/tests/test_transfer_client.py b/tests/test_transfer_client.py index 4a57c3bb3..d91752004 100644 --- a/tests/test_transfer_client.py +++ b/tests/test_transfer_client.py @@ -14,9 +14,7 @@ import runloop_api_client._base_client as _base_mod from runloop_api_client import Runloop, AsyncRunloop from runloop_api_client._base_client import ( - _shard_index, _is_transfer_path, - _pool_affinity_key, _is_background_path, ) @@ -77,15 +75,6 @@ def test_path_classification() -> None: assert _is_transfer_path("/v1/objects/obj_1/download") is False -def test_affinity_key_and_shard() -> None: - assert _pool_affinity_key("/v1/devboxes/dbx_1/upload_file") == "dbx_1" - assert _pool_affinity_key("/v1/devboxes/dbx_1/wait_for_status") == "dbx_1" - assert _pool_affinity_key("/v1/devboxes/dbx_1/executions/ex_9/wait_for_status") == "ex_9" - assert _shard_index("dbx_1", 1) == 0 - assert _shard_index("dbx_1", 2) in (0, 1) - assert _shard_index("dbx_1", 2) == _shard_index("dbx_1", 2) - - def test_api_background_transfer_use_distinct_transports() -> None: client = _make_client(shared_http_pool=True, background_pool_shards=2, transfer_pool_shards=2) try: @@ -108,26 +97,24 @@ def test_api_background_transfer_use_distinct_transports() -> None: client.close() -def test_same_resource_affinity_uses_same_shard() -> None: +def test_round_robin_spreads_requests_across_shards() -> None: client = _make_client(background_pool_shards=2, transfer_pool_shards=2) try: - w1 = client._send_client_for_request(httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_same/wait_for_status")) - w2 = client._send_client_for_request(httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_same/wait_for_status")) - assert w1 is w2 - finally: - client.close() - - -def test_different_resources_can_land_on_different_shards() -> None: - client = _make_client(background_pool_shards=2) - try: - seen: set[int] = set() - for i in range(40): - req = httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_{i}/wait_for_status") - c = client._send_client_for_request(req) - seen.add(id(c._transport)) # type: ignore[attr-defined] - assert len(seen) == 2 + wait_req = httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_same/wait_for_status") + w0 = client._send_client_for_request(wait_req) + w1 = client._send_client_for_request(wait_req) + w2 = client._send_client_for_request(wait_req) + assert w0 is not w1 + assert w0 is w2 + assert set(client._background_clients) == {0, 1} + + upload_req = httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_same/upload_file") + t0 = client._send_client_for_request(upload_req) + t1 = client._send_client_for_request(upload_req) + assert t0 is not t1 + assert set(client._transfer_clients) == {0, 1} assert len(_base_mod._shared_sync_background_transports) == 2 + assert len(_base_mod._shared_sync_transfer_transports) == 2 finally: client.close() @@ -145,15 +132,24 @@ def test_custom_http_client_skips_isolation() -> None: custom.close() -def test_shards_shared_across_sdk_clients() -> None: +def test_round_robin_is_per_client_while_transports_are_shared() -> None: c1 = _make_client(background_pool_shards=2) c2 = _make_client(background_pool_shards=2) try: req = httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_shared/wait_for_status") + # Each SDK client has its own counter; both start at shard 0. t1 = c1._send_client_for_request(req) t2 = c2._send_client_for_request(req) assert t1 is not t2 assert t1._transport is t2._transport # type: ignore[attr-defined] + assert c1._background_next == 1 + assert c2._background_next == 1 + + # Advancing one client does not affect the other. + t1b = c1._send_client_for_request(req) + assert t1b is not t1 + assert c1._background_next == 2 + assert c2._background_next == 1 finally: c1.close() c2.close() From 5900470dd83314d080ae7d5b0387f8f9a197ca42 Mon Sep 17 00:00:00 2001 From: Reflex Date: Tue, 28 Jul 2026 00:16:47 +0000 Subject: [PATCH 08/11] fix: randomize per-client H2 shard start offsets Avoid every new SDK client landing first wait/upload on global shard 0 when transports are process-shared. Co-authored-by: Cursor --- src/runloop_api_client/_base_client.py | 15 +++++++------ src/runloop_api_client/_client.py | 6 ++++-- tests/test_transfer_client.py | 29 ++++++++++++++++++++------ 3 files changed, 36 insertions(+), 14 deletions(-) diff --git a/src/runloop_api_client/_base_client.py b/src/runloop_api_client/_base_client.py index 33af8100b..3e48b2340 100644 --- a/src/runloop_api_client/_base_client.py +++ b/src/runloop_api_client/_base_client.py @@ -8,6 +8,7 @@ import asyncio import inspect import logging +import secrets import weakref import platform import warnings @@ -178,8 +179,8 @@ async def aclose(self) -> None: # Sharded H2 bulkheads: long-polls and file transfers stay off the API control-plane # connection. Each shard index maps to its own shared transport (≈ one H2 connection). -# Per-client round-robin spreads concurrent requests across shards. Removable once -# httpcore respects stream capacity when opening connections. +# Per-client round-robin (random start offset) spreads concurrent requests across +# shards. Removable once httpcore respects stream capacity when opening connections. _shared_sync_background_transports: dict[int, _SharedTransport] = {} _shared_sync_transfer_transports: dict[int, _SharedTransport] = {} _shared_async_background_transports: weakref.WeakKeyDictionary[ @@ -1054,8 +1055,9 @@ def __init__( self._bulkhead_lock = threading.Lock() self._background_pool_shards = background_pool_shards self._transfer_pool_shards = transfer_pool_shards - self._background_next = 0 - self._transfer_next = 0 + # Random start avoids every short-lived SDK client pinning global shard 0. + self._background_next = secrets.randbelow(background_pool_shards) + self._transfer_next = secrets.randbelow(transfer_pool_shards) # Custom http_client owns the full transport stack; don't invent sibling pools. self._isolate_workload_pools = http_client is None @@ -1786,8 +1788,9 @@ def __init__( self._transfer_clients = {} self._background_pool_shards = background_pool_shards self._transfer_pool_shards = transfer_pool_shards - self._background_next = 0 - self._transfer_next = 0 + # Random start avoids every short-lived SDK client pinning global shard 0. + self._background_next = secrets.randbelow(background_pool_shards) + self._transfer_next = secrets.randbelow(transfer_pool_shards) # Custom http_client owns the full transport stack; don't invent sibling pools. self._isolate_workload_pools = http_client is None diff --git a/src/runloop_api_client/_client.py b/src/runloop_api_client/_client.py index 5b5e555b9..36599b314 100644 --- a/src/runloop_api_client/_client.py +++ b/src/runloop_api_client/_client.py @@ -98,7 +98,8 @@ def __init__( # Set to False to create a private connection pool (old behavior). shared_http_pool: bool = True, # Separate H2 connection pools for long-polls (/wait_for_status) and file - # transfers. Each shard ≈ one H2 connection; requests round-robin per client. + # transfers. Each shard ≈ one H2 connection; requests round-robin per client + # from a randomized starting offset (avoids every new client hitting shard 0). background_pool_shards: int = DEFAULT_BACKGROUND_POOL_SHARDS, transfer_pool_shards: int = DEFAULT_TRANSFER_POOL_SHARDS, # Enable or disable schema validation for data returned by the API. @@ -406,7 +407,8 @@ def __init__( # Set to False to create a private connection pool (old behavior). shared_http_pool: bool = True, # Separate H2 connection pools for long-polls (/wait_for_status) and file - # transfers. Each shard ≈ one H2 connection; requests round-robin per client. + # transfers. Each shard ≈ one H2 connection; requests round-robin per client + # from a randomized starting offset (avoids every new client hitting shard 0). background_pool_shards: int = DEFAULT_BACKGROUND_POOL_SHARDS, transfer_pool_shards: int = DEFAULT_TRANSFER_POOL_SHARDS, # Enable or disable schema validation for data returned by the API. diff --git a/tests/test_transfer_client.py b/tests/test_transfer_client.py index d91752004..830993a19 100644 --- a/tests/test_transfer_client.py +++ b/tests/test_transfer_client.py @@ -119,6 +119,18 @@ def test_round_robin_spreads_requests_across_shards() -> None: client.close() +def test_many_clients_first_requests_use_both_shards() -> None: + clients = [_make_client(background_pool_shards=2) for _ in range(40)] + try: + req = httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_1/wait_for_status") + seen_transports = {id(c._send_client_for_request(req)._transport) for c in clients} # type: ignore[attr-defined] + assert len(seen_transports) == 2 + assert set(_base_mod._shared_sync_background_transports) == {0, 1} + finally: + for c in clients: + c.close() + + def test_custom_http_client_skips_isolation() -> None: custom = httpx.Client() client = _make_client(http_client=custom) @@ -137,19 +149,24 @@ def test_round_robin_is_per_client_while_transports_are_shared() -> None: c2 = _make_client(background_pool_shards=2) try: req = httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_shared/wait_for_status") - # Each SDK client has its own counter; both start at shard 0. + start1 = c1._background_next + start2 = c2._background_next t1 = c1._send_client_for_request(req) t2 = c2._send_client_for_request(req) assert t1 is not t2 - assert t1._transport is t2._transport # type: ignore[attr-defined] - assert c1._background_next == 1 - assert c2._background_next == 1 + # Same starting shard → same shared transport; different → different shards. + if start1 % 2 == start2 % 2: + assert t1._transport is t2._transport # type: ignore[attr-defined] + else: + assert t1._transport is not t2._transport # type: ignore[attr-defined] + assert c1._background_next == start1 + 1 + assert c2._background_next == start2 + 1 # Advancing one client does not affect the other. t1b = c1._send_client_for_request(req) assert t1b is not t1 - assert c1._background_next == 2 - assert c2._background_next == 1 + assert c1._background_next == start1 + 2 + assert c2._background_next == start2 + 1 finally: c1.close() c2.close() From e6ef85aa8e190d05757773f1b4c5bcb9b6bff84a Mon Sep 17 00:00:00 2001 From: Reflex Date: Tue, 28 Jul 2026 21:43:17 +0000 Subject: [PATCH 09/11] refactor: encapsulate sharded H2 pools; raise defaults Move shard registries onto SharedTransport, route via get_client_for_path, and shard the API pool too (defaults: API 8 / background 16 / transfer 2). Co-authored-by: Cursor --- src/runloop_api_client/_base_client.py | 532 +++++++++++++------------ src/runloop_api_client/_client.py | 20 +- src/runloop_api_client/_constants.py | 9 +- src/runloop_api_client/sdk/async_.py | 8 +- src/runloop_api_client/sdk/sync.py | 8 +- tests/test_h2_bulkhead_integration.py | 29 +- tests/test_shared_pool.py | 29 +- tests/test_transfer_client.py | 99 +++-- 8 files changed, 388 insertions(+), 346 deletions(-) diff --git a/src/runloop_api_client/_base_client.py b/src/runloop_api_client/_base_client.py index 3e48b2340..fbef8e7aa 100644 --- a/src/runloop_api_client/_base_client.py +++ b/src/runloop_api_client/_base_client.py @@ -25,6 +25,7 @@ Generic, Mapping, TypeVar, + Callable, Iterable, Iterator, Optional, @@ -79,6 +80,7 @@ DEFAULT_MAX_RETRIES, INITIAL_RETRY_DELAY, RAW_RESPONSE_HEADER, + DEFAULT_API_POOL_SHARDS, OVERRIDE_CAST_TO_HEADER, DEFAULT_CONNECTION_LIMITS, DEFAULT_TRANSFER_POOL_SHARDS, @@ -101,6 +103,10 @@ # when the last user releases it. # The async transport is keyed by event loop because connections bind to the # loop that created them and cannot be reused across asyncio.run() calls. +# +# httpx/httpcore use a single H2 connection per transport, so each workload +# category uses a sharded Registry of SharedTransport instances. Per-client +# round-robin spreads requests across shards. _pool_lock = threading.Lock() @@ -137,6 +143,34 @@ def close(self) -> None: if should_close: self._transport.close() + class Registry: + """Process-global map of shard index → shared H2 transport.""" + + def __init__(self) -> None: + self._shards: dict[int, _SharedTransport] = {} + + def acquire(self, shard: int) -> _SharedTransport: + with _pool_lock: + existing = self._shards.get(shard) + if existing is not None and existing.acquire(): + return existing + transport = _SharedTransport( + httpx.HTTPTransport(limits=DEFAULT_CONNECTION_LIMITS, http2=True), + ) + self._shards[shard] = transport + return transport + + def take_all(self) -> list[_SharedTransport]: + """Remove and return all transports (test cleanup).""" + with _pool_lock: + values = list(self._shards.values()) + self._shards.clear() + return values + + def shard_ids(self) -> set[int]: + with _pool_lock: + return set(self._shards) + class _SharedAsyncTransport(httpx.AsyncBaseTransport): """Async refcounted wrapper: delegates to a real async transport.""" @@ -171,24 +205,41 @@ async def aclose(self) -> None: if should_close: await self._transport.aclose() + class Registry: + """Per-event-loop map of shard index → shared async H2 transport.""" -_shared_sync_transport: _SharedTransport | None = None -_shared_async_transports: weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, _SharedAsyncTransport] = ( - weakref.WeakKeyDictionary() -) + def __init__(self) -> None: + self._by_loop: weakref.WeakKeyDictionary[ + asyncio.AbstractEventLoop, dict[int, _SharedAsyncTransport] + ] = weakref.WeakKeyDictionary() + + def acquire(self, loop: asyncio.AbstractEventLoop, shard: int) -> _SharedAsyncTransport: + with _pool_lock: + bucket = self._by_loop.get(loop) + if bucket is None: + bucket = {} + self._by_loop[loop] = bucket + existing = bucket.get(shard) + if existing is not None and existing.acquire(): + return existing + transport = _SharedAsyncTransport( + httpx.AsyncHTTPTransport(limits=DEFAULT_CONNECTION_LIMITS, http2=True), + ) + bucket[shard] = transport + return transport + + def clear(self) -> None: + with _pool_lock: + self._by_loop.clear() -# Sharded H2 bulkheads: long-polls and file transfers stay off the API control-plane -# connection. Each shard index maps to its own shared transport (≈ one H2 connection). -# Per-client round-robin (random start offset) spreads concurrent requests across -# shards. Removable once httpcore respects stream capacity when opening connections. -_shared_sync_background_transports: dict[int, _SharedTransport] = {} -_shared_sync_transfer_transports: dict[int, _SharedTransport] = {} -_shared_async_background_transports: weakref.WeakKeyDictionary[ - asyncio.AbstractEventLoop, dict[int, _SharedAsyncTransport] -] = weakref.WeakKeyDictionary() -_shared_async_transfer_transports: weakref.WeakKeyDictionary[ - asyncio.AbstractEventLoop, dict[int, _SharedAsyncTransport] -] = weakref.WeakKeyDictionary() + +# Process-global sharded registries (one ≈ H2 connection per shard index). +_shared_sync_api_transports = _SharedTransport.Registry() +_shared_sync_background_transports = _SharedTransport.Registry() +_shared_sync_transfer_transports = _SharedTransport.Registry() +_shared_async_api_transports = _SharedAsyncTransport.Registry() +_shared_async_background_transports = _SharedAsyncTransport.Registry() +_shared_async_transfer_transports = _SharedAsyncTransport.Registry() _BACKGROUND_PATH_SUFFIXES = ("/wait_for_status",) _TRANSFER_PATH_SUFFIXES = ("/upload_file", "/download_file") @@ -202,36 +253,100 @@ def _is_transfer_path(path: str) -> bool: return path.endswith(_TRANSFER_PATH_SUFFIXES) -def _acquire_shared_sync_transport(bucket: dict[int, _SharedTransport], shard: int) -> _SharedTransport: - with _pool_lock: - existing = bucket.get(shard) - if existing is not None and existing.acquire(): +class _SyncClientPool: + """Per-SDK-client round-robin over sharded httpx.Client wrappers.""" + + def __init__( + self, + *, + shards: int, + shared: bool, + registry: _SharedTransport.Registry | None, + make_client: Callable[[httpx.BaseTransport | None], httpx.Client], + ) -> None: + if shards < 1: + raise ValueError("shards must be >= 1") + self.shards = shards + self._shared = shared + self._registry = registry + self._make_client = make_client + self._clients: dict[int, httpx.Client] = {} + self._lock = threading.Lock() + self._next = secrets.randbelow(shards) + + def ensure(self, shard: int) -> httpx.Client: + existing = self._clients.get(shard) + if existing is not None: return existing - transport = _SharedTransport( - httpx.HTTPTransport(limits=DEFAULT_CONNECTION_LIMITS, http2=True), - ) - bucket[shard] = transport - return transport - - -def _acquire_shared_async_transport( - by_loop: weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, dict[int, _SharedAsyncTransport]], - loop: asyncio.AbstractEventLoop, - shard: int, -) -> _SharedAsyncTransport: - with _pool_lock: - bucket = by_loop.get(loop) - if bucket is None: - bucket = {} - by_loop[loop] = bucket - existing = bucket.get(shard) - if existing is not None and existing.acquire(): + with self._lock: + existing = self._clients.get(shard) + if existing is not None: + return existing + transport: httpx.BaseTransport | None = None + if self._shared and self._registry is not None: + transport = self._registry.acquire(shard) + client = self._make_client(transport) + self._clients[shard] = client + return client + + def next_client(self) -> httpx.Client: + with self._lock: + shard = self._next % self.shards + self._next += 1 + return self.ensure(shard) + + def close(self) -> None: + for client in self._clients.values(): + client.close() + self._clients.clear() + + +class _AsyncClientPool: + """Per-SDK-client round-robin over sharded httpx.AsyncClient wrappers.""" + + def __init__( + self, + *, + shards: int, + shared: bool, + registry: _SharedAsyncTransport.Registry | None, + make_client: Callable[[httpx.AsyncBaseTransport | None], httpx.AsyncClient], + ) -> None: + if shards < 1: + raise ValueError("shards must be >= 1") + self.shards = shards + self._shared = shared + self._registry = registry + self._make_client = make_client + self._clients: dict[int, httpx.AsyncClient] = {} + self._next = secrets.randbelow(shards) + + def ensure(self, shard: int) -> httpx.AsyncClient: + existing = self._clients.get(shard) + if existing is not None: return existing - transport = _SharedAsyncTransport( - httpx.AsyncHTTPTransport(limits=DEFAULT_CONNECTION_LIMITS, http2=True), - ) - bucket[shard] = transport - return transport + transport: httpx.AsyncBaseTransport | None = None + if self._shared and self._registry is not None: + try: + loop: asyncio.AbstractEventLoop | None = asyncio.get_running_loop() + except RuntimeError: + loop = None + if loop is not None: + transport = self._registry.acquire(loop, shard) + client = self._make_client(transport) + self._clients[shard] = client + return client + + def next_client(self) -> httpx.AsyncClient: + # Single-threaded event loop: counter bump needs no lock when there is no await. + shard = self._next % self.shards + self._next += 1 + return self.ensure(shard) + + async def aclose(self) -> None: + for client in self._clients.values(): + await client.aclose() + self._clients.clear() # TODO: make base page type vars covariant @@ -989,15 +1104,15 @@ def __del__(self) -> None: class SyncAPIClient(BaseClient[httpx.Client, Stream[Any]]): _client: httpx.Client - _background_clients: dict[int, httpx.Client] - _transfer_clients: dict[int, httpx.Client] + _api_pool: _SyncClientPool | None + _background_pool: _SyncClientPool | None + _transfer_pool: _SyncClientPool | None _default_stream_cls: type[Stream[Any]] | None = None _uses_shared_pool: bool _isolate_workload_pools: bool + _api_pool_shards: int _background_pool_shards: int _transfer_pool_shards: int - _background_next: int - _transfer_next: int _closed: bool def __init__( @@ -1012,6 +1127,7 @@ def __init__( custom_query: Mapping[str, object] | None = None, _strict_response_validation: bool, shared_http_pool: bool = True, + api_pool_shards: int = DEFAULT_API_POOL_SHARDS, background_pool_shards: int = DEFAULT_BACKGROUND_POOL_SHARDS, transfer_pool_shards: int = DEFAULT_TRANSFER_POOL_SHARDS, ) -> None: @@ -1033,6 +1149,8 @@ def __init__( f"Invalid `http_client` argument; Expected an instance of `httpx.Client` but got {type(http_client)}" ) + if api_pool_shards < 1: + raise ValueError("api_pool_shards must be >= 1") if background_pool_shards < 1: raise ValueError("background_pool_shards must be >= 1") if transfer_pool_shards < 1: @@ -1050,111 +1168,72 @@ def __init__( ) self._closed = False - self._background_clients = {} - self._transfer_clients = {} - self._bulkhead_lock = threading.Lock() + self._api_pool_shards = api_pool_shards self._background_pool_shards = background_pool_shards self._transfer_pool_shards = transfer_pool_shards - # Random start avoids every short-lived SDK client pinning global shard 0. - self._background_next = secrets.randbelow(background_pool_shards) - self._transfer_next = secrets.randbelow(transfer_pool_shards) # Custom http_client owns the full transport stack; don't invent sibling pools. self._isolate_workload_pools = http_client is None if http_client is not None: self._client = http_client self._uses_shared_pool = False - elif shared_http_pool: - global _shared_sync_transport - with _pool_lock: - if _shared_sync_transport is None or not _shared_sync_transport.acquire(): - _shared_sync_transport = _SharedTransport( - httpx.HTTPTransport(limits=DEFAULT_CONNECTION_LIMITS, http2=True), - ) - self._client = SyncHttpxClientWrapper( - base_url=base_url, - timeout=cast(Timeout, timeout), - transport=_shared_sync_transport, - ) - self._uses_shared_pool = True - else: - self._client = SyncHttpxClientWrapper( - base_url=base_url, - timeout=cast(Timeout, timeout), - ) - self._uses_shared_pool = False + self._api_pool = None + self._background_pool = None + self._transfer_pool = None + return - def _make_bulkhead_client(self, *, transport: httpx.BaseTransport | None) -> httpx.Client: - timeout = cast(Timeout, self.timeout) - if transport is not None: + self._uses_shared_pool = shared_http_pool + + def make_client(transport: httpx.BaseTransport | None) -> httpx.Client: + timeout_ = cast(Timeout, self.timeout) + if transport is not None: + return SyncHttpxClientWrapper( + base_url=self._base_url, + timeout=timeout_, + transport=transport, + ) return SyncHttpxClientWrapper( base_url=self._base_url, - timeout=timeout, - transport=transport, + timeout=timeout_, ) - return SyncHttpxClientWrapper( - base_url=self._base_url, - timeout=timeout, - ) - def _ensure_background_client(self, shard: int) -> httpx.Client: - existing = self._background_clients.get(shard) - if existing is not None: - return existing - with self._bulkhead_lock: - existing = self._background_clients.get(shard) - if existing is not None: - return existing - if self._uses_shared_pool: - transport: httpx.BaseTransport | None = _acquire_shared_sync_transport( - _shared_sync_background_transports, shard - ) - else: - transport = None - client = self._make_bulkhead_client(transport=transport) - self._background_clients[shard] = client - return client + api_registry = _shared_sync_api_transports if shared_http_pool else None + bg_registry = _shared_sync_background_transports if shared_http_pool else None + xfer_registry = _shared_sync_transfer_transports if shared_http_pool else None - def _ensure_transfer_client(self, shard: int) -> httpx.Client: - existing = self._transfer_clients.get(shard) - if existing is not None: - return existing - with self._bulkhead_lock: - existing = self._transfer_clients.get(shard) - if existing is not None: - return existing - if self._uses_shared_pool: - transport: httpx.BaseTransport | None = _acquire_shared_sync_transport( - _shared_sync_transfer_transports, shard - ) - else: - transport = None - client = self._make_bulkhead_client(transport=transport) - self._transfer_clients[shard] = client - return client - - def _next_background_client(self) -> httpx.Client: - # Select under the lock; ensure afterward so _ensure_* can take the same lock. - with self._bulkhead_lock: - shard = self._background_next % self._background_pool_shards - self._background_next += 1 - return self._ensure_background_client(shard) - - def _next_transfer_client(self) -> httpx.Client: - with self._bulkhead_lock: - shard = self._transfer_next % self._transfer_pool_shards - self._transfer_next += 1 - return self._ensure_transfer_client(shard) + self._api_pool = _SyncClientPool( + shards=api_pool_shards, + shared=shared_http_pool, + registry=api_registry, + make_client=make_client, + ) + self._background_pool = _SyncClientPool( + shards=background_pool_shards, + shared=shared_http_pool, + registry=bg_registry, + make_client=make_client, + ) + self._transfer_pool = _SyncClientPool( + shards=transfer_pool_shards, + shared=shared_http_pool, + registry=xfer_registry, + make_client=make_client, + ) + # Eager primary client (shard 0) for lifecycle / cookie-jar compatibility. + self._client = self._api_pool.ensure(0) - def _send_client_for_request(self, request: httpx.Request) -> httpx.Client: - if not self._isolate_workload_pools: + def _get_client_for_path(self, path: str) -> httpx.Client: + if not self._isolate_workload_pools or self._api_pool is None: return self._client - path = request.url.path + assert self._background_pool is not None and self._transfer_pool is not None if _is_background_path(path): - return self._next_background_client() + return self._background_pool.next_client() if _is_transfer_path(path): - return self._next_transfer_client() - return self._client + return self._transfer_pool.next_client() + return self._api_pool.next_client() + + def _send_client_for_request(self, request: httpx.Request) -> httpx.Client: + return self._get_client_for_path(request.url.path) def is_closed(self) -> bool: return self._closed or self._client.is_closed @@ -1169,13 +1248,14 @@ def close(self) -> None: if self._closed: return self._closed = True - self._client.close() - for client in self._background_clients.values(): - client.close() - self._background_clients.clear() - for client in self._transfer_clients.values(): - client.close() - self._transfer_clients.clear() + if self._api_pool is not None: + # Closes _client (api shard 0) along with any other API shards. + self._api_pool.close() + assert self._background_pool is not None and self._transfer_pool is not None + self._background_pool.close() + self._transfer_pool.close() + else: + self._client.close() def __enter__(self: _T) -> _T: return self @@ -1723,15 +1803,15 @@ def __del__(self) -> None: class AsyncAPIClient(BaseClient[httpx.AsyncClient, AsyncStream[Any]]): _client: httpx.AsyncClient - _background_clients: dict[int, httpx.AsyncClient] - _transfer_clients: dict[int, httpx.AsyncClient] + _api_pool: _AsyncClientPool | None + _background_pool: _AsyncClientPool | None + _transfer_pool: _AsyncClientPool | None _default_stream_cls: type[AsyncStream[Any]] | None = None _uses_shared_pool: bool _isolate_workload_pools: bool + _api_pool_shards: int _background_pool_shards: int _transfer_pool_shards: int - _background_next: int - _transfer_next: int _closed: bool def __init__( @@ -1746,6 +1826,7 @@ def __init__( custom_headers: Mapping[str, str] | None = None, custom_query: Mapping[str, object] | None = None, shared_http_pool: bool = True, + api_pool_shards: int = DEFAULT_API_POOL_SHARDS, background_pool_shards: int = DEFAULT_BACKGROUND_POOL_SHARDS, transfer_pool_shards: int = DEFAULT_TRANSFER_POOL_SHARDS, ) -> None: @@ -1767,6 +1848,8 @@ def __init__( f"Invalid `http_client` argument; Expected an instance of `httpx.AsyncClient` but got {type(http_client)}" ) + if api_pool_shards < 1: + raise ValueError("api_pool_shards must be >= 1") if background_pool_shards < 1: raise ValueError("background_pool_shards must be >= 1") if transfer_pool_shards < 1: @@ -1784,118 +1867,79 @@ def __init__( ) self._closed = False - self._background_clients = {} - self._transfer_clients = {} + self._api_pool_shards = api_pool_shards self._background_pool_shards = background_pool_shards self._transfer_pool_shards = transfer_pool_shards - # Random start avoids every short-lived SDK client pinning global shard 0. - self._background_next = secrets.randbelow(background_pool_shards) - self._transfer_next = secrets.randbelow(transfer_pool_shards) # Custom http_client owns the full transport stack; don't invent sibling pools. self._isolate_workload_pools = http_client is None if http_client is not None: self._client = http_client self._uses_shared_pool = False - elif shared_http_pool: + self._api_pool = None + self._background_pool = None + self._transfer_pool = None + return + + # Async shared transports require a running loop; without one, fall back to private. + can_share = shared_http_pool + if shared_http_pool: try: - loop: asyncio.AbstractEventLoop | None = asyncio.get_running_loop() + asyncio.get_running_loop() except RuntimeError: - loop = None - if loop is not None: - with _pool_lock: - existing = _shared_async_transports.get(loop) - if existing is not None and existing.acquire(): - transport: _SharedAsyncTransport = existing - else: - transport = _SharedAsyncTransport( - httpx.AsyncHTTPTransport(limits=DEFAULT_CONNECTION_LIMITS, http2=True), - ) - _shared_async_transports[loop] = transport - self._client = AsyncHttpxClientWrapper( - base_url=base_url, - timeout=cast(Timeout, timeout), + can_share = False + + self._uses_shared_pool = can_share + + def make_client(transport: httpx.AsyncBaseTransport | None) -> httpx.AsyncClient: + timeout_ = cast(Timeout, self.timeout) + if transport is not None: + return AsyncHttpxClientWrapper( + base_url=self._base_url, + timeout=timeout_, transport=transport, ) - self._uses_shared_pool = True - else: - self._client = AsyncHttpxClientWrapper( - base_url=base_url, - timeout=cast(Timeout, timeout), - ) - self._uses_shared_pool = False - else: - self._client = AsyncHttpxClientWrapper( - base_url=base_url, - timeout=cast(Timeout, timeout), - ) - self._uses_shared_pool = False - - def _make_bulkhead_client(self, *, transport: httpx.AsyncBaseTransport | None) -> httpx.AsyncClient: - timeout = cast(Timeout, self.timeout) - if transport is not None: return AsyncHttpxClientWrapper( base_url=self._base_url, - timeout=timeout, - transport=transport, + timeout=timeout_, ) - return AsyncHttpxClientWrapper( - base_url=self._base_url, - timeout=timeout, - ) - def _ensure_background_client(self, shard: int) -> httpx.AsyncClient: - existing = self._background_clients.get(shard) - if existing is not None: - return existing - transport: httpx.AsyncBaseTransport | None = None - if self._uses_shared_pool: - try: - loop: asyncio.AbstractEventLoop | None = asyncio.get_running_loop() - except RuntimeError: - loop = None - if loop is not None: - transport = _acquire_shared_async_transport(_shared_async_background_transports, loop, shard) - client = self._make_bulkhead_client(transport=transport) - self._background_clients[shard] = client - return client + api_registry = _shared_async_api_transports if can_share else None + bg_registry = _shared_async_background_transports if can_share else None + xfer_registry = _shared_async_transfer_transports if can_share else None - def _ensure_transfer_client(self, shard: int) -> httpx.AsyncClient: - existing = self._transfer_clients.get(shard) - if existing is not None: - return existing - transport: httpx.AsyncBaseTransport | None = None - if self._uses_shared_pool: - try: - loop: asyncio.AbstractEventLoop | None = asyncio.get_running_loop() - except RuntimeError: - loop = None - if loop is not None: - transport = _acquire_shared_async_transport(_shared_async_transfer_transports, loop, shard) - client = self._make_bulkhead_client(transport=transport) - self._transfer_clients[shard] = client - return client - - def _next_background_client(self) -> httpx.AsyncClient: - # Single-threaded event loop: counter bump needs no lock when there is no await. - shard = self._background_next % self._background_pool_shards - self._background_next += 1 - return self._ensure_background_client(shard) - - def _next_transfer_client(self) -> httpx.AsyncClient: - shard = self._transfer_next % self._transfer_pool_shards - self._transfer_next += 1 - return self._ensure_transfer_client(shard) + self._api_pool = _AsyncClientPool( + shards=api_pool_shards, + shared=can_share, + registry=api_registry, + make_client=make_client, + ) + self._background_pool = _AsyncClientPool( + shards=background_pool_shards, + shared=can_share, + registry=bg_registry, + make_client=make_client, + ) + self._transfer_pool = _AsyncClientPool( + shards=transfer_pool_shards, + shared=can_share, + registry=xfer_registry, + make_client=make_client, + ) + self._client = self._api_pool.ensure(0) - def _send_client_for_request(self, request: httpx.Request) -> httpx.AsyncClient: - if not self._isolate_workload_pools: + def _get_client_for_path(self, path: str) -> httpx.AsyncClient: + if not self._isolate_workload_pools or self._api_pool is None: return self._client - path = request.url.path + assert self._background_pool is not None and self._transfer_pool is not None if _is_background_path(path): - return self._next_background_client() + return self._background_pool.next_client() if _is_transfer_path(path): - return self._next_transfer_client() - return self._client + return self._transfer_pool.next_client() + return self._api_pool.next_client() + + def _send_client_for_request(self, request: httpx.Request) -> httpx.AsyncClient: + return self._get_client_for_path(request.url.path) def is_closed(self) -> bool: return self._closed or self._client.is_closed @@ -1910,13 +1954,13 @@ async def close(self) -> None: if self._closed: return self._closed = True - await self._client.aclose() - for client in self._background_clients.values(): - await client.aclose() - self._background_clients.clear() - for client in self._transfer_clients.values(): - await client.aclose() - self._transfer_clients.clear() + if self._api_pool is not None: + await self._api_pool.aclose() + assert self._background_pool is not None and self._transfer_pool is not None + await self._background_pool.aclose() + await self._transfer_pool.aclose() + else: + await self._client.aclose() async def __aenter__(self: _T) -> _T: return self diff --git a/src/runloop_api_client/_client.py b/src/runloop_api_client/_client.py index 36599b314..bf32e59d6 100644 --- a/src/runloop_api_client/_client.py +++ b/src/runloop_api_client/_client.py @@ -26,7 +26,7 @@ ) from ._compat import cached_property from ._version import __version__ -from ._constants import DEFAULT_TRANSFER_POOL_SHARDS, DEFAULT_BACKGROUND_POOL_SHARDS +from ._constants import DEFAULT_API_POOL_SHARDS, DEFAULT_TRANSFER_POOL_SHARDS, DEFAULT_BACKGROUND_POOL_SHARDS from ._streaming import Stream as Stream, AsyncStream as AsyncStream from ._exceptions import RunloopError, APIStatusError from ._base_client import ( @@ -97,9 +97,9 @@ def __init__( # Enables HTTP/2 multiplexing and avoids ConnectTimeout storms under high concurrency. # Set to False to create a private connection pool (old behavior). shared_http_pool: bool = True, - # Separate H2 connection pools for long-polls (/wait_for_status) and file - # transfers. Each shard ≈ one H2 connection; requests round-robin per client - # from a randomized starting offset (avoids every new client hitting shard 0). + # Sharded H2 pools by workload (API / long-polls / transfers). Each shard ≈ + # one H2 connection; requests round-robin per client from a random offset. + api_pool_shards: int = DEFAULT_API_POOL_SHARDS, background_pool_shards: int = DEFAULT_BACKGROUND_POOL_SHARDS, transfer_pool_shards: int = DEFAULT_TRANSFER_POOL_SHARDS, # Enable or disable schema validation for data returned by the API. @@ -148,6 +148,7 @@ def __init__( custom_query=default_query, _strict_response_validation=_strict_response_validation, shared_http_pool=shared_http_pool, + api_pool_shards=api_pool_shards, background_pool_shards=background_pool_shards, transfer_pool_shards=transfer_pool_shards, ) @@ -292,6 +293,7 @@ def copy( timeout: float | Timeout | None | NotGiven = not_given, http_client: httpx.Client | None = None, shared_http_pool: bool | None = None, + api_pool_shards: int | None = None, background_pool_shards: int | None = None, transfer_pool_shards: int | None = None, max_retries: int | NotGiven = not_given, @@ -335,6 +337,7 @@ def copy( timeout=self.timeout if isinstance(timeout, NotGiven) else timeout, http_client=http_client, shared_http_pool=resolved_shared, + api_pool_shards=(api_pool_shards if api_pool_shards is not None else self._api_pool_shards), background_pool_shards=( background_pool_shards if background_pool_shards is not None else self._background_pool_shards ), @@ -406,9 +409,9 @@ def __init__( # Enables HTTP/2 multiplexing and avoids ConnectTimeout storms under high concurrency. # Set to False to create a private connection pool (old behavior). shared_http_pool: bool = True, - # Separate H2 connection pools for long-polls (/wait_for_status) and file - # transfers. Each shard ≈ one H2 connection; requests round-robin per client - # from a randomized starting offset (avoids every new client hitting shard 0). + # Sharded H2 pools by workload (API / long-polls / transfers). Each shard ≈ + # one H2 connection; requests round-robin per client from a random offset. + api_pool_shards: int = DEFAULT_API_POOL_SHARDS, background_pool_shards: int = DEFAULT_BACKGROUND_POOL_SHARDS, transfer_pool_shards: int = DEFAULT_TRANSFER_POOL_SHARDS, # Enable or disable schema validation for data returned by the API. @@ -457,6 +460,7 @@ def __init__( custom_query=default_query, _strict_response_validation=_strict_response_validation, shared_http_pool=shared_http_pool, + api_pool_shards=api_pool_shards, background_pool_shards=background_pool_shards, transfer_pool_shards=transfer_pool_shards, ) @@ -601,6 +605,7 @@ def copy( timeout: float | Timeout | None | NotGiven = not_given, http_client: httpx.AsyncClient | None = None, shared_http_pool: bool | None = None, + api_pool_shards: int | None = None, background_pool_shards: int | None = None, transfer_pool_shards: int | None = None, max_retries: int | NotGiven = not_given, @@ -644,6 +649,7 @@ def copy( timeout=self.timeout if isinstance(timeout, NotGiven) else timeout, http_client=http_client, shared_http_pool=resolved_shared, + api_pool_shards=(api_pool_shards if api_pool_shards is not None else self._api_pool_shards), background_pool_shards=( background_pool_shards if background_pool_shards is not None else self._background_pool_shards ), diff --git a/src/runloop_api_client/_constants.py b/src/runloop_api_client/_constants.py index b91e63d90..6061b7708 100644 --- a/src/runloop_api_client/_constants.py +++ b/src/runloop_api_client/_constants.py @@ -10,10 +10,11 @@ DEFAULT_MAX_RETRIES = 5 DEFAULT_CONNECTION_LIMITS = httpx.Limits(max_connections=20, max_keepalive_connections=10) -# Separate H2 connection pools for long-polls and file transfers. Each "shard" is -# its own httpx client / transport (≈ one H2 connection). Default 2 shards ≈ two -# H2 connections so waits/uploads do not share the API control-plane connection. -DEFAULT_BACKGROUND_POOL_SHARDS = 2 +# Separate H2 connection pools by workload. Each "shard" is its own httpx +# client / transport (≈ one H2 connection). Defaults sized for Jetty's ~128 +# streams/connection: enough API/wait concurrency without oversized transfer. +DEFAULT_API_POOL_SHARDS = 8 +DEFAULT_BACKGROUND_POOL_SHARDS = 16 DEFAULT_TRANSFER_POOL_SHARDS = 2 INITIAL_RETRY_DELAY = 1.0 diff --git a/src/runloop_api_client/sdk/async_.py b/src/runloop_api_client/sdk/async_.py index 8639771a4..a23f35e2d 100644 --- a/src/runloop_api_client/sdk/async_.py +++ b/src/runloop_api_client/sdk/async_.py @@ -42,7 +42,7 @@ from .._client import DEFAULT_MAX_RETRIES, AsyncRunloop from ._helpers import detect_content_type from .async_axon import AsyncAxon -from .._constants import DEFAULT_TRANSFER_POOL_SHARDS, DEFAULT_BACKGROUND_POOL_SHARDS +from .._constants import DEFAULT_API_POOL_SHARDS, DEFAULT_TRANSFER_POOL_SHARDS, DEFAULT_BACKGROUND_POOL_SHARDS from .async_agent import AsyncAgent from .async_devbox import AsyncDevbox from .async_scorer import AsyncScorer @@ -1330,6 +1330,7 @@ def __init__( default_headers: Mapping[str, str] | None = None, default_query: Mapping[str, object] | None = None, http_client: httpx.AsyncClient | None = None, + api_pool_shards: int = DEFAULT_API_POOL_SHARDS, background_pool_shards: int = DEFAULT_BACKGROUND_POOL_SHARDS, transfer_pool_shards: int = DEFAULT_TRANSFER_POOL_SHARDS, ) -> None: @@ -1349,7 +1350,9 @@ def __init__( :type default_query: Mapping[str, object] | None, optional :param http_client: Custom ``httpx.AsyncClient`` instance to reuse, defaults to None :type http_client: httpx.AsyncClient | None, optional - :param background_pool_shards: H2 shards for long-polls (round-robin), defaults to 2 + :param api_pool_shards: H2 shards for short RPCs (round-robin), defaults to 8 + :type api_pool_shards: int, optional + :param background_pool_shards: H2 shards for long-polls (round-robin), defaults to 16 :type background_pool_shards: int, optional :param transfer_pool_shards: H2 shards for upload/download (round-robin), defaults to 2 :type transfer_pool_shards: int, optional @@ -1362,6 +1365,7 @@ def __init__( default_headers=default_headers, default_query=default_query, http_client=http_client, + api_pool_shards=api_pool_shards, background_pool_shards=background_pool_shards, transfer_pool_shards=transfer_pool_shards, ) diff --git a/src/runloop_api_client/sdk/sync.py b/src/runloop_api_client/sdk/sync.py index 29a91303a..ce10d1457 100644 --- a/src/runloop_api_client/sdk/sync.py +++ b/src/runloop_api_client/sdk/sync.py @@ -50,7 +50,7 @@ from .benchmark import Benchmark from .blueprint import Blueprint from .mcp_config import McpConfig -from .._constants import DEFAULT_TRANSFER_POOL_SHARDS, DEFAULT_BACKGROUND_POOL_SHARDS +from .._constants import DEFAULT_API_POOL_SHARDS, DEFAULT_TRANSFER_POOL_SHARDS, DEFAULT_BACKGROUND_POOL_SHARDS from .gateway_config import GatewayConfig from .network_policy import NetworkPolicy from .storage_object import StorageObject @@ -1355,6 +1355,7 @@ def __init__( default_headers: Mapping[str, str] | None = None, default_query: Mapping[str, object] | None = None, http_client: httpx.Client | None = None, + api_pool_shards: int = DEFAULT_API_POOL_SHARDS, background_pool_shards: int = DEFAULT_BACKGROUND_POOL_SHARDS, transfer_pool_shards: int = DEFAULT_TRANSFER_POOL_SHARDS, ) -> None: @@ -1374,7 +1375,9 @@ def __init__( :type default_query: Mapping[str, object] | None, optional :param http_client: Custom ``httpx.Client`` instance to reuse, defaults to None :type http_client: httpx.Client | None, optional - :param background_pool_shards: H2 shards for long-polls (round-robin), defaults to 2 + :param api_pool_shards: H2 shards for short RPCs (round-robin), defaults to 8 + :type api_pool_shards: int, optional + :param background_pool_shards: H2 shards for long-polls (round-robin), defaults to 16 :type background_pool_shards: int, optional :param transfer_pool_shards: H2 shards for upload/download (round-robin), defaults to 2 :type transfer_pool_shards: int, optional @@ -1387,6 +1390,7 @@ def __init__( default_headers=default_headers, default_query=default_query, http_client=http_client, + api_pool_shards=api_pool_shards, background_pool_shards=background_pool_shards, transfer_pool_shards=transfer_pool_shards, ) diff --git a/tests/test_h2_bulkhead_integration.py b/tests/test_h2_bulkhead_integration.py index 2af683e97..170651c9c 100644 --- a/tests/test_h2_bulkhead_integration.py +++ b/tests/test_h2_bulkhead_integration.py @@ -36,23 +36,18 @@ def _clear_pool_state() -> None: - with _base_mod._pool_lock: - old_sync = _base_mod._shared_sync_transport - old_bg = list(_base_mod._shared_sync_background_transports.values()) - old_xfer = list(_base_mod._shared_sync_transfer_transports.values()) - _base_mod._shared_sync_transport = None - _base_mod._shared_sync_background_transports.clear() - _base_mod._shared_sync_transfer_transports.clear() - _base_mod._shared_async_transports.clear() - _base_mod._shared_async_background_transports.clear() - _base_mod._shared_async_transfer_transports.clear() - for transport in [old_sync, *old_bg, *old_xfer]: - if transport is not None: - try: - transport._transport.close() - except Exception: - pass - + old = [] + old.extend(_base_mod._shared_sync_api_transports.take_all()) + old.extend(_base_mod._shared_sync_background_transports.take_all()) + old.extend(_base_mod._shared_sync_transfer_transports.take_all()) + _base_mod._shared_async_api_transports.clear() + _base_mod._shared_async_background_transports.clear() + _base_mod._shared_async_transfer_transports.clear() + for transport in old: + try: + transport._transport.close() + except Exception: + pass @pytest.fixture(autouse=True) def _reset_pools() -> Iterator[None]: # pyright: ignore[reportUnusedFunction] diff --git a/tests/test_shared_pool.py b/tests/test_shared_pool.py index b798b5baf..c60552fc7 100644 --- a/tests/test_shared_pool.py +++ b/tests/test_shared_pool.py @@ -28,23 +28,18 @@ def _reset_shared_pool() -> Iterator[None]: # pyright: ignore[reportUnusedFunct def _clear_pool_state() -> None: - with _base_mod._pool_lock: - old_sync = _base_mod._shared_sync_transport - old_bg = list(_base_mod._shared_sync_background_transports.values()) - old_xfer = list(_base_mod._shared_sync_transfer_transports.values()) - _base_mod._shared_sync_transport = None - _base_mod._shared_sync_background_transports.clear() - _base_mod._shared_sync_transfer_transports.clear() - _base_mod._shared_async_transports.clear() - _base_mod._shared_async_background_transports.clear() - _base_mod._shared_async_transfer_transports.clear() - for transport in [old_sync, *old_bg, *old_xfer]: - if transport is not None: - try: - transport._transport.close() - except Exception: - pass - + old = [] + old.extend(_base_mod._shared_sync_api_transports.take_all()) + old.extend(_base_mod._shared_sync_background_transports.take_all()) + old.extend(_base_mod._shared_sync_transfer_transports.take_all()) + _base_mod._shared_async_api_transports.clear() + _base_mod._shared_async_background_transports.clear() + _base_mod._shared_async_transfer_transports.clear() + for transport in old: + try: + transport._transport.close() + except Exception: + pass def _make_client(**kwargs: Any) -> Runloop: kwargs.setdefault("base_url", base_url) diff --git a/tests/test_transfer_client.py b/tests/test_transfer_client.py index 830993a19..1027c6888 100644 --- a/tests/test_transfer_client.py +++ b/tests/test_transfer_client.py @@ -42,22 +42,18 @@ def _reset_shared_pool() -> Iterator[None]: # pyright: ignore[reportUnusedFunct def _clear_pool_state() -> None: - with _base_mod._pool_lock: - old_sync = _base_mod._shared_sync_transport - old_bg = list(_base_mod._shared_sync_background_transports.values()) - old_xfer = list(_base_mod._shared_sync_transfer_transports.values()) - _base_mod._shared_sync_transport = None - _base_mod._shared_sync_background_transports.clear() - _base_mod._shared_sync_transfer_transports.clear() - _base_mod._shared_async_transports.clear() - _base_mod._shared_async_background_transports.clear() - _base_mod._shared_async_transfer_transports.clear() - for transport in [old_sync, *old_bg, *old_xfer]: - if transport is not None: - try: - transport._transport.close() - except Exception: - pass + old = [] + old.extend(_base_mod._shared_sync_api_transports.take_all()) + old.extend(_base_mod._shared_sync_background_transports.take_all()) + old.extend(_base_mod._shared_sync_transfer_transports.take_all()) + _base_mod._shared_async_api_transports.clear() + _base_mod._shared_async_background_transports.clear() + _base_mod._shared_async_transfer_transports.clear() + for transport in old: + try: + transport._transport.close() + except Exception: + pass def _make_client(**kwargs: Any) -> Runloop: @@ -76,17 +72,16 @@ def test_path_classification() -> None: def test_api_background_transfer_use_distinct_transports() -> None: - client = _make_client(shared_http_pool=True, background_pool_shards=2, transfer_pool_shards=2) + client = _make_client(shared_http_pool=True, api_pool_shards=2, background_pool_shards=2, transfer_pool_shards=2) try: api_req = httpx.Request("POST", f"{base_url}/v1/devboxes") wait_req = httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_a/wait_for_status") upload_req = httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_a/upload_file") - api = client._send_client_for_request(api_req) - wait = client._send_client_for_request(wait_req) - upload = client._send_client_for_request(upload_req) + api = client._get_client_for_path(api_req.url.path) + wait = client._get_client_for_path(wait_req.url.path) + upload = client._get_client_for_path(upload_req.url.path) - assert api is client._client assert wait is not api assert upload is not api assert wait is not upload @@ -99,22 +94,23 @@ def test_api_background_transfer_use_distinct_transports() -> None: def test_round_robin_spreads_requests_across_shards() -> None: client = _make_client(background_pool_shards=2, transfer_pool_shards=2) + assert client._background_pool is not None and client._transfer_pool is not None try: wait_req = httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_same/wait_for_status") - w0 = client._send_client_for_request(wait_req) - w1 = client._send_client_for_request(wait_req) - w2 = client._send_client_for_request(wait_req) + w0 = client._get_client_for_path(wait_req.url.path) + w1 = client._get_client_for_path(wait_req.url.path) + w2 = client._get_client_for_path(wait_req.url.path) assert w0 is not w1 assert w0 is w2 - assert set(client._background_clients) == {0, 1} + assert set(client._background_pool._clients) == {0, 1} upload_req = httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_same/upload_file") - t0 = client._send_client_for_request(upload_req) - t1 = client._send_client_for_request(upload_req) + t0 = client._get_client_for_path(upload_req.url.path) + t1 = client._get_client_for_path(upload_req.url.path) assert t0 is not t1 - assert set(client._transfer_clients) == {0, 1} - assert len(_base_mod._shared_sync_background_transports) == 2 - assert len(_base_mod._shared_sync_transfer_transports) == 2 + assert set(client._transfer_pool._clients) == {0, 1} + assert _base_mod._shared_sync_background_transports.shard_ids() == {0, 1} + assert _base_mod._shared_sync_transfer_transports.shard_ids() == {0, 1} finally: client.close() @@ -122,10 +118,10 @@ def test_round_robin_spreads_requests_across_shards() -> None: def test_many_clients_first_requests_use_both_shards() -> None: clients = [_make_client(background_pool_shards=2) for _ in range(40)] try: - req = httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_1/wait_for_status") - seen_transports = {id(c._send_client_for_request(req)._transport) for c in clients} # type: ignore[attr-defined] + path = "/v1/devboxes/dbx_1/wait_for_status" + seen_transports = {id(c._get_client_for_path(path)._transport) for c in clients} # type: ignore[attr-defined] assert len(seen_transports) == 2 - assert set(_base_mod._shared_sync_background_transports) == {0, 1} + assert _base_mod._shared_sync_background_transports.shard_ids() == {0, 1} finally: for c in clients: c.close() @@ -136,9 +132,8 @@ def test_custom_http_client_skips_isolation() -> None: client = _make_client(http_client=custom) try: assert client._isolate_workload_pools is False - req = httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_1/wait_for_status") - assert client._send_client_for_request(req) is custom - assert client._background_clients == {} + assert client._get_client_for_path("/v1/devboxes/dbx_1/wait_for_status") is custom + assert client._background_pool is None finally: client.close() custom.close() @@ -147,26 +142,25 @@ def test_custom_http_client_skips_isolation() -> None: def test_round_robin_is_per_client_while_transports_are_shared() -> None: c1 = _make_client(background_pool_shards=2) c2 = _make_client(background_pool_shards=2) + assert c1._background_pool is not None and c2._background_pool is not None try: - req = httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_shared/wait_for_status") - start1 = c1._background_next - start2 = c2._background_next - t1 = c1._send_client_for_request(req) - t2 = c2._send_client_for_request(req) + path = "/v1/devboxes/dbx_shared/wait_for_status" + start1 = c1._background_pool._next + start2 = c2._background_pool._next + t1 = c1._get_client_for_path(path) + t2 = c2._get_client_for_path(path) assert t1 is not t2 - # Same starting shard → same shared transport; different → different shards. if start1 % 2 == start2 % 2: assert t1._transport is t2._transport # type: ignore[attr-defined] else: assert t1._transport is not t2._transport # type: ignore[attr-defined] - assert c1._background_next == start1 + 1 - assert c2._background_next == start2 + 1 + assert c1._background_pool._next == start1 + 1 + assert c2._background_pool._next == start2 + 1 - # Advancing one client does not affect the other. - t1b = c1._send_client_for_request(req) + t1b = c1._get_client_for_path(path) assert t1b is not t1 - assert c1._background_next == start1 + 2 - assert c2._background_next == start2 + 1 + assert c1._background_pool._next == start1 + 2 + assert c2._background_pool._next == start2 + 1 finally: c1.close() c2.close() @@ -182,10 +176,8 @@ async def test_async_bulkheads() -> None: transfer_pool_shards=2, ) try: - wait = client._send_client_for_request( - httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_1/executions/ex_1/wait_for_status") - ) - upload = client._send_client_for_request(httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_1/upload_file")) + wait = client._get_client_for_path("/v1/devboxes/dbx_1/executions/ex_1/wait_for_status") + upload = client._get_client_for_path("/v1/devboxes/dbx_1/upload_file") assert wait is not client._client assert upload is not client._client assert wait is not upload @@ -195,6 +187,7 @@ async def test_async_bulkheads() -> None: def test_concurrent_background_client_init_is_singleton() -> None: client = _make_client(shared_http_pool=False, background_pool_shards=1) + assert client._background_pool is not None barrier = threading.Barrier(16) results: list[httpx.Client] = [] errors: list[BaseException] = [] @@ -202,7 +195,7 @@ def test_concurrent_background_client_init_is_singleton() -> None: def worker() -> None: try: barrier.wait() - results.append(client._ensure_background_client(0)) + results.append(client._background_pool.ensure(0)) # type: ignore[union-attr] except BaseException as exc: # noqa: BLE001 errors.append(exc) From e7042d79d23f8d3503078aad5d78079e9cebfc77 Mon Sep 17 00:00:00 2001 From: Reflex Date: Tue, 28 Jul 2026 22:04:53 +0000 Subject: [PATCH 10/11] chore: mark Registry test helpers; fix pyright on pool cleanup Co-authored-by: Cursor --- src/runloop_api_client/_base_client.py | 4 +++- tests/test_h2_bulkhead_integration.py | 3 ++- tests/test_shared_pool.py | 3 ++- tests/test_transfer_client.py | 2 +- 4 files changed, 8 insertions(+), 4 deletions(-) diff --git a/src/runloop_api_client/_base_client.py b/src/runloop_api_client/_base_client.py index fbef8e7aa..5401b4442 100644 --- a/src/runloop_api_client/_base_client.py +++ b/src/runloop_api_client/_base_client.py @@ -161,13 +161,14 @@ def acquire(self, shard: int) -> _SharedTransport: return transport def take_all(self) -> list[_SharedTransport]: - """Remove and return all transports (test cleanup).""" + """Test-only: remove and return all transports for fixture cleanup.""" with _pool_lock: values = list(self._shards.values()) self._shards.clear() return values def shard_ids(self) -> set[int]: + """Test-only: which shard indexes currently hold a shared transport.""" with _pool_lock: return set(self._shards) @@ -229,6 +230,7 @@ def acquire(self, loop: asyncio.AbstractEventLoop, shard: int) -> _SharedAsyncTr return transport def clear(self) -> None: + """Test-only: drop all per-loop shard maps (does not close transports).""" with _pool_lock: self._by_loop.clear() diff --git a/tests/test_h2_bulkhead_integration.py b/tests/test_h2_bulkhead_integration.py index 170651c9c..581241a31 100644 --- a/tests/test_h2_bulkhead_integration.py +++ b/tests/test_h2_bulkhead_integration.py @@ -36,7 +36,7 @@ def _clear_pool_state() -> None: - old = [] + old: list[_base_mod._SharedTransport] = [] old.extend(_base_mod._shared_sync_api_transports.take_all()) old.extend(_base_mod._shared_sync_background_transports.take_all()) old.extend(_base_mod._shared_sync_transfer_transports.take_all()) @@ -49,6 +49,7 @@ def _clear_pool_state() -> None: except Exception: pass + @pytest.fixture(autouse=True) def _reset_pools() -> Iterator[None]: # pyright: ignore[reportUnusedFunction] _clear_pool_state() diff --git a/tests/test_shared_pool.py b/tests/test_shared_pool.py index c60552fc7..0300d3625 100644 --- a/tests/test_shared_pool.py +++ b/tests/test_shared_pool.py @@ -28,7 +28,7 @@ def _reset_shared_pool() -> Iterator[None]: # pyright: ignore[reportUnusedFunct def _clear_pool_state() -> None: - old = [] + old: list[_base_mod._SharedTransport] = [] old.extend(_base_mod._shared_sync_api_transports.take_all()) old.extend(_base_mod._shared_sync_background_transports.take_all()) old.extend(_base_mod._shared_sync_transfer_transports.take_all()) @@ -41,6 +41,7 @@ def _clear_pool_state() -> None: except Exception: pass + def _make_client(**kwargs: Any) -> Runloop: kwargs.setdefault("base_url", base_url) kwargs.setdefault("bearer_token", bearer_token) diff --git a/tests/test_transfer_client.py b/tests/test_transfer_client.py index 1027c6888..b7fb81c9c 100644 --- a/tests/test_transfer_client.py +++ b/tests/test_transfer_client.py @@ -42,7 +42,7 @@ def _reset_shared_pool() -> Iterator[None]: # pyright: ignore[reportUnusedFunct def _clear_pool_state() -> None: - old = [] + old: list[_base_mod._SharedTransport] = [] old.extend(_base_mod._shared_sync_api_transports.take_all()) old.extend(_base_mod._shared_sync_background_transports.take_all()) old.extend(_base_mod._shared_sync_transfer_transports.take_all()) From a0f69b3adc8aa9eac9122d3ec713fd02adadfac3 Mon Sep 17 00:00:00 2001 From: Reflex Date: Tue, 28 Jul 2026 22:12:02 +0000 Subject: [PATCH 11/11] chore: ruff format _base_client.py Co-authored-by: Cursor --- src/runloop_api_client/_base_client.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/runloop_api_client/_base_client.py b/src/runloop_api_client/_base_client.py index 5401b4442..da359063d 100644 --- a/src/runloop_api_client/_base_client.py +++ b/src/runloop_api_client/_base_client.py @@ -210,9 +210,9 @@ class Registry: """Per-event-loop map of shard index → shared async H2 transport.""" def __init__(self) -> None: - self._by_loop: weakref.WeakKeyDictionary[ - asyncio.AbstractEventLoop, dict[int, _SharedAsyncTransport] - ] = weakref.WeakKeyDictionary() + self._by_loop: weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, dict[int, _SharedAsyncTransport]] = ( + weakref.WeakKeyDictionary() + ) def acquire(self, loop: asyncio.AbstractEventLoop, shard: int) -> _SharedAsyncTransport: with _pool_lock: