From b98324ee0a687aa1194300f5e038550bc1162c7b Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 19 Aug 2026 16:56:19 +0000 Subject: [PATCH 1/3] fix(python-sdk): stop buffering streamed request bodies on the retrying transports pyqwest's retry middleware keeps a request replayable, and for a body that is not already bytes it does that by growing a copy of it as the body is sent. A streamed upload was therefore sent incrementally and mirrored whole in RAM, so peak memory scaled with file size for files.write of a file-like object, volume.write_file and template context uploads (SDK-332). Declare pyqwest's unbuffered retry mode on the shared retrying transports instead: a streamed body is replayed only while nothing has been read from it, which is all the connect-only policy needs, since a connect error means the request was never written. bytes bodies keep their retries either way. Releases whose retry middleware has no RetryMode keep the buffered behavior they ship, where going unbuffered would cost the connect retries instead. Co-authored-by: Mish Ushakov --- ...python-unbuffered-retry-streamed-bodies.md | 14 + packages/python-sdk/e2b/api/__init__.py | 29 ++- .../e2b/api/client_async/__init__.py | 11 +- .../e2b/api/client_sync/__init__.py | 11 +- .../tests/test_retry_streamed_bodies.py | 244 ++++++++++++++++++ 5 files changed, 304 insertions(+), 5 deletions(-) create mode 100644 .changeset/python-unbuffered-retry-streamed-bodies.md create mode 100644 packages/python-sdk/tests/test_retry_streamed_bodies.py diff --git a/.changeset/python-unbuffered-retry-streamed-bodies.md b/.changeset/python-unbuffered-retry-streamed-bodies.md new file mode 100644 index 0000000000..4ff375bc55 --- /dev/null +++ b/.changeset/python-unbuffered-retry-streamed-bodies.md @@ -0,0 +1,14 @@ +--- +"@e2b/python-sdk": patch +--- + +Stop mirroring streamed request bodies in memory while they are being sent. The +shared retrying transports now declare pyqwest's unbuffered retry policy, so a +streamed upload — `files.write` of a file-like object, `volume.write_file`, +template build contexts — is no longer copied into a replay buffer as it goes to +the wire, and peak memory stays flat instead of scaling with file size. Connect +retries are unchanged: a connect error means the body was never read, so the +same body is handed to the next attempt, and `bytes` bodies (unary RPC payloads, +in-memory writes) were already replayable for free. Requires a pyqwest whose +retry middleware exposes `RetryMode`; on older releases the SDK keeps that +release's buffered behavior. diff --git a/packages/python-sdk/e2b/api/__init__.py b/packages/python-sdk/e2b/api/__init__.py index 2db0fa80c6..95ca7bd146 100644 --- a/packages/python-sdk/e2b/api/__init__.py +++ b/packages/python-sdk/e2b/api/__init__.py @@ -4,12 +4,13 @@ import re from dataclasses import dataclass from types import TracebackType -from typing import NamedTuple, Optional, Protocol, Tuple, Union +from typing import Any, NamedTuple, Optional, Protocol, Tuple, Union from urllib.parse import quote import httpx from httpx import AsyncBaseTransport, BaseTransport, Timeout from pyqwest import Proxy +from pyqwest.middleware import retry as retry_middleware from e2b.api.client.client import AuthenticatedClient from e2b.api.client.types import Response @@ -85,6 +86,32 @@ async def on_response(response: Response) -> None: pool_idle_timeout = float(os.getenv("E2B_KEEPALIVE_EXPIRY") or "300") pool_max_idle_per_host = int(os.getenv("E2B_MAX_KEEPALIVE_CONNECTIONS") or "20") +# The retry policy the shared transports declare per request (see +# `ConnectionRetryTransport` in client_sync/client_async). +# +# pyqwest's retry middleware keeps a request replayable, and for a body that +# isn't already `bytes` it does that by growing a copy of it in memory as it is +# sent. A streamed upload therefore went to the wire incrementally *and* was +# mirrored whole in RAM, so peak memory scaled with file size for +# `files.write`, `volume.write_file` and template context uploads even though +# nothing in the SDK buffers (SDK-332). +# +# `RetryMode.UNBUFFERED` drops the copy: a streamed body is replayed only while +# nothing has been read from it, which is all a connect-only policy ever needs, +# since a connect error means the request was never written. `bytes` bodies — +# every unary RPC payload and every in-memory write — are replayable as they +# are and keep their retries either way. +# +# Releases without `RetryMode` keep the buffered behavior they ship — `True` is +# what the middleware understood before the mode existed — because going +# unbuffered there would trade the copy for the connect retries rather than get +# them both: reqwest used to read ahead into its body channel while connecting, +# so a connect error surfaced with the body already started (measured: the one +# chunk of a unary RPC body, up to a whole small-chunked upload) and neither +# buffered nor rewindable. +_retry_mode = getattr(retry_middleware, "RetryMode", None) +unbuffered_retries: Any = True if _retry_mode is None else _retry_mode.UNBUFFERED + class ProxyConfig(NamedTuple): """The ``proxy`` connection option in the shape pyqwest transports take. diff --git a/packages/python-sdk/e2b/api/client_async/__init__.py b/packages/python-sdk/e2b/api/client_async/__init__.py index 678cc29e6d..2fc175c660 100644 --- a/packages/python-sdk/e2b/api/client_async/__init__.py +++ b/packages/python-sdk/e2b/api/client_async/__init__.py @@ -1,5 +1,5 @@ import threading -from typing import Dict, Optional, Tuple, Union +from typing import Any, Dict, Optional, Tuple, Union import httpx @@ -15,6 +15,7 @@ pool_idle_timeout, pool_max_idle_per_host, proxy_to_config, + unbuffered_retries, ) from e2b.connection_config import READ_TIMEOUT, ConnectionConfig @@ -32,7 +33,13 @@ class ConnectionRetryTransport(RetryTransport): 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.""" + for idempotent methods. + + Streamed bodies are retried without being buffered, so an upload is not + mirrored in memory as it is sent (see ``unbuffered_retries``).""" + + def should_retry_request(self, request: Request) -> Any: + return unbuffered_retries def should_retry_response( self, request: Request, response: Union[Response, Exception] diff --git a/packages/python-sdk/e2b/api/client_sync/__init__.py b/packages/python-sdk/e2b/api/client_sync/__init__.py index 73c279acec..f004a01d00 100644 --- a/packages/python-sdk/e2b/api/client_sync/__init__.py +++ b/packages/python-sdk/e2b/api/client_sync/__init__.py @@ -1,4 +1,4 @@ -from typing import Dict, Optional, Tuple, Union +from typing import Any, Dict, Optional, Tuple, Union import httpx import threading @@ -15,6 +15,7 @@ pool_idle_timeout, pool_max_idle_per_host, proxy_to_config, + unbuffered_retries, ) from e2b.connection_config import READ_TIMEOUT, ConnectionConfig @@ -32,7 +33,13 @@ class ConnectionRetryTransport(SyncRetryTransport): 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.""" + for idempotent methods. + + Streamed bodies are retried without being buffered, so an upload is not + mirrored in memory as it is sent (see ``unbuffered_retries``).""" + + def should_retry_request(self, request: SyncRequest) -> Any: + return unbuffered_retries def should_retry_response( self, request: SyncRequest, response: Union[SyncResponse, Exception] diff --git a/packages/python-sdk/tests/test_retry_streamed_bodies.py b/packages/python-sdk/tests/test_retry_streamed_bodies.py new file mode 100644 index 0000000000..3fc9f1f168 --- /dev/null +++ b/packages/python-sdk/tests/test_retry_streamed_bodies.py @@ -0,0 +1,244 @@ +"""Streamed request bodies pass through the shared retrying transports without +being mirrored in memory (SDK-332). + +pyqwest's retry middleware keeps a request replayable, and for a body that is +not already ``bytes`` it does that by growing a copy of the body as it is sent. +An upload therefore reached the wire incrementally *and* was held whole in RAM, +so peak memory scaled with file size for ``files.write``, ``volume.write_file`` +and template context uploads. The transports now declare +``unbuffered_retries``, which drops the copy and replays a streamed body only +while nothing has been read from it — all the connect-only policy needs, since +pyqwest starts reading a body when hyper first writes it and a connect error +means it never got that far. + +``bytes`` bodies (every unary RPC payload, every in-memory write) are replayable +as they are, so their retries are unaffected; the tests for those live in +``test_envd_retry_transport.py``. + +The tests that pin the unbuffered half are skipped on a pyqwest whose retry +middleware has no ``RetryMode``: there the SDK keeps that release's buffered +behavior, because going unbuffered without pyqwest's deferred body reads would +trade the copy for the connect retries. +""" + +import socket +import tracemalloc +from typing import AsyncIterator, Iterator, List + +import pytest +from pyqwest import ( + HTTPTransport, + Request, + Response, + SyncHTTPTransport, + SyncRequest, + SyncResponse, +) + +from e2b.api import unbuffered_retries +from e2b.api.client_async import ConnectionRetryTransport +from e2b.api.client_sync import ( + ConnectionRetryTransport as SyncConnectionRetryTransport, +) + +CHUNK_SIZE = 256 * 1024 +CHUNK_COUNT = 64 +BODY_SIZE = CHUNK_SIZE * CHUNK_COUNT +"""16 MiB in 256 KiB chunks: large enough that a mirrored copy dwarfs the +per-chunk allocations, small enough to stay quick.""" + +unbuffered_only = pytest.mark.skipif( + unbuffered_retries is True, + reason="pyqwest without RetryMode buffers streamed bodies to replay them", +) + +MAX_TRACED_PEAK = BODY_SIZE // 4 + + +def _retrying(inner) -> ConnectionRetryTransport: + # Keep the exponential backoff out of test wall-clock time. + return ConnectionRetryTransport( + inner, initial_interval=0.001, max_interval=0.002, max_retries=3 + ) + + +def _retrying_sync(inner) -> SyncConnectionRetryTransport: + return SyncConnectionRetryTransport( + inner, initial_interval=0.001, max_interval=0.002, max_retries=3 + ) + + +def _chunks(pulled: List[int]) -> Iterator[bytes]: + """A streamed body that records how much of it was read.""" + for _ in range(CHUNK_COUNT): + pulled.append(CHUNK_SIZE) + yield b"x" * CHUNK_SIZE + + +async def _achunks(pulled: List[int]) -> AsyncIterator[bytes]: + for _ in range(CHUNK_COUNT): + pulled.append(CHUNK_SIZE) + yield b"x" * CHUNK_SIZE + + +def _refused_port() -> int: + """A port nothing listens on, so connecting to it fails immediately.""" + sock = socket.socket() + sock.bind(("127.0.0.1", 0)) + port = sock.getsockname()[1] + sock.close() + return port + + +class SendingTransport: + """Inner async transport that reads the request body the way reqwest does, + after failing the first ``failures`` attempts the way a refused connect + does — before the body was touched.""" + + def __init__(self, failures: int = 0, fail_mid_body: bool = False): + self.failures = failures + self.fail_mid_body = fail_mid_body + self.attempts = 0 + self.received = 0 + + async def execute(self, request: Request) -> Response: + self.attempts += 1 + if self.attempts <= self.failures and not self.fail_mid_body: + raise ConnectionError("tcp connect error") + content = request.content + if isinstance(content, bytes): + self.received += len(content) + else: + async for chunk in content: + self.received += len(chunk) + if self.attempts <= self.failures and self.fail_mid_body: + raise ConnectionError("tcp connect error") + return Response(status=200, content=b"ok") + + +class SendingSyncTransport: + def __init__(self, failures: int = 0, fail_mid_body: bool = False): + self.failures = failures + self.fail_mid_body = fail_mid_body + self.attempts = 0 + self.received = 0 + + def execute_sync(self, request: SyncRequest) -> SyncResponse: + self.attempts += 1 + if self.attempts <= self.failures and not self.fail_mid_body: + raise ConnectionError("tcp connect error") + content = request.content + if isinstance(content, bytes): + self.received += len(content) + else: + for chunk in content: + self.received += len(chunk) + if self.attempts <= self.failures and self.fail_mid_body: + raise ConnectionError("tcp connect error") + return SyncResponse(status=200, content=b"ok") + + +@unbuffered_only +async def test_async_streamed_body_is_not_mirrored_in_memory(): + inner = SendingTransport() + pulled: List[int] = [] + request = Request("PUT", "http://sandbox.test/upload", content=_achunks(pulled)) + tracemalloc.start() + try: + response = await _retrying(inner).execute(request) + peak = tracemalloc.get_traced_memory()[1] + finally: + tracemalloc.stop() + assert response.status == 200 + # The whole body reached the transport, but only a chunk of it at a time. + assert inner.received == BODY_SIZE + assert sum(pulled) == BODY_SIZE + assert peak < MAX_TRACED_PEAK, peak + + +@unbuffered_only +def test_sync_streamed_body_is_not_mirrored_in_memory(): + inner = SendingSyncTransport() + pulled: List[int] = [] + request = SyncRequest("PUT", "http://sandbox.test/upload", content=_chunks(pulled)) + tracemalloc.start() + try: + response = _retrying_sync(inner).execute_sync(request) + peak = tracemalloc.get_traced_memory()[1] + finally: + tracemalloc.stop() + assert response.status == 200 + assert inner.received == BODY_SIZE + assert sum(pulled) == BODY_SIZE + assert peak < MAX_TRACED_PEAK, peak + + +async def test_async_failed_connect_replays_an_untouched_streamed_body(): + inner = SendingTransport(failures=1) + pulled: List[int] = [] + request = Request("PUT", "http://sandbox.test/upload", content=_achunks(pulled)) + response = await _retrying(inner).execute(request) + assert response.status == 200 + assert inner.attempts == 2 + # The retry sent the body from the start: nothing had been read from it. + assert inner.received == BODY_SIZE + + +def test_sync_failed_connect_replays_an_untouched_streamed_body(): + inner = SendingSyncTransport(failures=1) + pulled: List[int] = [] + request = SyncRequest("PUT", "http://sandbox.test/upload", content=_chunks(pulled)) + response = _retrying_sync(inner).execute_sync(request) + assert response.status == 200 + assert inner.attempts == 2 + assert inner.received == BODY_SIZE + + +@unbuffered_only +async def test_async_streamed_body_is_not_replayed_once_it_was_read(): + # A body with no copy behind it cannot be rewound, so a failure that + # arrives after the first chunk went out surfaces instead of replaying a + # truncated request. + inner = SendingTransport(failures=1, fail_mid_body=True) + pulled: List[int] = [] + request = Request("PUT", "http://sandbox.test/upload", content=_achunks(pulled)) + with pytest.raises(ConnectionError): + await _retrying(inner).execute(request) + assert inner.attempts == 1 + + +@unbuffered_only +def test_sync_streamed_body_is_not_replayed_once_it_was_read(): + inner = SendingSyncTransport(failures=1, fail_mid_body=True) + pulled: List[int] = [] + request = SyncRequest("PUT", "http://sandbox.test/upload", content=_chunks(pulled)) + with pytest.raises(ConnectionError): + _retrying_sync(inner).execute_sync(request) + assert inner.attempts == 1 + + +# Through a real pyqwest transport: the retries above are only safe because a +# connect failure leaves the body unread, which is what lets the middleware +# hand the same iterator to the next attempt. + + +@unbuffered_only +async def test_async_refused_connect_leaves_a_streamed_body_unread(): + url = f"http://127.0.0.1:{_refused_port()}/upload" + pulled: List[int] = [] + with pytest.raises(ConnectionError): + await _retrying(HTTPTransport()).execute( + Request("PUT", url, content=_achunks(pulled)) + ) + assert pulled == [] + + +@unbuffered_only +def test_sync_refused_connect_leaves_a_streamed_body_unread(): + url = f"http://127.0.0.1:{_refused_port()}/upload" + pulled: List[int] = [] + with pytest.raises(ConnectionError): + _retrying_sync(SyncHTTPTransport()).execute_sync( + SyncRequest("PUT", url, content=_chunks(pulled)) + ) + assert pulled == [] From 7f12167acf51899bc613300f722c32d0fdf55095 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 19 Aug 2026 17:08:23 +0000 Subject: [PATCH 2/3] test(python-sdk): assert the streamed body is read once across a retry Co-authored-by: Mish Ushakov --- packages/python-sdk/tests/test_retry_streamed_bodies.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/packages/python-sdk/tests/test_retry_streamed_bodies.py b/packages/python-sdk/tests/test_retry_streamed_bodies.py index 3fc9f1f168..ee402f859b 100644 --- a/packages/python-sdk/tests/test_retry_streamed_bodies.py +++ b/packages/python-sdk/tests/test_retry_streamed_bodies.py @@ -180,8 +180,10 @@ async def test_async_failed_connect_replays_an_untouched_streamed_body(): response = await _retrying(inner).execute(request) assert response.status == 200 assert inner.attempts == 2 - # The retry sent the body from the start: nothing had been read from it. + # The retry sent the body from the start, reading the source exactly once: + # nothing had been read from it when the connect failed. assert inner.received == BODY_SIZE + assert sum(pulled) == BODY_SIZE def test_sync_failed_connect_replays_an_untouched_streamed_body(): @@ -192,6 +194,7 @@ def test_sync_failed_connect_replays_an_untouched_streamed_body(): assert response.status == 200 assert inner.attempts == 2 assert inner.received == BODY_SIZE + assert sum(pulled) == BODY_SIZE @unbuffered_only @@ -205,6 +208,7 @@ async def test_async_streamed_body_is_not_replayed_once_it_was_read(): with pytest.raises(ConnectionError): await _retrying(inner).execute(request) assert inner.attempts == 1 + assert inner.received == CHUNK_SIZE @unbuffered_only @@ -215,6 +219,7 @@ def test_sync_streamed_body_is_not_replayed_once_it_was_read(): with pytest.raises(ConnectionError): _retrying_sync(inner).execute_sync(request) assert inner.attempts == 1 + assert inner.received == CHUNK_SIZE # Through a real pyqwest transport: the retries above are only safe because a From 945624ce8ac6ffad07708e85073b4cc00e55651d Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 19 Aug 2026 18:12:11 +0000 Subject: [PATCH 3/3] fix(python-sdk): correct the affected paths and pin both retry policies Review found three overclaims and a gap: - template context uploads were never mirrored by the retry middleware, since they build their own transport without it, so drop them from the changeset, the module comment and the test docstring; - the transports' docstrings promised unbuffered streaming unconditionally, which is not true of the pinned pyqwest, so condition them the way the module comment already does; - the policy constant read as a boolean flag while holding a policy, and spelled its buffered value `True`, so rename it to _retry_request_policy and prefix it like the other internals in this family. The gap: on the pinned pyqwest the inherited hook already returns `True`, so nothing failed if the override or the RetryMode lookup were removed. Resolving the policy is now a function, and the tests pin both of its arms, the presence of both overrides, which regime is live, the buffered replay the fallback relies on, and the streamed body shape the httpx adapter hands the middleware. Co-authored-by: Mish Ushakov --- ...python-unbuffered-retry-streamed-bodies.md | 20 +- packages/python-sdk/e2b/api/__init__.py | 20 +- .../e2b/api/client_async/__init__.py | 10 +- .../e2b/api/client_sync/__init__.py | 10 +- .../tests/test_retry_streamed_bodies.py | 206 ++++++++++++++++-- 5 files changed, 230 insertions(+), 36 deletions(-) diff --git a/.changeset/python-unbuffered-retry-streamed-bodies.md b/.changeset/python-unbuffered-retry-streamed-bodies.md index 4ff375bc55..e69e834ae8 100644 --- a/.changeset/python-unbuffered-retry-streamed-bodies.md +++ b/.changeset/python-unbuffered-retry-streamed-bodies.md @@ -2,13 +2,13 @@ "@e2b/python-sdk": patch --- -Stop mirroring streamed request bodies in memory while they are being sent. The -shared retrying transports now declare pyqwest's unbuffered retry policy, so a -streamed upload — `files.write` of a file-like object, `volume.write_file`, -template build contexts — is no longer copied into a replay buffer as it goes to -the wire, and peak memory stays flat instead of scaling with file size. Connect -retries are unchanged: a connect error means the body was never read, so the -same body is handed to the next attempt, and `bytes` bodies (unary RPC payloads, -in-memory writes) were already replayable for free. Requires a pyqwest whose -retry middleware exposes `RetryMode`; on older releases the SDK keeps that -release's buffered behavior. +Stop mirroring streamed request bodies in memory while they are being sent, on a +pyqwest whose retry middleware exposes `RetryMode` (older releases keep the +buffered replay they ship, so this is inert until the `pyqwest` dependency +allows one). The shared retrying transports declare pyqwest's unbuffered retry +policy, so a streamed `files.write` of a file-like object or `volume.write_file` +is no longer copied into a replay buffer as it goes to the wire, and peak memory +stays flat instead of scaling with file size. Connect retries are unchanged: a +connect error means the body was never read, so the same body is handed to the +next attempt, and `bytes` bodies (unary RPC payloads, in-memory writes) were +already replayable for free. diff --git a/packages/python-sdk/e2b/api/__init__.py b/packages/python-sdk/e2b/api/__init__.py index 95ca7bd146..2390f3b639 100644 --- a/packages/python-sdk/e2b/api/__init__.py +++ b/packages/python-sdk/e2b/api/__init__.py @@ -86,15 +86,17 @@ async def on_response(response: Response) -> None: pool_idle_timeout = float(os.getenv("E2B_KEEPALIVE_EXPIRY") or "300") pool_max_idle_per_host = int(os.getenv("E2B_MAX_KEEPALIVE_CONNECTIONS") or "20") + # The retry policy the shared transports declare per request (see # `ConnectionRetryTransport` in client_sync/client_async). # # pyqwest's retry middleware keeps a request replayable, and for a body that # isn't already `bytes` it does that by growing a copy of it in memory as it is # sent. A streamed upload therefore went to the wire incrementally *and* was -# mirrored whole in RAM, so peak memory scaled with file size for -# `files.write`, `volume.write_file` and template context uploads even though -# nothing in the SDK buffers (SDK-332). +# mirrored whole in RAM, so peak memory scaled with file size for `files.write` +# of a file-like object and for `volume.write_file`, even though nothing in the +# SDK buffers (SDK-332). Template context uploads were never affected: they +# build their own transport with no retry middleware in it. # # `RetryMode.UNBUFFERED` drops the copy: a streamed body is replayed only while # nothing has been read from it, which is all a connect-only policy ever needs, @@ -109,8 +111,16 @@ async def on_response(response: Response) -> None: # so a connect error surfaced with the body already started (measured: the one # chunk of a unary RPC body, up to a whole small-chunked upload) and neither # buffered nor rewindable. -_retry_mode = getattr(retry_middleware, "RetryMode", None) -unbuffered_retries: Any = True if _retry_mode is None else _retry_mode.UNBUFFERED +# +# Typed `Any` rather than the union it is: the pinned release declares +# `should_retry_request` as `-> bool`, so annotating the override with a union +# that admits `RetryMode` fails the Liskov check until the pin moves. +def _resolve_retry_request_policy(middleware: Any) -> Any: + retry_mode = getattr(middleware, "RetryMode", None) + return True if retry_mode is None else retry_mode.UNBUFFERED + + +_retry_request_policy: Any = _resolve_retry_request_policy(retry_middleware) class ProxyConfig(NamedTuple): diff --git a/packages/python-sdk/e2b/api/client_async/__init__.py b/packages/python-sdk/e2b/api/client_async/__init__.py index 2fc175c660..58700b1a67 100644 --- a/packages/python-sdk/e2b/api/client_async/__init__.py +++ b/packages/python-sdk/e2b/api/client_async/__init__.py @@ -10,12 +10,12 @@ from e2b.api import ( AsyncApiClient, ProxyConfig, + _retry_request_policy, connection_retries, make_async_logging_event_hooks, pool_idle_timeout, pool_max_idle_per_host, proxy_to_config, - unbuffered_retries, ) from e2b.connection_config import READ_TIMEOUT, ConnectionConfig @@ -35,11 +35,13 @@ class ConnectionRetryTransport(RetryTransport): default policy would otherwise also retry I/O errors and 429/5xx responses for idempotent methods. - Streamed bodies are retried without being buffered, so an upload is not - mirrored in memory as it is sent (see ``unbuffered_retries``).""" + On a pyqwest whose retry middleware exposes ``RetryMode``, streamed bodies + are retried without being buffered, so an upload is not mirrored in memory + as it is sent; earlier releases keep the buffered replay they ship (see + ``_retry_request_policy``).""" def should_retry_request(self, request: Request) -> Any: - return unbuffered_retries + return _retry_request_policy def should_retry_response( self, request: Request, response: Union[Response, Exception] diff --git a/packages/python-sdk/e2b/api/client_sync/__init__.py b/packages/python-sdk/e2b/api/client_sync/__init__.py index f004a01d00..775b1c00d3 100644 --- a/packages/python-sdk/e2b/api/client_sync/__init__.py +++ b/packages/python-sdk/e2b/api/client_sync/__init__.py @@ -10,12 +10,12 @@ from e2b.api import ( ApiClient, ProxyConfig, + _retry_request_policy, connection_retries, make_logging_event_hooks, pool_idle_timeout, pool_max_idle_per_host, proxy_to_config, - unbuffered_retries, ) from e2b.connection_config import READ_TIMEOUT, ConnectionConfig @@ -35,11 +35,13 @@ class ConnectionRetryTransport(SyncRetryTransport): default policy would otherwise also retry I/O errors and 429/5xx responses for idempotent methods. - Streamed bodies are retried without being buffered, so an upload is not - mirrored in memory as it is sent (see ``unbuffered_retries``).""" + On a pyqwest whose retry middleware exposes ``RetryMode``, streamed bodies + are retried without being buffered, so an upload is not mirrored in memory + as it is sent; earlier releases keep the buffered replay they ship (see + ``_retry_request_policy``).""" def should_retry_request(self, request: SyncRequest) -> Any: - return unbuffered_retries + return _retry_request_policy def should_retry_response( self, request: SyncRequest, response: Union[SyncResponse, Exception] diff --git a/packages/python-sdk/tests/test_retry_streamed_bodies.py b/packages/python-sdk/tests/test_retry_streamed_bodies.py index ee402f859b..550615d55c 100644 --- a/packages/python-sdk/tests/test_retry_streamed_bodies.py +++ b/packages/python-sdk/tests/test_retry_streamed_bodies.py @@ -4,27 +4,34 @@ pyqwest's retry middleware keeps a request replayable, and for a body that is not already ``bytes`` it does that by growing a copy of the body as it is sent. An upload therefore reached the wire incrementally *and* was held whole in RAM, -so peak memory scaled with file size for ``files.write``, ``volume.write_file`` -and template context uploads. The transports now declare -``unbuffered_retries``, which drops the copy and replays a streamed body only -while nothing has been read from it — all the connect-only policy needs, since -pyqwest starts reading a body when hyper first writes it and a connect error -means it never got that far. +so peak memory scaled with file size for ``files.write`` of a file-like object +and for ``volume.write_file`` (template context uploads build their own +transport with no retry middleware, so they were never affected). The transports +now declare ``_retry_request_policy``, which drops the copy and replays a +streamed body only while nothing has been read from it — all the connect-only +policy needs, since pyqwest starts reading a body when hyper first writes it and +a connect error means it never got that far. ``bytes`` bodies (every unary RPC payload, every in-memory write) are replayable as they are, so their retries are unaffected; the tests for those live in ``test_envd_retry_transport.py``. -The tests that pin the unbuffered half are skipped on a pyqwest whose retry -middleware has no ``RetryMode``: there the SDK keeps that release's buffered -behavior, because going unbuffered without pyqwest's deferred body reads would -trade the copy for the connect retries. +Which policy is in force depends on the installed pyqwest, so the tests come in +two halves keyed on whether its retry middleware exposes ``RetryMode``: +``unbuffered_only`` pins the unbuffered behavior, ``buffered_only`` pins the +buffered replay older releases ship, and one always-running test reports which +half is live so a pyqwest bump cannot switch them over unnoticed. The rest — +the policy resolution, the overrides being declared at all, and what the httpx +adapter hands the middleware — run on every release. """ +import io import socket import tracemalloc +from types import SimpleNamespace from typing import AsyncIterator, Iterator, List +import httpx import pytest from pyqwest import ( HTTPTransport, @@ -34,23 +41,34 @@ SyncRequest, SyncResponse, ) +from pyqwest.httpx import AsyncPyqwestTransport, PyqwestTransport +from pyqwest.middleware import retry as retry_middleware -from e2b.api import unbuffered_retries +from e2b.api import _resolve_retry_request_policy, _retry_request_policy from e2b.api.client_async import ConnectionRetryTransport from e2b.api.client_sync import ( ConnectionRetryTransport as SyncConnectionRetryTransport, ) +PYQWEST_RETRY_MODE = getattr(retry_middleware, "RetryMode", None) +"""``pyqwest.middleware.retry.RetryMode``, or ``None`` on a release without it.""" + CHUNK_SIZE = 256 * 1024 CHUNK_COUNT = 64 BODY_SIZE = CHUNK_SIZE * CHUNK_COUNT """16 MiB in 256 KiB chunks: large enough that a mirrored copy dwarfs the per-chunk allocations, small enough to stay quick.""" +# Keyed on the capability rather than on `_retry_request_policy`'s value, so a +# wrong value fails a test instead of quietly turning it into a skip. unbuffered_only = pytest.mark.skipif( - unbuffered_retries is True, + PYQWEST_RETRY_MODE is None, reason="pyqwest without RetryMode buffers streamed bodies to replay them", ) +buffered_only = pytest.mark.skipif( + PYQWEST_RETRY_MODE is not None, + reason="pyqwest with RetryMode replays streamed bodies only while unread", +) MAX_TRACED_PEAK = BODY_SIZE // 4 @@ -93,19 +111,24 @@ def _refused_port() -> int: class SendingTransport: """Inner async transport that reads the request body the way reqwest does, after failing the first ``failures`` attempts the way a refused connect - does — before the body was touched.""" + does — before the body was touched, unless ``fail_mid_body``. + + ``streamed`` records, per attempt, whether the body arrived as something + other than ``bytes``: that is the shape the middleware replays by copying.""" def __init__(self, failures: int = 0, fail_mid_body: bool = False): self.failures = failures self.fail_mid_body = fail_mid_body self.attempts = 0 self.received = 0 + self.streamed: List[bool] = [] async def execute(self, request: Request) -> Response: self.attempts += 1 if self.attempts <= self.failures and not self.fail_mid_body: raise ConnectionError("tcp connect error") content = request.content + self.streamed.append(not isinstance(content, bytes)) if isinstance(content, bytes): self.received += len(content) else: @@ -122,12 +145,14 @@ def __init__(self, failures: int = 0, fail_mid_body: bool = False): self.fail_mid_body = fail_mid_body self.attempts = 0 self.received = 0 + self.streamed: List[bool] = [] def execute_sync(self, request: SyncRequest) -> SyncResponse: self.attempts += 1 if self.attempts <= self.failures and not self.fail_mid_body: raise ConnectionError("tcp connect error") content = request.content + self.streamed.append(not isinstance(content, bytes)) if isinstance(content, bytes): self.received += len(content) else: @@ -138,6 +163,46 @@ def execute_sync(self, request: SyncRequest) -> SyncResponse: return SyncResponse(status=200, content=b"ok") +# What the transports declare, on every pyqwest release: the resolution of the +# policy, and the fact that they override the hook at all. Neither depends on +# which policy is in force, and both are invisible to the behavioral tests +# below on a release where the inherited default happens to match. + + +def test_policy_is_pyqwests_unbuffered_mode_when_the_release_has_it(): + mode = SimpleNamespace(UNBUFFERED=object()) + assert _resolve_retry_request_policy(SimpleNamespace(RetryMode=mode)) is ( + mode.UNBUFFERED + ) + + +def test_policy_falls_back_to_buffered_retries_without_retry_mode(): + # `True` is what the middleware understood before the mode existed. + assert _resolve_retry_request_policy(SimpleNamespace()) is True + + +def test_policy_matches_the_installed_pyqwest(): + # The two halves of this module are gated on the same capability, so this + # says out loud which half is live: a pyqwest bump flips it rather than + # quietly switching six tests from skipped to running. + if PYQWEST_RETRY_MODE is None: + assert _retry_request_policy is True + else: + assert _retry_request_policy is PYQWEST_RETRY_MODE.UNBUFFERED + + +@pytest.mark.parametrize( + "transport", + [SyncConnectionRetryTransport, ConnectionRetryTransport], + ids=["sync", "async"], +) +def test_transports_declare_the_retry_policy(transport): + # Inheriting pyqwest's default would restore the buffered replay silently, + # so pin the override itself and not just its result. + assert "should_retry_request" in transport.__dict__ + assert transport(None).should_retry_request(None) is _retry_request_policy + + @unbuffered_only async def test_async_streamed_body_is_not_mirrored_in_memory(): inner = SendingTransport() @@ -222,6 +287,37 @@ def test_sync_streamed_body_is_not_replayed_once_it_was_read(): assert inner.received == CHUNK_SIZE +# The other half: on a pyqwest without `RetryMode` the copy the middleware +# grows is what makes the same failure retryable, which is why the fallback +# keeps it rather than going unbuffered without pyqwest's deferred body reads. + + +@buffered_only +async def test_async_buffered_streamed_body_is_replayed_from_the_copy(): + inner = SendingTransport(failures=1, fail_mid_body=True) + pulled: List[int] = [] + request = Request("PUT", "http://sandbox.test/upload", content=_achunks(pulled)) + response = await _retrying(inner).execute(request) + assert response.status == 200 + assert inner.attempts == 2 + # The first attempt sent one chunk; the replay resent it from the copy and + # then streamed the rest, so the chunk went out twice. + assert inner.received == CHUNK_SIZE + BODY_SIZE + assert sum(pulled) == BODY_SIZE + + +@buffered_only +def test_sync_buffered_streamed_body_is_replayed_from_the_copy(): + inner = SendingSyncTransport(failures=1, fail_mid_body=True) + pulled: List[int] = [] + request = SyncRequest("PUT", "http://sandbox.test/upload", content=_chunks(pulled)) + response = _retrying_sync(inner).execute_sync(request) + assert response.status == 200 + assert inner.attempts == 2 + assert inner.received == CHUNK_SIZE + BODY_SIZE + assert sum(pulled) == BODY_SIZE + + # Through a real pyqwest transport: the retries above are only safe because a # connect failure leaves the body unread, which is what lets the middleware # hand the same iterator to the next attempt. @@ -247,3 +343,87 @@ def test_sync_refused_connect_leaves_a_streamed_body_unread(): SyncRequest("PUT", url, content=_chunks(pulled)) ) assert pulled == [] + + +# Through the httpx adapter, the way every affected call site reaches the +# middleware: `files.write` passes a file-like object as `content` or as a +# multipart part, `volume.write_file` passes `IO[bytes]` or an iterator. The +# adapter is what decides whether the middleware sees a replayable `bytes` body +# or the streamed one this policy is about, so the join is worth pinning — the +# transport-level tests above would all pass on a body the adapter had already +# flattened. + +UPLOAD_URL = "http://sandbox.test/files?path=/home/user/upload.bin" + + +def _put_sync(client: httpx.Client, shape: str) -> httpx.Response: + """``files.write``'s two upload shapes: the body as ``content``, and the + same file-like object as a multipart part.""" + if shape == "content": + return client.put(UPLOAD_URL, content=_chunks([])) + return client.put( + UPLOAD_URL, files={"file": ("upload.bin", io.BytesIO(b"x" * BODY_SIZE))} + ) + + +async def _put_async(client: httpx.AsyncClient, shape: str) -> httpx.Response: + if shape == "content": + return await client.put(UPLOAD_URL, content=_achunks([])) + return await client.put( + UPLOAD_URL, files={"file": ("upload.bin", io.BytesIO(b"x" * BODY_SIZE))} + ) + + +@pytest.mark.parametrize("shape", ["content", "multipart"]) +def test_sync_httpx_hands_a_streamed_upload_to_the_middleware(shape): + inner = SendingSyncTransport() + with httpx.Client(transport=PyqwestTransport(_retrying_sync(inner))) as client: + response = _put_sync(client, shape) + assert response.status_code == 200 + assert inner.streamed == [True] + # Multipart framing adds a little to the part itself. + assert inner.received >= BODY_SIZE + + +@pytest.mark.parametrize("shape", ["content", "multipart"]) +async def test_async_httpx_hands_a_streamed_upload_to_the_middleware(shape): + inner = SendingTransport() + async with httpx.AsyncClient( + transport=AsyncPyqwestTransport(_retrying(inner)) + ) as client: + response = await _put_async(client, shape) + assert response.status_code == 200 + assert inner.streamed == [True] + assert inner.received >= BODY_SIZE + + +@unbuffered_only +def test_sync_httpx_upload_is_not_mirrored_in_memory(): + inner = SendingSyncTransport() + with httpx.Client(transport=PyqwestTransport(_retrying_sync(inner))) as client: + tracemalloc.start() + try: + response = client.put(UPLOAD_URL, content=_chunks([])) + peak = tracemalloc.get_traced_memory()[1] + finally: + tracemalloc.stop() + assert response.status_code == 200 + assert inner.received == BODY_SIZE + assert peak < MAX_TRACED_PEAK, peak + + +@unbuffered_only +async def test_async_httpx_upload_is_not_mirrored_in_memory(): + inner = SendingTransport() + async with httpx.AsyncClient( + transport=AsyncPyqwestTransport(_retrying(inner)) + ) as client: + tracemalloc.start() + try: + response = await client.put(UPLOAD_URL, content=_achunks([])) + peak = tracemalloc.get_traced_memory()[1] + finally: + tracemalloc.stop() + assert response.status_code == 200 + assert inner.received == BODY_SIZE + assert peak < MAX_TRACED_PEAK, peak