fix(python-sdk): keep streamed uploads off the retrying transport

.changeset/python-unify-pyqwest-connection-pools.md

@@ -0,0 +1,16 @@
+---
+"@e2b/python-sdk": patch
+---
+
+Run every persistent HTTP stack in the SDK on one shared pyqwest connection pool
+per proxy instead of four: the control-plane REST API, the envd HTTP API, the
+envd RPC clients, and the volume content API now all draw from
+`e2b.api.client_sync`/`client_async`. reqwest pools per host internally, so one
+pool serves the API host and every per-sandbox host without interference — and
+since envd RPC and the envd HTTP API hit the same host, an active sandbox now
+needs a single HTTP/2 connection instead of one per stack. Streamed downloads
+keep a pool of their own, the only one carrying the idle `read_timeout`:
+reqwest's read timer runs during body send and TTFB, so on a shared pool it
+would cut off long uploads. Layers above the pool are unchanged — the RPC
+stack's plain-HTTP error normalization still wraps it, and connect-only retries
+are now part of it.

.changeset/python-uploads-skip-connect-retries.md

@@ -0,0 +1,16 @@
+---
+"@e2b/python-sdk": patch
+---
+
+Stop mirroring streamed file uploads in memory. To be able to replay a request,
+pyqwest's retry middleware copies a streamed request body in full as it is sent,
+so an upload was streamed to the wire *and* held whole in RAM: writing a 64 MiB
+file with `volume.write_file`, `sandbox.files.write` or a template context upload
+peaked at ~67 MB of it. A streamed upload now goes out on an httpx transport over
+the same connection pool with the retry layer left out, so it keeps the pooled
+connections while the body is no longer copied — a 32 MiB upload peaks at 1.5 MB
+instead of 38 MB. What those uploads give up is the connect retry, which fires
+before any of the body was written, so a failure to connect surfaces to the
+caller as it would have anyway; writes of in-memory data are replayable for free
+and keep their retries. Template context uploads no longer build a connection
+pool of their own per build either.

packages/python-sdk/e2b/api/client_async/__init__.py

@@ -24,119 +24,196 @@ def get_api_client(config: ConnectionConfig, **kwargs) -> AsyncApiClient:
 
 
 class ConnectionRetryTransport(RetryTransport):
-    """Retry only failures establishing the connection — shared by the REST
-    API and envd RPC stacks: pyqwest raises the builtin ``ConnectionError``
-    only before the request was written, so these retries can never replay a
-    request the server may have received (a delivered REST call or unary RPC
-    like ``SendInput``). This matches the connect-only ``retries`` of the
-    httpx transports this replaced; the retry middleware's default policy
-    would otherwise also retry I/O errors and 429/5xx responses for
-    idempotent methods."""
+    """Retry only failures establishing the connection — part of the shared
+    transport stack, so it covers the REST API, the envd HTTP API, the envd RPC
+    clients and the volume content API alike: pyqwest raises the builtin
+    ``ConnectionError`` only before the request was written, so these retries
+    can never replay a request the server may have received (a delivered REST
+    call or unary RPC like ``SendInput``). This matches the connect-only
+    ``retries`` of the httpx transports this replaced; the retry middleware's
+    default policy would otherwise also retry I/O errors and 429/5xx responses
+    for idempotent methods.
+
+    Retries cost memory on a streamed request body: to be able to replay it,
+    the middleware copies the body in full as it is sent (reqwest reads ahead
+    into it while connecting, so a connect error leaves the iterator already
+    started and unrewindable). File uploads therefore skip this layer and go
+    straight to the pool underneath — see :func:`get_upload_transport`."""
 
     def should_retry_response(
         self, request: Request, response: Union[Response, Exception]
     ) -> bool:
         return isinstance(response, ConnectionError)
 
 
