From eb45e231cc9583b211f5dab6c3ecd84f805c43d5 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Fri, 3 Jul 2026 00:02:23 -0700 Subject: [PATCH] fix(vertex_ai/files): stream batch uploads via a GCS resumable chunked session Batch JSONL uploads were being sent as a single uploadType=media POST with a streamed body. This puts them back on a GCS resumable session: open the session with a POST, then PUT the transformed body in bounded 8 MiB chunks with Content-Range, so a large upload is never held in memory in full The async path assembles each chunk on a worker thread so the CPU-bound OpenAI to Vertex row transform never blocks the event loop between chunk PUTs --- litellm/llms/custom_httpx/llm_http_handler.py | 256 +++++++++++++----- .../llms/vertex_ai/files/transformation.py | 23 +- litellm/types/files.py | 17 +- .../test_openai_batches_and_files.py | 53 ++-- .../test_vertex_ai_binary_file_upload.py | 12 +- .../files/test_vertex_ai_files_streaming.py | 244 ++++++++++++----- 6 files changed, 432 insertions(+), 173 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 3c10239f868..ce50b8f6215 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -89,7 +89,7 @@ from litellm.types.containers.main import ( ContainerObject, DeleteContainerResult, ) -from litellm.types.files import StreamingMediaUploadConfig, TwoStepFileUploadConfig +from litellm.types.files import ResumableChunkedUploadConfig, TwoStepFileUploadConfig from litellm.types.integrations.custom_logger import ( AgenticLoopPlan, AgenticLoopRequestPatch, @@ -3300,15 +3300,13 @@ class BaseLLMHTTPHandler: data=presigned_request["data"], timeout=timeout, ) - elif isinstance(transformed_request, dict) and "streaming_media_upload" in transformed_request: - media_cfg = cast(StreamingMediaUploadConfig, transformed_request["streaming_media_upload"]) + elif isinstance(transformed_request, dict) and "resumable_chunked_upload" in transformed_request: try: - upload_response = self._upload_media( + upload_response = self._resumable_chunked_upload( client=sync_httpx_client, - url=api_base, + initiate_url=api_base, base_headers=headers, - body_stream=cast(BaseFileUploadStream, media_cfg["body_stream"]), - content_type=media_cfg.get("content_type") or "application/octet-stream", + config=cast(ResumableChunkedUploadConfig, transformed_request["resumable_chunked_upload"]), timeout=timeout, ) except Exception as e: @@ -3393,8 +3391,8 @@ class BaseLLMHTTPHandler: # huge) upload payload; logging deep-copies additional_args, so log # a placeholder instead of re-materializing the payload. "complete_input_dict": ( - "" - if isinstance(transformed_request, dict) and "streaming_media_upload" in transformed_request + "" + if isinstance(transformed_request, dict) and "resumable_chunked_upload" in transformed_request else transformed_request ), "api_base": api_base, @@ -3462,15 +3460,13 @@ class BaseLLMHTTPHandler: data=presigned_request["data"], timeout=timeout, ) - elif isinstance(transformed_request, dict) and "streaming_media_upload" in transformed_request: - media_cfg = cast(StreamingMediaUploadConfig, transformed_request["streaming_media_upload"]) + elif isinstance(transformed_request, dict) and "resumable_chunked_upload" in transformed_request: try: - upload_response = await self._aupload_media( + upload_response = await self._aresumable_chunked_upload( client=async_httpx_client, - url=api_base, + initiate_url=api_base, base_headers=headers, - body_stream=cast(BaseFileUploadStream, media_cfg["body_stream"]), - content_type=media_cfg.get("content_type") or "application/octet-stream", + config=cast(ResumableChunkedUploadConfig, transformed_request["resumable_chunked_upload"]), timeout=timeout, ) except Exception as e: @@ -3515,81 +3511,211 @@ class BaseLLMHTTPHandler: litellm_params=litellm_params, ) - # The fine-grained transform stream (one piece per JSONL row) is regrouped - # into blocks of this size before upload, so the request yields a manageable - # number of chunks; never more than one block is buffered. - _MEDIA_UPLOAD_BLOCK_SIZE = 4 * 1024 * 1024 + # 8 MiB; a 256 KiB multiple, which GCS requires for every non-final chunk. + _RESUMABLE_CHUNK_SIZE = 8 * 1024 * 1024 @staticmethod - def _iter_in_blocks(byte_iter: Iterator[bytes], block_size: int) -> Iterator[bytes]: + def _iter_resumable_chunks(byte_iter: Iterator[bytes], chunk_size: int) -> Iterator[bytes]: + """Regroup a byte stream into ``chunk_size`` pieces followed by a final + piece of whatever remains (possibly empty). Every piece but the last is + exactly ``chunk_size`` bytes, so the caller can keep that a 256 KiB + multiple and never buffers more than one chunk. + """ buf = bytearray() for piece in byte_iter: buf.extend(piece) - while len(buf) >= block_size: - yield bytes(buf[:block_size]) - del buf[:block_size] - if buf: - yield bytes(buf) + while len(buf) >= chunk_size: + yield bytes(buf[:chunk_size]) + del buf[:chunk_size] + yield bytes(buf) - def _check_media_upload_response(self, resp: httpx.Response) -> None: - if resp.status_code not in (200, 201): - resp.raise_for_status() - raise ValueError(f"media upload: unexpected status {resp.status_code}") + @staticmethod + def _resumable_content_range(offset: int, data_len: int, is_final: bool) -> str: + if not is_final: + return f"bytes {offset}-{offset + data_len - 1}/*" + total = offset + data_len + if data_len == 0: + return f"bytes */{total}" + return f"bytes {offset}-{total - 1}/{total}" - def _upload_media( + @staticmethod + def _build_resumable_request( + httpx_client: Union[httpx.Client, httpx.AsyncClient], + method: str, + url: str, + headers: Dict[str, str], + content: bytes, + timeout: Optional[Union[float, httpx.Timeout]], + ) -> httpx.Request: + # Passing timeout=None to httpx means "no timeout"; use the client-default + # sentinel instead so an unset caller timeout keeps the client's default. + effective_timeout = timeout if timeout is not None else httpx.USE_CLIENT_DEFAULT + return httpx_client.build_request(method, url, headers=headers, content=content, timeout=effective_timeout) + + def _resumable_chunked_upload( self, *, client: HTTPHandler, - url: str, + initiate_url: str, base_headers: Dict[str, str], - body_stream: BaseFileUploadStream, - content_type: str, + config: ResumableChunkedUploadConfig, timeout: Optional[Union[float, httpx.Timeout]], ) -> httpx.Response: - headers = {**base_headers, "Content-Type": content_type} - kwargs: Dict[str, Any] = { - "headers": headers, - "content": self._iter_in_blocks(body_stream.iter_bytes(), self._MEDIA_UPLOAD_BLOCK_SIZE), + """Open a GCS resumable session, then PUT the body in bounded chunks so a + large upload is never held in memory in full.""" + stream = cast(BaseFileUploadStream, config["body_stream"]) + chunk_size = config.get("chunk_size", self._RESUMABLE_CHUNK_SIZE) + session_url_header = config.get("session_url_header", "location") + initiate_headers: Dict[str, str] = config.get("initiate_headers") or {} + httpx_client = client.client + + init_headers = {**base_headers, **initiate_headers} + init_req = self._build_resumable_request(httpx_client, "POST", initiate_url, init_headers, b"", timeout) + init_resp = httpx_client.send(init_req, follow_redirects=False) + init_resp.read() + if init_resp.status_code not in (200, 201): + # raise_for_status only raises on 4xx/5xx; surface an unexpected 2xx/3xx + # (e.g. 202) here instead of falling through to a misleading missing-URL error. + init_resp.raise_for_status() + raise ValueError(f"resumable upload: unexpected session-init status {init_resp.status_code}") + session_url = cast(str, init_resp.headers.get(session_url_header) or "") + if not session_url: + raise ValueError(f"resumable upload: no session URL in '{session_url_header}' header") + + offset = 0 + pending: Optional[bytes] = None + for chunk in self._iter_resumable_chunks(stream.iter_bytes(), chunk_size): + if pending is not None: + self._send_resumable_chunk( + httpx_client, + session_url, + base_headers, + pending, + offset, + is_final=False, + timeout=timeout, + ) + offset += len(pending) + pending = chunk + return self._send_resumable_chunk( + httpx_client, + session_url, + base_headers, + pending or b"", + offset, + is_final=True, + timeout=timeout, + ) + + def _send_resumable_chunk( + self, + httpx_client: httpx.Client, + url: str, + base_headers: Dict[str, str], + data: bytes, + offset: int, + *, + is_final: bool, + timeout: Optional[Union[float, httpx.Timeout]], + ) -> httpx.Response: + headers = { + **base_headers, + "Content-Range": self._resumable_content_range(offset, len(data), is_final), } - if timeout is not None: - kwargs["timeout"] = timeout - resp = client.client.post(url, **kwargs) - self._check_media_upload_response(resp) + req = self._build_resumable_request(httpx_client, "PUT", url, headers, data, timeout) + resp = httpx_client.send(req, follow_redirects=False) + resp.read() + if resp.status_code not in ((200, 201) if is_final else (308,)): + # 4xx/5xx raise here; the ValueError catches an unexpected success + # status (e.g. a 200 where the protocol expects a 308 between chunks). + resp.raise_for_status() + raise ValueError(f"resumable upload: unexpected status {resp.status_code}") return resp - async def _aupload_media( + async def _aresumable_chunked_upload( self, *, client: AsyncHTTPHandler, - url: str, + initiate_url: str, base_headers: Dict[str, str], - body_stream: BaseFileUploadStream, - content_type: str, + config: ResumableChunkedUploadConfig, timeout: Optional[Union[float, httpx.Timeout]], ) -> httpx.Response: - """Stream the transformed body straight to a single media upload. Each - block is produced on a worker thread (the transform never runs on the - event loop) and sent with chunked transfer-encoding, so the body is - neither buffered in memory nor staged to disk, and the upload is one - continuous request rather than the many sequential round-trips of the - resumable path that overran client/LB timeouts.""" - headers = {**base_headers, "Content-Type": content_type} - block_iter = iter(self._iter_in_blocks(body_stream.iter_bytes(), self._MEDIA_UPLOAD_BLOCK_SIZE)) + stream = cast(BaseFileUploadStream, config["body_stream"]) + chunk_size = config.get("chunk_size", self._RESUMABLE_CHUNK_SIZE) + session_url_header = config.get("session_url_header", "location") + initiate_headers: Dict[str, str] = config.get("initiate_headers") or {} + httpx_client = client.client + + init_headers = {**base_headers, **initiate_headers} + init_req = self._build_resumable_request(httpx_client, "POST", initiate_url, init_headers, b"", timeout) + init_resp = await httpx_client.send(init_req, follow_redirects=False) + await init_resp.aread() + if init_resp.status_code not in (200, 201): + # raise_for_status only raises on 4xx/5xx; surface an unexpected 2xx/3xx + # (e.g. 202) here instead of falling through to a misleading missing-URL error. + init_resp.raise_for_status() + raise ValueError(f"resumable upload: unexpected session-init status {init_resp.status_code}") + session_url = cast(str, init_resp.headers.get(session_url_header) or "") + if not session_url: + raise ValueError(f"resumable upload: no session URL in '{session_url_header}' header") + + # The OpenAI->Vertex transform that produces each chunk is CPU-bound; + # pull chunks on a worker thread so assembling the next one never blocks + # the event loop while the previous chunk's PUT is in flight. + chunk_iter = iter(self._iter_resumable_chunks(stream.iter_bytes(), chunk_size)) done = object() + offset = 0 + pending: Optional[bytes] = None + while True: + chunk = await asyncio.to_thread(next, chunk_iter, done) + if chunk is done: + break + if pending is not None: + await self._asend_resumable_chunk( + httpx_client, + session_url, + base_headers, + pending, + offset, + is_final=False, + timeout=timeout, + ) + offset += len(pending) + pending = cast(bytes, chunk) + return await self._asend_resumable_chunk( + httpx_client, + session_url, + base_headers, + pending or b"", + offset, + is_final=True, + timeout=timeout, + ) - async def _abody() -> AsyncIterator[bytes]: - while True: - block = await asyncio.to_thread(next, block_iter, done) - if block is done: - break - yield cast(bytes, block) - - kwargs: Dict[str, Any] = {"headers": headers, "content": _abody()} - if timeout is not None: - kwargs["timeout"] = timeout - resp = await client.client.post(url, **kwargs) + async def _asend_resumable_chunk( + self, + httpx_client: httpx.AsyncClient, + url: str, + base_headers: Dict[str, str], + data: bytes, + offset: int, + *, + is_final: bool, + timeout: Optional[Union[float, httpx.Timeout]], + ) -> httpx.Response: + headers = { + **base_headers, + "Content-Range": self._resumable_content_range(offset, len(data), is_final), + } + req = self._build_resumable_request(httpx_client, "PUT", url, headers, data, timeout) + resp = await httpx_client.send(req, follow_redirects=False) await resp.aread() - self._check_media_upload_response(resp) + if resp.status_code not in ((200, 201) if is_final else (308,)): + # 4xx/5xx raise here; the ValueError catches an unexpected success + # status (e.g. a 200 where the protocol expects a 308 between chunks). + resp.raise_for_status() + raise ValueError(f"resumable upload: unexpected status {resp.status_code}") return resp def create_batch( diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index dd877b52eb8..fcb7617d97e 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -60,7 +60,7 @@ from litellm.types.llms.openai import ( OpenAIFileObject, PathLike, ) -from litellm.types.files import StreamingMediaUploadConfig +from litellm.types.files import ResumableChunkedUploadConfig from litellm.types.llms.vertex_ai import GcsBucketResponse from litellm.types.utils import LlmProviders, ModelResponse @@ -380,11 +380,19 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): raise ValueError("file is required") if purpose is None: raise ValueError("purpose is required") + _, content_type = extract_file_metadata(file_data) object_name = self.get_object_name(file_data, purpose) if object_prefix: object_name = f"{object_prefix}/{object_name}" encoded_object_name = encode_gcs_object_name_for_url(object_name) - endpoint = f"upload/storage/v1/b/{bucket_name}/o?uploadType=media&name={encoded_object_name}" + # Batch jsonl is streamed via a resumable session (bounded memory on + # large uploads); everything else is a single simple-media upload. + upload_type = ( + "resumable" + if FilesAPIUtils.is_batch_jsonl_request(create_file_data=data, content_type=content_type) + else "media" + ) + endpoint = f"upload/storage/v1/b/{bucket_name}/o?uploadType={upload_type}&name={encoded_object_name}" api_base = api_base or "https://storage.googleapis.com" if not api_base: raise ValueError("api_base is required") @@ -434,9 +442,8 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): """ 2 Cases: 1. Handle basic file upload - 2. Handle batch file upload (.jsonl), staged to a temp file and uploaded - in a single media request so large uploads stay memory-bounded without - the per-chunk round-trips of a resumable session. + 2. Handle batch file upload (.jsonl), streamed to a GCS resumable + session so large uploads stay memory-bounded. """ file_data = create_file_data.get("file") if file_data is None: @@ -448,12 +455,14 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): content_type=content_type, ): return { - "streaming_media_upload": StreamingMediaUploadConfig( + "resumable_chunked_upload": ResumableChunkedUploadConfig( body_stream=_OpenAIToVertexBatchUploadStream( file_data, self._map_openai_to_vertex_params, ), - content_type="application/json", + initiate_headers={ + "X-Upload-Content-Type": "application/json", + }, ) } diff --git a/litellm/types/files.py b/litellm/types/files.py index 97c4dbb3abc..1b2d7e30f1f 100644 --- a/litellm/types/files.py +++ b/litellm/types/files.py @@ -323,18 +323,19 @@ class TwoStepFileUploadConfig(TypedDict, total=False): upload_url_key: str -class StreamingMediaUploadConfig(TypedDict, total=False): - """Drives a memory-bounded single-request upload (GCS simple/media upload). +class ResumableChunkedUploadConfig(TypedDict, total=False): + """Drives a memory-bounded resumable upload (GCS JSON API). - The handler stages ``body_stream`` to a temp file off the event loop (so peak - memory stays bounded), then PUTs/POSTs it in one request with a known - Content-Length. Unlike a resumable chunked upload this incurs no per-chunk - round-trips, so a multi-GB upload finishes in one continuous transfer instead - of hundreds of sequential PUTs that overrun client/LB timeouts. + The handler POSTs to the upload URL to open a session, reads the session URI + from ``session_url_header``, then PUTs ``body_stream`` to that URI in + ``chunk_size``-byte chunks (a 256 KiB multiple) using Content-Range, so the + payload is never buffered in full and the transfer is resumable. ``body_stream`` is a ``BaseFileUploadStream``; it is typed ``Any`` here to avoid importing the llms layer into types. """ body_stream: Required[Any] - content_type: str + chunk_size: int + session_url_header: str + initiate_headers: Dict[str, str] diff --git a/tests/batches_tests/test_openai_batches_and_files.py b/tests/batches_tests/test_openai_batches_and_files.py index 0a49b3d77d1..8c984ca77b7 100644 --- a/tests/batches_tests/test_openai_batches_and_files.py +++ b/tests/batches_tests/test_openai_batches_and_files.py @@ -513,11 +513,28 @@ async def test_avertex_batch_prediction(monkeypatch): mock_response.status_code = 200 return mock_response - # Batch jsonl creation now stages the body to a temp file and issues a single - # uploadType=media POST against the raw httpx.AsyncClient (client.client) inside - # _astage_and_upload_media, not AsyncHTTPHandler.post. Patch that raw POST so the - # real staging/upload + response transform run while the GCS object response is - # mocked; AsyncHTTPHandler.post still handles the batch-prediction call. + # Batch jsonl creation streams the body to a GCS resumable session on the raw + # httpx.AsyncClient (client.client): one POST opens the session (URI in the + # Location header), then the body is PUT in chunks. Those go through + # AsyncClient.send, not AsyncHTTPHandler.post, so patch send to run the real + # upload + response transform against a mocked GCS session; AsyncHTTPHandler.post + # still handles the batch-prediction call. + gcs_session_url = ( + "https://storage.googleapis.com/upload/storage/v1/b/litellm-local/o?uploadType=resumable&upload_id=SID" + ) + + async def mock_gcs_send(request, **kwargs): + if request.method == "POST": + return httpx.Response(200, headers={"location": gcs_session_url}, request=request) + # A non-final chunk carries an unknown total ("bytes X-Y/*") and must get a + # 308; only the final chunk ("bytes X-Y/TOTAL") returns the object resource. + # Returning 200 for every PUT would mask the handler's 308 requirement and + # break the moment a payload spans more than one chunk. + content_range = request.headers["content-range"] + if content_range.rsplit("/", 1)[-1] == "*": + return httpx.Response(308, headers={"range": "bytes=0-*"}, request=request) + return httpx.Response(200, json=mock_file_response, request=request) + with ( patch( "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", @@ -525,13 +542,9 @@ async def test_avertex_batch_prediction(monkeypatch): ), patch.object( httpx.AsyncClient, - "post", + "send", new_callable=AsyncMock, - return_value=httpx.Response( - 200, - json=mock_file_response, - request=httpx.Request("POST", "https://storage.googleapis.com/upload"), - ), + side_effect=mock_gcs_send, ) as mock_gcs_upload, ): litellm.set_verbose = True @@ -553,14 +566,16 @@ async def test_avertex_batch_prediction(monkeypatch): == "gs://litellm-local/litellm-vertex-files/publishers/google/models/gemini-1.5-flash-001/5f7b99ad-9203-4430-98bf-3b45451af4cb" ) - mock_gcs_upload.assert_awaited_once() - upload_url = str(mock_gcs_upload.call_args.args[0]) - assert "uploadType=media" in upload_url - assert "/b/litellm-local/o" in upload_url - assert ( - mock_gcs_upload.call_args.kwargs["headers"]["Content-Type"] - == "application/json" - ) + # Session-open POST then at least one chunk PUT. + assert mock_gcs_upload.await_count >= 2 + init_request = mock_gcs_upload.call_args_list[0].args[0] + assert init_request.method == "POST" + assert "uploadType=resumable" in str(init_request.url) + assert "/b/litellm-local/o" in str(init_request.url) + assert init_request.headers["X-Upload-Content-Type"] == "application/json" + chunk_request = mock_gcs_upload.call_args_list[1].args[0] + assert chunk_request.method == "PUT" + assert str(chunk_request.url) == gcs_session_url # Create batch create_batch_response = await litellm.acreate_batch( diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py index d2ee9d7d659..0e76d007a08 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py @@ -133,9 +133,9 @@ class TestVertexAIBinaryFileUpload: @pytest.mark.asyncio async def test_jsonl_file_upload_returns_streaming_body(self): """ - Test that JSONL batch files are transformed into a streaming-media config + Test that JSONL batch files are transformed into a resumable-upload config carrying a streaming body (not a buffered bytes payload), so the handler - can stage the upload to a temp file and send it in one media request. + can chunk the upload to a GCS resumable session. """ # Create mock JSONL content mock_jsonl_content = ( @@ -158,11 +158,11 @@ class TestVertexAIBinaryFileUpload: litellm_params={}, ) - assert isinstance(transformed_request, dict) and "streaming_media_upload" in transformed_request, ( - f"Expected a streaming media upload config for JSONL, got {type(transformed_request)}" + assert isinstance(transformed_request, dict) and "resumable_chunked_upload" in transformed_request, ( + f"Expected a resumable upload config for JSONL, got {type(transformed_request)}" ) - stream = transformed_request["streaming_media_upload"]["body_stream"] + stream = transformed_request["resumable_chunked_upload"]["body_stream"] decoded = json.loads(b"".join(stream.iter_bytes()).decode("utf-8")) assert "request" in decoded, "JSONL transform must wrap each row in {'request': ...}" @@ -205,7 +205,7 @@ class TestVertexAIBinaryFileUpload: optional_params={}, litellm_params={}, ) - assert isinstance(result2, dict) and "streaming_media_upload" in result2 + assert isinstance(result2, dict) and "resumable_chunked_upload" in result2 # Test 3: Upload another binary file binary_content2 = b"\xc4\xe5\xf2\xe5\xeb" diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py index 2e3280c0ed1..1a8be10ca3f 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py @@ -1,12 +1,14 @@ """ -Tests for the streaming OpenAI -> Vertex JSONL batch transform. +Tests for the streaming OpenAI -> Vertex JSONL batch transform and its resumable +chunked upload to GCS. The transform converts batch uploads entry-by-entry rather than materializing the payload in full intermediate lists (decoded str, parsed dicts, transformed -dicts, joined output), which keeps peak memory bounded on large uploads. +dicts, joined output), and the handler streams the result to a GCS resumable +session in bounded chunks, so peak memory stays bounded on large uploads. These tests lock in the behaviour that would regress if the streaming path were -replaced by a list-based pipeline: +replaced by a list-based pipeline or a buffered single-request upload: 1. Byte-for-byte output parity with a list pipeline (wire format). 2. The streaming transform peaks at a clear fraction of a list pipeline on the same input (relative differential, robust to GC noise). @@ -15,12 +17,14 @@ replaced by a list-based pipeline: 4. A tuple-wrapped file handle uploaded through the real create_file ordering keeps every row, including entry 0 (no partial upload from a consumed cursor). + 5. The resumable upload chunks the body, labels Content-Range correctly, and + finalizes even when the payload lands on a chunk boundary. """ import gc import io import json -import tempfile +import threading import time import tracemalloc @@ -43,13 +47,13 @@ from litellm.types.llms.openai import CreateFileRequest def _upload_stream(transformed) -> BaseFileUploadStream: - """Pull the streaming body out of the upload transform result.""" - return transformed["streaming_media_upload"]["body_stream"] + """Pull the streaming body out of a resumable-upload transform result.""" + return transformed["resumable_chunked_upload"]["body_stream"] def _join_upload_body(transformed) -> bytes: """Materialize a transform result's upload body for byte-level assertions.""" - if isinstance(transformed, dict) and "streaming_media_upload" in transformed: + if isinstance(transformed, dict) and "resumable_chunked_upload" in transformed: return b"".join(_upload_stream(transformed).iter_bytes()) if isinstance(transformed, BaseFileUploadStream): return b"".join(transformed.iter_bytes()) @@ -90,7 +94,7 @@ def _reference_vertex_jsonl_string(cfg: VertexAIFilesConfig, content: str) -> st class TestStreamingOutputParity: - def test_transform_create_file_request_returns_streaming_body_parity(self): + def test_transform_create_file_request_returns_resumable_stream_parity(self): cfg = VertexAIFilesConfig() raw = _make_openai_jsonl_bytes(300) request: CreateFileRequest = { @@ -102,10 +106,10 @@ class TestStreamingOutputParity: model="", create_file_data=request, optional_params={}, litellm_params={} ) - # A batch upload must be a streaming-media config carrying a streaming - # body, so the handler can stream it to GCS; a buffered bytes/str return + # A batch upload must be a resumable-upload config carrying a streaming + # body, so the handler can chunk it to GCS; a buffered bytes/str return # would defeat the OOM fix. - assert isinstance(out, dict) and "streaming_media_upload" in out + assert isinstance(out, dict) and "resumable_chunked_upload" in out assert isinstance(_upload_stream(out), BaseFileUploadStream) assert _join_upload_body(out).decode("utf-8") == _reference_vertex_jsonl_string(cfg, raw.decode("utf-8")) @@ -323,10 +327,10 @@ class TestPathSourcedStreaming: litellm_params={"gcs_bucket_name": "test-bucket"}, data=data, ) - assert "uploadType=media" in url + assert "uploadType=resumable" in url out = cfg.transform_create_file_request(model="", create_file_data=data, optional_params={}, litellm_params={}) - assert isinstance(out, dict) and "streaming_media_upload" in out + assert isinstance(out, dict) and "resumable_chunked_upload" in out body = _join_upload_body(out).decode("utf-8") assert body == _reference_vertex_jsonl_string(cfg, raw.decode("utf-8")) lines = body.splitlines() @@ -411,20 +415,26 @@ def _logging_obj() -> Logging: ) -def _gcs_media_mock(status: int = 200): - """A fake GCS simple-media endpoint: one request carries the whole object; - capture the body and headers and return the object resource.""" - state = {"received": bytearray(), "methods": [], "urls": [], "headers": [], "timeouts": []} +def _gcs_resumable_mock(session_url: str, final_status: int = 200): + """A fake GCS resumable endpoint: POST opens a session (URI in Location), + each PUT appends and returns 308 until the final chunk returns 200/201.""" + state = {"received": bytearray(), "ranges": [], "methods": [], "urls": [], "timeouts": []} async def handler(request: httpx.Request) -> httpx.Response: state["methods"].append(request.method) state["urls"].append(str(request.url)) - state["headers"].append(dict(request.headers)) # httpx records the resolved per-request timeout here, so the test can # assert the caller's timeout was forwarded rather than the client default. state["timeouts"].append(request.extensions.get("timeout")) - state["received"].extend(await request.aread()) - return httpx.Response(status, json=_GCS_OBJECT_JSON) + if request.method == "POST": + return httpx.Response(200, headers={"location": session_url}) + body = await request.aread() + content_range = request.headers["content-range"] + state["ranges"].append(content_range) + state["received"].extend(body) + if content_range.rsplit("/", 1)[-1] == "*": + return httpx.Response(308, headers={"range": f"bytes=0-{len(state['received']) - 1}"}) + return httpx.Response(final_status, json=_GCS_OBJECT_JSON) return handler, state @@ -435,8 +445,8 @@ def _async_handler_with(mock) -> AsyncHTTPHandler: return handler -class TestUploadUrl: - def test_batch_jsonl_uses_media_upload_type(self): +class TestResumableUploadUrl: + def test_batch_jsonl_uses_resumable_upload_type(self): cfg = VertexAIFilesConfig() request: CreateFileRequest = { "file": ("batch.jsonl", _make_openai_jsonl_bytes(3), "application/jsonl"), @@ -450,12 +460,13 @@ class TestUploadUrl: litellm_params={"gcs_bucket_name": "test-bucket"}, data=request, ) - # A single media upload is one continuous transfer (no per-chunk - # round-trips), which is what keeps large uploads under client/LB timeouts. - assert "uploadType=media" in url - assert "uploadType=resumable" not in url + # A resumable session chunks the body with Content-Range, so a multi-GB + # upload is never buffered in full; that is what keeps large uploads + # memory-bounded on the proxy. + assert "uploadType=resumable" in url + assert "uploadType=media" not in url - def test_binary_upload_uses_media_upload_type(self): + def test_binary_upload_stays_simple_media(self): cfg = VertexAIFilesConfig() request: CreateFileRequest = { "file": ("doc.pdf", b"%PDF-1.4 binary", "application/pdf"), @@ -473,7 +484,7 @@ class TestUploadUrl: assert "uploadType=resumable" not in url -class TestUploadStreamBody: +class TestResumableStreamBody: def test_stream_matches_legacy_pipeline(self): cfg = VertexAIFilesConfig() raw = _make_openai_jsonl_bytes(120) @@ -502,13 +513,35 @@ class TestUploadStreamBody: assert first == second and len(first) > 0 -@pytest.mark.asyncio -class TestStreamingMediaUpload: - """End-to-end against a faked GCS media endpoint. These fail if the handler - buffers the payload in memory, drops bytes, omits Content-Length (which would - flip httpx to chunked transfer-encoding), or makes more than one request.""" +class TestResumableChunking: + def test_intermediate_chunks_are_exactly_chunk_size(self): + pieces = list(BaseLLMHTTPHandler._iter_resumable_chunks(iter([b"x" * 10]), 4)) + assert pieces == [b"xxxx", b"xxxx", b"xx"] - async def _run(self, raw: bytes, status: int = 200, timeout=None): + def test_exact_multiple_yields_trailing_empty_for_finalize(self): + pieces = list(BaseLLMHTTPHandler._iter_resumable_chunks(iter([b"x" * 8]), 4)) + assert pieces == [b"xxxx", b"xxxx", b""] + + def test_default_chunk_size_is_256kib_multiple(self): + assert BaseLLMHTTPHandler._RESUMABLE_CHUNK_SIZE % (256 * 1024) == 0 + + def test_content_range_intermediate_uses_star_total(self): + assert BaseLLMHTTPHandler._resumable_content_range(0, 4096, is_final=False) == "bytes 0-4095/*" + + def test_content_range_final_uses_real_total(self): + assert BaseLLMHTTPHandler._resumable_content_range(8192, 100, is_final=True) == "bytes 8192-8291/8292" + + def test_content_range_empty_finalize(self): + assert BaseLLMHTTPHandler._resumable_content_range(8192, 0, is_final=True) == "bytes */8192" + + +@pytest.mark.asyncio +class TestResumableUploadProtocol: + """End-to-end against a faked GCS resumable endpoint. These are the tests + that fail if the handler buffers the whole body, drops bytes, mislabels a + Content-Range, follows the 308 instead of continuing, or skips finalize.""" + + async def _run(self, raw: bytes, chunk_size: int, final_status: int = 200, timeout=None): cfg = VertexAIFilesConfig() request: CreateFileRequest = { "file": ("batch.jsonl", raw, "application/jsonl"), @@ -525,8 +558,11 @@ class TestStreamingMediaUpload: transformed = cfg.transform_create_file_request( model="", create_file_data=request, optional_params={}, litellm_params={} ) + transformed["resumable_chunked_upload"]["chunk_size"] = chunk_size expected = _join_upload_body(transformed) - mock, state = _gcs_media_mock(status=status) + + session_url = "https://storage.googleapis.com/upload/sess?upload_id=SID" + mock, state = _gcs_resumable_mock(session_url, final_status=final_status) response = await BaseLLMHTTPHandler().async_create_file( transformed_request=transformed, litellm_params={}, @@ -537,50 +573,122 @@ class TestStreamingMediaUpload: client=_async_handler_with(mock), timeout=timeout, ) - return expected, state, response + return expected, state, response, session_url, api_base - async def test_single_request_carries_whole_payload(self): + async def test_streams_in_chunks_and_reassembles(self): raw = _make_openai_jsonl_bytes(300) - expected, state, response = await self._run(raw) + chunk_size = 4096 + expected, state, response, session_url, api_base = await self._run(raw, chunk_size) - # Exactly one request (the single media upload), and it lands on the - # media endpoint, not a resumable session. - assert state["methods"] == ["POST"] - assert "uploadType=media" in state["urls"][0] + # One session-open POST, then a sequence of chunk PUTs. + assert state["methods"][0] == "POST" + assert set(state["methods"][1:]) == {"PUT"} + assert state["methods"].count("PUT") >= 2, "payload must span multiple chunks" - # The body is streamed with chunked transfer-encoding and no - # Content-Length, which is what proves it is neither buffered in memory - # nor staged to a temp file (the disk-exhaustion guard) before sending. - headers = state["headers"][0] - assert headers.get("transfer-encoding") == "chunked" - assert "content-length" not in headers - # httpx reassembles the chunked body; GCS receives exactly the transform. + # POST opens a resumable session; every chunk goes to the session URI. + assert "uploadType=resumable" in state["urls"][0] + assert all(u == session_url for u in state["urls"][1:]) + + # Every non-final chunk is exactly chunk_size with an unknown-total range; + # the final chunk carries the real total. + intermediate = state["ranges"][:-1] + for index, content_range in enumerate(intermediate): + assert content_range == f"bytes {index * chunk_size}-{(index + 1) * chunk_size - 1}/*" + total = len(expected) + last_offset = len(intermediate) * chunk_size + if last_offset == total: # payload landed on a chunk boundary + assert state["ranges"][-1] == f"bytes */{total}" + else: + assert state["ranges"][-1] == f"bytes {last_offset}-{total - 1}/{total}" + + # The bytes GCS received are exactly the transformed batch payload. assert bytes(state["received"]) == expected assert response.object == "file" - async def test_failed_upload_raises(self): + async def test_exact_multiple_finalizes_with_empty_chunk(self): + # Build a body that is an exact multiple of the chunk size so the stream + # ends on a chunk boundary; the upload must still finalize (bytes */TOTAL). + chunk_size = 256 + stream = _FixedBytesStream(b"a" * (chunk_size * 3)) + config = {"body_stream": stream, "chunk_size": chunk_size} + session_url = "https://storage.googleapis.com/upload/sess?upload_id=SID" + mock, state = _gcs_resumable_mock(session_url) + + response = await BaseLLMHTTPHandler()._aresumable_chunked_upload( + client=_async_handler_with(mock), + initiate_url="https://storage.googleapis.com/upload?uploadType=resumable", + base_headers={"Authorization": "Bearer x"}, + config=config, + timeout=None, + ) + + assert state["ranges"][-1] == f"bytes */{chunk_size * 3}" + assert bytes(state["received"]) == b"a" * (chunk_size * 3) + assert response.status_code == 200 + + async def test_failed_chunk_raises(self): raw = _make_openai_jsonl_bytes(80) with pytest.raises(Exception): - await self._run(raw, status=403) + await self._run(raw, chunk_size=4096, final_status=403) + + async def test_transform_runs_off_the_event_loop(self): + # The per-row OpenAI->Vertex transform is CPU-bound; the async upload must + # assemble chunks on a worker thread so it never blocks the event loop + # between chunk PUTs. This records the thread each body piece is produced + # on and asserts none of them is the event-loop thread; reverting the + # offload to a plain sync for-loop makes this fail. + class _ThreadRecordingStream(BaseFileUploadStream): + def __init__(self, data: bytes, piece: int = 64): + self._data = data + self._piece = piece + self.producer_threads: set[int] = set() + + def iter_bytes(self): + for i in range(0, len(self._data), self._piece): + self.producer_threads.add(threading.get_ident()) + yield self._data[i : i + self._piece] + + stream = _ThreadRecordingStream(b"a" * 2048) + mock, _ = _gcs_resumable_mock("https://storage.googleapis.com/upload/sess?upload_id=SID") + await BaseLLMHTTPHandler()._aresumable_chunked_upload( + client=_async_handler_with(mock), + initiate_url="https://storage.googleapis.com/upload?uploadType=resumable", + base_headers={"Authorization": "Bearer x"}, + config={"body_stream": stream, "chunk_size": 256}, + timeout=None, + ) + + assert stream.producer_threads, "stream must have produced at least one piece" + assert threading.get_ident() not in stream.producer_threads, ( + "chunk assembly (the CPU-bound transform) must run on a worker thread, not the event loop" + ) async def test_request_timeout_is_forwarded(self): - # The caller's per-request timeout must reach the GCS upload; every other - # upload branch forwards it. httpx records the resolved timeout in + # The caller's per-request timeout must reach every GCS request (session + # open and each chunk PUT). httpx records the resolved timeout in # request.extensions["timeout"]; a dropped timeout would show the client # default instead of the value passed here. - raw = _make_openai_jsonl_bytes(20) - _, state, _ = await self._run(raw, timeout=httpx.Timeout(137.0)) - forwarded = state["timeouts"][0] - assert forwarded is not None - assert forwarded.get("read") == 137.0 and forwarded.get("write") == 137.0 + raw = _make_openai_jsonl_bytes(300) + _, state, _, _, _ = await self._run(raw, chunk_size=4096, timeout=httpx.Timeout(137.0)) + assert len(state["timeouts"]) >= 2 + for forwarded in state["timeouts"]: + assert forwarded is not None + assert forwarded.get("read") == 137.0 and forwarded.get("write") == 137.0 - async def test_upload_does_not_stage_to_disk(self, monkeypatch): - # Disk-exhaustion guard: the transformed body must stream to GCS, never be - # written to a temp file first. If any tempfile is created during the - # upload, an attacker could fill the proxy's temp volume with large - # concurrent uploads. - created = [] - real_tempfile = tempfile.TemporaryFile - monkeypatch.setattr(tempfile, "TemporaryFile", lambda *a, **k: (created.append(1), real_tempfile(*a, **k))[1]) - await self._run(_make_openai_jsonl_bytes(50)) - assert created == [] + async def test_unexpected_session_init_status_raises(self): + # raise_for_status only raises on 4xx/5xx, so a 2xx that is not 200/201 + # (e.g. 202) would otherwise fall through to a misleading "no session URL" + # error. The handler must surface the actual status instead. + async def bad_init(request: httpx.Request) -> httpx.Response: + if request.method == "POST": + return httpx.Response(202, request=request) + return httpx.Response(200, json=_GCS_OBJECT_JSON, request=request) + + with pytest.raises(ValueError, match="unexpected session-init status 202"): + await BaseLLMHTTPHandler()._aresumable_chunked_upload( + client=_async_handler_with(bad_init), + initiate_url="https://storage.googleapis.com/upload?uploadType=resumable", + base_headers={"Authorization": "Bearer x"}, + config={"body_stream": _FixedBytesStream(b"a" * 128), "chunk_size": 256}, + timeout=None, + )