diff --git a/.changeset/python-unbuffered-retry-streamed-bodies.md b/.changeset/python-unbuffered-retry-streamed-bodies.md new file mode 100644 index 0000000000..e69e834ae8 --- /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, 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 2db0fa80c6..2390f3b639 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 @@ -86,6 +87,42 @@ async def on_response(response: Response) -> None: 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` +# 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, +# 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. +# +# 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): """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..58700b1a67 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 @@ -10,6 +10,7 @@ from e2b.api import ( AsyncApiClient, ProxyConfig, + _retry_request_policy, connection_retries, make_async_logging_event_hooks, pool_idle_timeout, @@ -32,7 +33,15 @@ 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. + + 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 _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 73c279acec..775b1c00d3 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 @@ -10,6 +10,7 @@ from e2b.api import ( ApiClient, ProxyConfig, + _retry_request_policy, connection_retries, make_logging_event_hooks, pool_idle_timeout, @@ -32,7 +33,15 @@ 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. + + 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 _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 new file mode 100644 index 0000000000..550615d55c --- /dev/null +++ b/packages/python-sdk/tests/test_retry_streamed_bodies.py @@ -0,0 +1,429 @@ +"""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`` 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``. + +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, + Request, + Response, + SyncHTTPTransport, + SyncRequest, + SyncResponse, +) +from pyqwest.httpx import AsyncPyqwestTransport, PyqwestTransport +from pyqwest.middleware import retry as retry_middleware + +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( + 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 + + +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, 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: + 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 + 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: + 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") + + +# 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() + 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, 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(): + 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 + assert sum(pulled) == 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 + assert inner.received == CHUNK_SIZE + + +@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 + 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. + + +@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 == [] + + +# 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