-def retrying_http_transport(
-    proxy: Optional[ProxyConfig], read_timeout: Optional[float] = None
-) -> ConnectionRetryTransport:
-    """A fresh pyqwest transport (= its own connection pool) with the SDK's
-    shared tuning — system CA certs (without which TLS through an
-    intercepting proxy fails), the httpx-equivalent pool limits, and
-    connect-only retries. The REST API, envd RPC, and envd HTTP API stacks
-    each cache their own instances (pool unification is a follow-up).
+TransportKey = Tuple[Optional[ProxyConfig], Optional[float]]
+"""Cache key of the shared transports: the proxy and the idle read bound."""
 
-    ``read_timeout`` bounds every read on the transport's connections; see
-    :func:`get_envd_transport` for when that is (and isn't) appropriate.
+_transport_lock = threading.Lock()
+# One pyqwest transport — one reqwest connection pool — per proxy and idle read
+# bound; `None` is the direct pool. Every HTTP stack in the SDK draws from
+# these: the control-plane REST API, the envd HTTP API, the envd RPC clients
+# (`e2b.envd.client_async`) and the volume content API
+# (`e2b.volume.client_async`). reqwest pools per host internally, so a single
+# pool serves the API host and every per-sandbox host without interference —
+# and because envd RPC and the envd HTTP API share it, a sandbox needs one
+# HTTP/2 connection instead of one per stack.
+#
+# pyqwest's I/O runs on its own Rust runtime, so unlike the httpx transports
+# they replaced, the transports are not bound to an event loop and the caches
+# are process-global rather than per-loop.
+_pools: Dict[TransportKey, HTTPTransport] = {}
+# The connect-retrying stack over each pool.
+_transports: Dict[TransportKey, ConnectionRetryTransport] = {}
+# The httpx adapter over each pool, shared by every httpx client on it — one per
+# pool with the retries, one without for uploads.
+_httpx_transports: Dict[TransportKey, AsyncPyqwestTransport] = {}
+_upload_transports: Dict[TransportKey, AsyncPyqwestTransport] = {}
+
+
+def get_pool(
+    proxy: Optional[ProxyConfig], read_timeout: Optional[float] = None
+) -> HTTPTransport:
+    """The shared connection pool with the SDK's tuning — system CA certs
+    (without which TLS through an intercepting proxy fails) and the
+    httpx-equivalent pool limits. For TLS connections ALPN negotiates the HTTP
+    version (HTTP/2 against the E2B API and envd), like the http2-enabled httpx
+    transports this replaced.
+
+    Layer concerns above it rather than into it — connect retries wrap it, so
+    does the RPC stack's plain-HTTP-error normalization, and headers and codecs
+    are per-request — so that the pool stays shareable.
+
+    ``read_timeout`` bounds every read on the pool's connections, which is why
+    it is part of the cache key: reqwest's read timer keeps running while a
+    request body is sent and while waiting for the response head, so a pool
+    carrying one would cut off long uploads and slow responses. Only streamed
+    downloads ask for it, as an idle bound (see :func:`get_transport`).
 
     Requests are logged by pyqwest itself on the ``pyqwest.access`` and
     ``pyqwest`` loggers at ``DEBUG`` (off unless enabled) — the transport-level
     diagnostics httpcore used to provide. The SDK's own ``logger`` option is
     separate and sits above this, on the httpx client."""
-    return ConnectionRetryTransport(
-        HTTPTransport(
-            tls_include_system_certs=True,
-            proxy=proxy.to_pyqwest() if proxy is not None else None,
-            pool_idle_timeout=pool_idle_timeout,
-            pool_max_idle_per_host=pool_max_idle_per_host,
-            read_timeout=read_timeout,
-            # Redirects belong to the httpx client above (which the generated
-            # clients leave off), not to reqwest.
-            follow_redirects=False,
-        ),
-        max_retries=connection_retries,
-    )
+    key = (proxy, read_timeout)
+    with _transport_lock:
+        pool = _pools.get(key)
+        if pool is None:
+            pool = HTTPTransport(
+                tls_include_system_certs=True,
+                proxy=proxy.to_pyqwest() if proxy is not None else None,
+                pool_idle_timeout=pool_idle_timeout,
+                pool_max_idle_per_host=pool_max_idle_per_host,
+                read_timeout=read_timeout,
+                # Redirects belong to the httpx client above (which the
+                # generated clients leave off), not to reqwest.
+                follow_redirects=False,
+            )
+            _pools[key] = pool
+        return pool
 
 
-_transport_lock = threading.Lock()
-# One transport (= one connection pool) per proxy; None is the direct pool.
-# pyqwest's I/O runs on its own Rust runtime, so unlike the httpx transports
-# they replaced, the transports are not bound to an event loop and the
-# caches are process-global rather than per-loop.
-_transports: Dict[Optional[ProxyConfig], AsyncPyqwestTransport] = {}
+def get_pyqwest_transport(
+    proxy: Optional[ProxyConfig], read_timeout: Optional[float] = None
+) -> ConnectionRetryTransport:
+    """The shared pool of :func:`get_pool` behind connect-only retries, for
+    consumers speaking pyqwest natively (the envd RPC clients). Consumers
+    speaking httpx take :func:`get_httpx_transport`, the adapter over the very
+    same stack."""
+    key = (proxy, read_timeout)
+    # Resolve the pool before taking the lock: it takes the same one.
+    pool = get_pool(proxy, read_timeout)
+    with _transport_lock:
+        transport = _transports.get(key)
+        if transport is None:
+            transport = ConnectionRetryTransport(pool, max_retries=connection_retries)
+            _transports[key] = transport
+        return transport
 
 
-def get_transport(config: ConnectionConfig) -> AsyncPyqwestTransport:
-    """The shared pyqwest-backed httpx transport for REST API calls. For TLS
-    connections ALPN negotiates the HTTP version (HTTP/2 against the E2B
-    API), like the http2-enabled httpx transport this replaced."""
-    proxy = proxy_to_config(config.proxy)
+def get_httpx_transport(
+    proxy: Optional[ProxyConfig], read_timeout: Optional[float] = None
+) -> AsyncPyqwestTransport:
+    """The httpx adapter over the retrying stack of
+    :func:`get_pyqwest_transport`, for the generated httpx clients (control
+    plane, envd HTTP API, volume content). The adapter holds no state of its
+    own and does not close the pool, so closing an httpx client leaves the
+    pool intact for the other clients on it."""
+    key = (proxy, read_timeout)
+    stack = get_pyqwest_transport(proxy, read_timeout)
     with _transport_lock:
-        transport = _transports.get(proxy)
+        transport = _httpx_transports.get(key)
         if transport is None:
-            transport = AsyncPyqwestTransport(retrying_http_transport(proxy))
-            _transports[proxy] = transport
+            transport = AsyncPyqwestTransport(stack)
+            _httpx_transports[key] = transport
         return transport
 
 
-# One transport per (proxy, streaming) pair, separate from the REST API
-# pools — envd traffic goes to per-sandbox hosts.
-_envd_transports: Dict[Tuple[Optional[ProxyConfig], bool], AsyncPyqwestTransport] = {}
+def get_upload_transport(
+    proxy: Optional[ProxyConfig], read_timeout: Optional[float] = None
+) -> AsyncPyqwestTransport:
+    """The httpx adapter over the *same pool* as :func:`get_httpx_transport`
+    but without the retry layer, for file uploads: envd ``files.write``, volume
+    ``write_file`` and template context uploads.
+
+    Those are the requests whose body is streamed rather than held in memory,
+    and the retry middleware makes a body replayable by copying it in full as
+    it is sent — so uploading a file on the retrying stack would mirror the
+    whole file in RAM (SDK-332). Skipping the layer, not the pool, keeps the
+    connection reuse: an upload still travels the sandbox's or the volume
+    host's pooled connections. What it gives up is the connect retry, which
+    fires before any of the body was written, so the caller sees the connection
+    error and can retry the upload itself."""
+    key = (proxy, read_timeout)
+    pool = get_pool(proxy, read_timeout)
+    with _transport_lock:
+        transport = _upload_transports.get(key)
+        if transport is None:
+            transport = AsyncPyqwestTransport(pool)
+            _upload_transports[key] = transport
+        return transport
 
 
-def get_envd_transport(
-    config: ConnectionConfig, *, for_streaming: bool = False
+def get_transport(
+    config: ConnectionConfig, *, for_streaming: bool = False, for_upload: bool = False
 ) -> AsyncPyqwestTransport:
-    """The shared pyqwest-backed httpx transports for the envd HTTP API
-    (file transfers, health checks).
-
-    The streaming transport carries ``read_timeout``, the idle bound on
-    every read: it resets after each successful read, so it caps how long a
-    streamed download may stall without limiting total transfer time. It is
-    fixed per transport — the adapter's per-request timeouts are
-    whole-request deadlines. Only streamed downloads use it: reqwest's read
-    timer keeps running while a request body is sent and while waiting for
-    the response head, so on the regular transport it would cut off uploads
-    and slow unary responses longer than the idle bound (those stay bounded
-    by their whole-request deadlines instead).
+    """The shared httpx transport for the control-plane REST API and the envd
+    HTTP API (file transfers, health checks) — one pool serves both, keyed by
+    the connection's proxy.
+
+    ``for_streaming`` selects the pool carrying ``READ_TIMEOUT``, the idle
+    bound on every read: it resets after each successful read, so it caps how
+    long a streamed download may stall without limiting total transfer time.
+    It is fixed per pool — the adapter's per-request timeouts are
+    whole-request deadlines rather than idle bounds — so only streamed
+    downloads take it, and they get their own pool (see :func:`get_pool`).
+
+    ``for_upload`` selects the same pool without the retry layer, so that a
+    streamed upload is not copied into memory to keep it replayable
+    (see :func:`get_upload_transport`).
     """
     proxy = proxy_to_config(config.proxy)
-    key = (proxy, for_streaming)
-    with _transport_lock:
-        transport = _envd_transports.get(key)
-        if transport is None:
-            transport = AsyncPyqwestTransport(
-                retrying_http_transport(
-                    proxy,
-                    read_timeout=READ_TIMEOUT if for_streaming else None,
-                )
-            )
-            _envd_transports[key] = transport
-        return transport
+    read_timeout = READ_TIMEOUT if for_streaming else None
+    if for_upload:
+        return get_upload_transport(proxy, read_timeout)
+    return get_httpx_transport(proxy, read_timeout)
 
 
 def get_envd_api(
-    config: ConnectionConfig, base_url: str, *, for_streaming: bool = False
+    config: ConnectionConfig,
+    base_url: str,
+    *,
+    for_streaming: bool = False,
+    for_upload: bool = False,
 ) -> httpx.AsyncClient:
     """An httpx client for a sandbox's envd HTTP API (file transfers, health
-    checks) on the shared pyqwest transports. The client itself is a cheap
-    stateless wrapper — one per consumer is fine — while the pooled transport
-    underneath is shared and loop-independent."""
+    checks) on the shared transports. The client itself is a cheap stateless
+    wrapper — one per consumer is fine — while the pooled transport underneath
+    is shared and loop-independent."""
     return httpx.AsyncClient(
         base_url=base_url,
-        transport=get_envd_transport(config, for_streaming=for_streaming),
+        transport=get_transport(
+            config, for_streaming=for_streaming, for_upload=for_upload
+        ),
         headers=config.sandbox_headers,
         event_hooks=make_async_logging_event_hooks(config.logger),
     )