scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 87

fix(vertex_ai/files): single media upload for batch files to fix 499s on large uploads (#31653)

Browse source 
* fix(vertex_ai/files): upload batch files in a single media request to fix 499s on large uploads

PR #31036 switched the vertex batch file upload from a single GCS media
upload to a chunked resumable session. The resumable path sends the body as
many sequential PUTs, each waiting a full round-trip to GCS before the next,
so a multi-GB upload accumulates hundreds of round-trips and overruns the
client/load-balancer request timeout, surfacing as 499s (client closed
connection) on files as small as 500MB. This was a regression from the
last-known-good commit, where the upload completed as one continuous request.

Revert the batch upload to a single uploadType=media request, but stage the
transformed payload to a temp file first so peak memory stays bounded (the
goal of the resumable rewrite) without the per-chunk round-trips. The temp
file is closed deterministically (TemporaryFile unlinks on close), not left
to the GC. The now-unused resumable chunked-upload plumbing is removed.

Also swap the per-row transform's stdlib json for orjson (parse + serialize),
which is ~4x faster on this hot path; the streaming body now emits compact
orjson bytes.

The request stays synchronous, so the returned file object is real and
POST /v1/batches keeps working immediately against the uploaded object.

Tests: single media request carries the whole payload with a real
Content-Length (no chunked transfer-encoding); failed upload raises; the
staged temp file is closed deterministically; byte-for-byte transform parity.

* test(vertex_ai/files): mock single media upload POST instead of removed resumable method

test_avertex_batch_prediction patched BaseLLMHTTPHandler._aresumable_chunked_upload, which was removed when the batch jsonl upload moved from a chunked resumable GCS session to a single uploadType=media request. Patch the raw httpx.AsyncClient.post that _astage_and_upload_media issues so the real staging, upload and response transform run while the GCS object response is mocked, and assert the media URL and Content-Type.

* fix(vertex_ai/files): forward request timeout to media upload, drop orjson, sort imports

Forward the per-request timeout through _stage_and_upload_media /
_astage_and_upload_media to the GCS POST. Every other upload branch forwards
it; the new media path was dropping it, so a caller-provided timeout was
silently ignored (the files path passes 600s by default, but a custom
request_timeout would not have reached this upload). Regression test asserts
the resolved timeout reaches the request (mutation-verified).

Revert the orjson swap in the batch transform: importing orjson at module load
in this core-path file broke `import litellm` on environments without orjson
(the Windows import test). Back to stdlib json; the upload leg dominates large
uploads anyway, so the transform-side win was marginal.

Fix import ordering in llm_http_handler.py (I001) introduced by the new imports.

* fix(vertex_ai/files): stream batch upload to GCS instead of staging to a temp file

Addresses a disk-exhaustion concern: staging the full transformed batch body to
a local temp file before the GCS request meant an authenticated user could fill
the proxy's temp volume with large concurrent uploads (on top of Starlette's
input spool).

GCS's simple/media upload accepts chunked transfer-encoding, so stream the
transform straight to the single media request instead. Each block is produced
on a worker thread (the transform never runs on the event loop) and sent
chunked, so the body is neither buffered in memory nor written to disk, and the
upload is still one continuous request (no per-chunk round-trips, no 499). Drops
the temp-file staging, the tempfile/IO imports, and Content-Length computation.

Regression test asserts the upload streams (chunked transfer-encoding, no
Content-Length) and creates no temp file; mutation-verified that reintroducing
staging fails it.

(cherry picked from commit 85840aef51)
This commit is contained in:
mubashir1osmani 2026-06-29 17:31:32 -07:00 • committed by Yuneng Jiang
parent 19c236aaac
commit 7b026008b9
No known key found for this signature in database
6 changed files with 217 additions and 485 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

285
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -41,7 +41,10 @@ from litellm.llms.base_llm.chat.transformation import BaseConfig
from litellm.llms.base_llm.containers.transformation import BaseContainerConfig
from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig
from litellm.llms.base_llm.evals.transformation import BaseEvalsAPIConfig
from litellm.llms.base_llm.files.transformation import BaseFilesConfig
from litellm.llms.base_llm.files.transformation import (
BaseFilesConfig,
BaseFileUploadStream,
)
from litellm.llms.base_llm.google_genai.transformation import (
BaseGoogleGenAIGenerateContentConfig,
)
@ -80,7 +83,7 @@ from litellm.types.containers.main import (
ContainerObject,
DeleteContainerResult,
)
from litellm.types.files import TwoStepFileUploadConfig
from litellm.types.files import StreamingMediaUploadConfig, TwoStepFileUploadConfig
from litellm.types.integrations.custom_logger import (
AgenticLoopPlan,
AgenticLoopRequestPatch,
@ -3232,18 +3235,15 @@ class BaseLLMHTTPHandler:
data=presigned_request["data"],
timeout=timeout,
)
elif (
isinstance(transformed_request, dict)
and "resumable_chunked_upload" in transformed_request
):
elif isinstance(transformed_request, dict) and "streaming_media_upload" in transformed_request:
media_cfg = cast(StreamingMediaUploadConfig, transformed_request["streaming_media_upload"])
try:
upload_response = self._resumable_chunked_upload(
upload_response = self._upload_media(
client=sync_httpx_client,
initiate_url=api_base,
url=api_base,
base_headers=headers,
config=cast(Dict[str, Any], transformed_request)[
"resumable_chunked_upload"
],
body_stream=cast(BaseFileUploadStream, media_cfg["body_stream"]),
content_type=media_cfg.get("content_type") or "application/octet-stream",
timeout=timeout,
)
except Exception as e:
@ -3330,13 +3330,12 @@ class BaseLLMHTTPHandler:
input="",
api_key="",
additional_args={
# A resumable upload config holds a reference to the (potentially
# A streaming upload config holds a reference to the (potentially
# huge) upload payload; logging deep-copies additional_args, so log
# a placeholder instead of re-materializing the payload.
"complete_input_dict": (
"<resumable chunked upload>"
if isinstance(transformed_request, dict)
and "resumable_chunked_upload" in transformed_request
"<streaming media upload>"
if isinstance(transformed_request, dict) and "streaming_media_upload" in transformed_request
else transformed_request
),
"api_base": api_base,
@ -3415,18 +3414,15 @@ class BaseLLMHTTPHandler:
data=presigned_request["data"],
timeout=timeout,
)
elif (
isinstance(transformed_request, dict)
and "resumable_chunked_upload" in transformed_request
):
elif isinstance(transformed_request, dict) and "streaming_media_upload" in transformed_request:
media_cfg = cast(StreamingMediaUploadConfig, transformed_request["streaming_media_upload"])
try:
upload_response = await self._aresumable_chunked_upload(
upload_response = await self._aupload_media(
client=async_httpx_client,
initiate_url=api_base,
url=api_base,
base_headers=headers,
config=cast(Dict[str, Any], transformed_request)[
"resumable_chunked_upload"
],
body_stream=cast(BaseFileUploadStream, media_cfg["body_stream"]),
content_type=media_cfg.get("content_type") or "application/octet-stream",
timeout=timeout,
)
except Exception as e:
@ -3475,222 +3471,81 @@ class BaseLLMHTTPHandler:
litellm_params=litellm_params,
)
# 8 MiB; a 256 KiB multiple, which GCS requires for every non-final chunk.
_RESUMABLE_CHUNK_SIZE = 8 * 1024 * 1024
# 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
@staticmethod
def _iter_resumable_chunks(
byte_iter: Iterator[bytes], chunk_size: int
) -> Iterator[bytes]:
"""Regroup a byte stream into ``chunk_size`` pieces, yielding a final
partial piece only when it is non-empty. Every full piece is exactly
``chunk_size`` bytes (kept a 256 KiB multiple for GCS) and never more than
one chunk is buffered. An exactly chunk-aligned stream yields only full
chunks, so the upload finalizes on its last data chunk instead of making
an extra empty request; a 0-byte stream yields nothing and the caller
finalizes with a single empty request.
"""
def _iter_in_blocks(byte_iter: Iterator[bytes], block_size: int) -> Iterator[bytes]:
buf = bytearray()
for piece in byte_iter:
buf.extend(piece)
while len(buf) >= chunk_size:
yield bytes(buf[:chunk_size])
del buf[:chunk_size]
while len(buf) >= block_size:
yield bytes(buf[:block_size])
del buf[:block_size]
if buf:
yield bytes(buf)
@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 _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_request_kwargs(
headers: dict,
content: bytes,
timeout: Optional[Union[float, httpx.Timeout]],
) -> dict:
kwargs: Dict[str, Any] = {"headers": headers, "content": content}
if timeout is not None:
kwargs["timeout"] = timeout
return kwargs
def _resumable_chunked_upload(
def _upload_media(
self,
*,
client: HTTPHandler,
initiate_url: str,
base_headers: dict,
config: dict,
timeout: Optional[Union[float, httpx.Timeout]],
) -> httpx.Response:
"""Open a GCS resumable session, then PUT the body in bounded chunks so a
large upload is never held in memory in full."""
stream = config["body_stream"]
chunk_size = config.get("chunk_size", self._RESUMABLE_CHUNK_SIZE)
session_url_header = config.get("session_url_header", "location")
httpx_client = client.client
init_headers = {**base_headers, **config.get("initiate_headers", {})}
init_req = httpx_client.build_request(
"POST",
initiate_url,
**self._resumable_request_kwargs(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):
init_resp.raise_for_status()
session_url = init_resp.headers.get(session_url_header)
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,
data: bytes,
offset: int,
*,
is_final: bool,
base_headers: Dict[str, str],
body_stream: BaseFileUploadStream,
content_type: str,
timeout: Optional[Union[float, httpx.Timeout]],
) -> httpx.Response:
headers = {
**base_headers,
"Content-Range": self._resumable_content_range(offset, len(data), is_final),
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),
}
req = httpx_client.build_request(
"PUT", url, **self._resumable_request_kwargs(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}")
if timeout is not None:
kwargs["timeout"] = timeout
resp = client.client.post(url, **kwargs)
self._check_media_upload_response(resp)
return resp
async def _aresumable_chunked_upload(
async def _aupload_media(
self,
*,
client: AsyncHTTPHandler,
initiate_url: str,
base_headers: dict,
config: dict,
timeout: Optional[Union[float, httpx.Timeout]],
) -> httpx.Response:
stream = config["body_stream"]
chunk_size = config.get("chunk_size", self._RESUMABLE_CHUNK_SIZE)
session_url_header = config.get("session_url_header", "location")
httpx_client = client.client
init_headers = {**base_headers, **config.get("initiate_headers", {})}
init_req = httpx_client.build_request(
"POST",
initiate_url,
**self._resumable_request_kwargs(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):
init_resp.raise_for_status()
session_url = init_resp.headers.get(session_url_header)
if not session_url:
raise ValueError(
f"resumable upload: no session URL in '{session_url_header}' header"
)
offset = 0
pending: Optional[bytes] = None
# Producing each chunk runs the synchronous per-row transform for that
# chunk's worth of rows. Pull it off the event loop thread so a large
# upload does not block other concurrent requests between PUTs.
chunk_iter = self._iter_resumable_chunks(stream.iter_bytes(), chunk_size)
done = object()
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 = chunk
return await self._asend_resumable_chunk(
httpx_client,
session_url,
base_headers,
pending or b"",
offset,
is_final=True,
timeout=timeout,
)
async def _asend_resumable_chunk(
self,
httpx_client: httpx.AsyncClient,
url: str,
base_headers: dict,
data: bytes,
offset: int,
*,
is_final: bool,
base_headers: Dict[str, str],
body_stream: BaseFileUploadStream,
content_type: str,
timeout: Optional[Union[float, httpx.Timeout]],
) -> httpx.Response:
headers = {
**base_headers,
"Content-Range": self._resumable_content_range(offset, len(data), is_final),
}
req = httpx_client.build_request(
"PUT", url, **self._resumable_request_kwargs(headers, data, timeout)
)
resp = await httpx_client.send(req, follow_redirects=False)
"""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))
done = object()
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)
await resp.aread()
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}")
self._check_media_upload_response(resp)
return resp
def create_batch(

25
litellm/llms/vertex_ai/files/transformation.py
View file

@ -60,7 +60,7 @@ from litellm.types.llms.openai import (
OpenAIFileObject,
PathLike,
)
from litellm.types.files import ResumableChunkedUploadConfig
from litellm.types.files import StreamingMediaUploadConfig
from litellm.types.llms.vertex_ai import GcsBucketResponse
from litellm.types.utils import LlmProviders, ModelResponse
@ -389,21 +389,11 @@ 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)
# 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}"
endpoint = f"upload/storage/v1/b/{bucket_name}/o?uploadType=media&name={encoded_object_name}"
api_base = api_base or "https://storage.googleapis.com"
if not api_base:
raise ValueError("api_base is required")
@ -455,8 +445,9 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
"""
2 Cases:
1. Handle basic file upload
2. Handle batch file upload (.jsonl), streamed to a GCS resumable
session so large uploads stay memory-bounded.
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.
"""
file_data = create_file_data.get("file")
if file_data is None:
@ -468,14 +459,12 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
content_type=content_type,
):
return {
"resumable_chunked_upload": ResumableChunkedUploadConfig(
"streaming_media_upload": StreamingMediaUploadConfig(
body_stream=_OpenAIToVertexBatchUploadStream(
file_data,
self._map_openai_to_vertex_params,
),
initiate_headers={
"X-Upload-Content-Type": "application/json",
},
content_type="application/json",
)
}

17
litellm/types/files.py
View file

@ -323,19 +323,18 @@ class TwoStepFileUploadConfig(TypedDict, total=False):
upload_url_key: str
class ResumableChunkedUploadConfig(TypedDict, total=False):
"""Drives a memory-bounded resumable upload (GCS JSON API).
class StreamingMediaUploadConfig(TypedDict, total=False):
"""Drives a memory-bounded single-request upload (GCS simple/media upload).
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.
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.
``body_stream`` is a ``BaseFileUploadStream``; it is typed ``Any`` here to
avoid importing the llms layer into types.
"""
body_stream: Required[Any]
chunk_size: int
session_url_header: str
initiate_headers: Dict[str, str]
content_type: str

38
tests/batches_tests/test_openai_batches_and_files.py
View file

@ -513,25 +513,26 @@ async def test_avertex_batch_prediction(monkeypatch):
mock_response.status_code = 200
return mock_response
# Batch jsonl file creation now streams to a GCS resumable session via
# _aresumable_chunked_upload (httpx send), not AsyncHTTPHandler.post, so mock
# that entry point to return the GCS object response. The resumable protocol
# itself is covered in test_vertex_ai_files_streaming.py.
mock_upload_response = httpx.Response(
200,
json=mock_file_response,
request=httpx.Request("PUT", "https://storage.googleapis.com/upload"),
)
# 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.
with (
patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
side_effect=mock_side_effect,
) as mock_global_post,
patch(
"litellm.llms.custom_httpx.llm_http_handler.BaseLLMHTTPHandler._aresumable_chunked_upload",
new_callable=AsyncMock,
return_value=mock_upload_response,
),
patch.object(
httpx.AsyncClient,
"post",
new_callable=AsyncMock,
return_value=httpx.Response(
200,
json=mock_file_response,
request=httpx.Request("POST", "https://storage.googleapis.com/upload"),
),
) as mock_gcs_upload,
):
litellm.set_verbose = True
litellm._turn_on_debug()
@ -552,6 +553,15 @@ 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"
)
# Create batch
create_batch_response = await litellm.acreate_batch(
completion_window="24h",

43
tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py
View file

@ -37,9 +37,7 @@ class TestVertexAIBinaryFileUpload:
# Create mock PDF binary data (with non-UTF-8 bytes)
# PDF files start with %PDF- and contain binary data
mock_pdf_content = b"%PDF-1.4\n%\xc4\xe5\xf2\xe5\xeb\xa7\xf3\xa0\xd0\xc4\xc6\n"
mock_pdf_content += (
b"\x00\x01\x02\x03\xff\xfe\xfd" * 100
) # Add more binary data
mock_pdf_content += b"\x00\x01\x02\x03\xff\xfe\xfd" * 100 # Add more binary data
# Create file object
file_obj = io.BytesIO(mock_pdf_content)
@ -60,14 +58,12 @@ class TestVertexAIBinaryFileUpload:
)
# Verify the transformation returns bytes (not string)
assert isinstance(
transformed_request, bytes
), f"Expected bytes for binary file, got {type(transformed_request)}"
assert isinstance(transformed_request, bytes), (
f"Expected bytes for binary file, got {type(transformed_request)}"
)
# Verify the bytes match the original content
assert (
transformed_request == mock_pdf_content
), "Transformed request should preserve binary content exactly"
assert transformed_request == mock_pdf_content, "Transformed request should preserve binary content exactly"
# Verify that the bytes contain non-UTF-8 characters
# This should raise UnicodeDecodeError if we try to decode
@ -132,16 +128,14 @@ class TestVertexAIBinaryFileUpload:
pytest.fail(f"httpx should accept bytes in data parameter: {e}")
# Document the expected behavior
assert isinstance(
mock_binary_data, bytes
), "Binary file data should remain as bytes"
assert isinstance(mock_binary_data, bytes), "Binary file data should remain as bytes"
@pytest.mark.asyncio
async def test_jsonl_file_upload_returns_resumable_stream(self):
async def test_jsonl_file_upload_returns_streaming_body(self):
"""
Test that JSONL batch files are transformed into a resumable-upload config
Test that JSONL batch files are transformed into a streaming-media config
carrying a streaming body (not a buffered bytes payload), so the handler
can stream the upload to GCS in bounded chunks.
can stage the upload to a temp file and send it in one media request.
"""
# Create mock JSONL content
mock_jsonl_content = (
@ -164,16 +158,13 @@ class TestVertexAIBinaryFileUpload:
litellm_params={},
)
assert (
isinstance(transformed_request, dict)
and "resumable_chunked_upload" in transformed_request
), f"Expected a resumable upload config for JSONL, got {type(transformed_request)}"
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)}"
)
stream = transformed_request["resumable_chunked_upload"]["body_stream"]
stream = transformed_request["streaming_media_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': ...}"
assert "request" in decoded, "JSONL transform must wrap each row in {'request': ...}"
@pytest.mark.asyncio
async def test_mixed_file_types_in_sequence(self):
@ -214,7 +205,7 @@ class TestVertexAIBinaryFileUpload:
optional_params={},
litellm_params={},
)
assert isinstance(result2, dict) and "resumable_chunked_upload" in result2
assert isinstance(result2, dict) and "streaming_media_upload" in result2
# Test 3: Upload another binary file
binary_content2 = b"\xc4\xe5\xf2\xe5\xeb"
@ -264,7 +255,5 @@ class TestVertexAIBinaryFileUpload:
},
}
assert (
expected_behavior["binary_files"]["encoding"] == "none - preserve raw bytes"
)
assert expected_behavior["binary_files"]["encoding"] == "none - preserve raw bytes"
assert expected_behavior["text_files"]["encoding"] == "UTF-8"

294
tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py
View file

@ -20,6 +20,7 @@ replaced by a list-based pipeline:
import gc
import io
import json
import tempfile
import time
import tracemalloc
@ -41,15 +42,15 @@ from litellm.llms.vertex_ai.files.transformation import (
from litellm.types.llms.openai import CreateFileRequest
def _resumable_stream(transformed) -> BaseFileUploadStream:
"""Pull the streaming body out of a resumable-upload transform result."""
return transformed["resumable_chunked_upload"]["body_stream"]
def _upload_stream(transformed) -> BaseFileUploadStream:
"""Pull the streaming body out of the upload transform result."""
return transformed["streaming_media_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 "resumable_chunked_upload" in transformed:
return b"".join(_resumable_stream(transformed).iter_bytes())
if isinstance(transformed, dict) and "streaming_media_upload" in transformed:
return b"".join(_upload_stream(transformed).iter_bytes())
if isinstance(transformed, BaseFileUploadStream):
return b"".join(transformed.iter_bytes())
if isinstance(transformed, str):
@ -83,17 +84,13 @@ def _reference_vertex_jsonl_string(cfg: VertexAIFilesConfig, content: str) -> st
transform, so the streaming path can be checked against it for parity."""
entries = [json.loads(line) for line in content.splitlines() if line.strip()]
return "\n".join(
json.dumps(
_openai_batch_jsonl_entry_to_vertex_wrapped_request(
entry, cfg._map_openai_to_vertex_params
)
)
json.dumps(_openai_batch_jsonl_entry_to_vertex_wrapped_request(entry, cfg._map_openai_to_vertex_params))
for entry in entries
)
class TestStreamingOutputParity:
def test_transform_create_file_request_returns_resumable_stream_parity(self):
def test_transform_create_file_request_returns_streaming_body_parity(self):
cfg = VertexAIFilesConfig()
raw = _make_openai_jsonl_bytes(300)
request: CreateFileRequest = {
@ -105,14 +102,12 @@ class TestStreamingOutputParity:
model="", create_file_data=request, optional_params={}, litellm_params={}
)
# A batch upload must be a resumable-upload config carrying a streaming
# body, so the handler can chunk it; a buffered bytes/str return would
# defeat the OOM fix.
assert isinstance(out, dict) and "resumable_chunked_upload" in out
assert isinstance(_resumable_stream(out), BaseFileUploadStream)
assert _join_upload_body(out).decode("utf-8") == _reference_vertex_jsonl_string(
cfg, raw.decode("utf-8")
)
# 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
# would defeat the OOM fix.
assert isinstance(out, dict) and "streaming_media_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"))
class TestFileLikeInputNotPartiallyConsumed:
@ -215,9 +210,7 @@ class TestStreamingLineIterator:
def seek(self, *args):
raise io.UnsupportedOperation("not seekable")
handle = _NonSeekable(
b'{"custom_id": "request-0"}\n{"custom_id": "request-1"}\n'
)
handle = _NonSeekable(b'{"custom_id": "request-0"}\n{"custom_id": "request-1"}\n')
with pytest.raises(ValueError, match="seekable"):
list(_iter_openai_jsonl_lines(handle))
@ -235,13 +228,8 @@ class TestGetObjectNameLazyParse:
cfg = VertexAIFilesConfig()
# Tail rows are deliberately not valid JSON. Parsing the whole payload
# would raise here; a first-row-only parse must not.
raw = (
b'{"custom_id": "r-0", "body": {"model": "gemini-2.5-flash"}}\n'
b"garbage line that is not json\n"
)
object_name = cfg.get_object_name(
("batch.jsonl", raw, "application/jsonl"), purpose="batch"
)
raw = b'{"custom_id": "r-0", "body": {"model": "gemini-2.5-flash"}}\ngarbage line that is not json\n'
object_name = cfg.get_object_name(("batch.jsonl", raw, "application/jsonl"), purpose="batch")
assert "gemini-2.5-flash" in object_name
@ -278,15 +266,11 @@ class TestStreamingPeakMemory:
def drain_stream():
# Consume the upload body one row at a time, as the chunked uploader
# does, without accumulating it.
for _ in _OpenAIToVertexBatchUploadStream(
raw, cfg._map_openai_to_vertex_params
).iter_bytes():
for _ in _OpenAIToVertexBatchUploadStream(raw, cfg._map_openai_to_vertex_params).iter_bytes():
pass
streaming_peak = self._measure(drain_stream)
list_peak = self._measure(
lambda: _reference_vertex_jsonl_string(cfg, content_str)
)
list_peak = self._measure(lambda: _reference_vertex_jsonl_string(cfg, content_str))
# Core guard: the lazily consumed streaming body peaks well under a list
# pipeline that materializes every transformed row. Building full
@ -305,9 +289,7 @@ class TestStreamingPeakMemory:
# first-row parse should allocate only a small fraction of the payload;
# parsing every row would blow past this bound.
peak = self._measure(lambda: cfg.get_object_name(file_data, purpose="batch"))
assert (
peak / len(raw) < 2.0
), "get_object_name should not copy the whole payload"
assert peak / len(raw) < 2.0, "get_object_name should not copy the whole payload"
class TestPathSourcedStreaming:
@ -341,12 +323,10 @@ class TestPathSourcedStreaming:
litellm_params={"gcs_bucket_name": "test-bucket"},
data=data,
)
assert "uploadType=resumable" in url
assert "uploadType=media" in url
out = cfg.transform_create_file_request(
model="", create_file_data=data, optional_params={}, litellm_params={}
)
assert isinstance(out, dict) and "resumable_chunked_upload" in out
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
body = _join_upload_body(out).decode("utf-8")
assert body == _reference_vertex_jsonl_string(cfg, raw.decode("utf-8"))
lines = body.splitlines()
@ -371,7 +351,7 @@ class TestPathSourcedStreaming:
out = cfg.transform_create_file_request(
model="", create_file_data=data, optional_params={}, litellm_params={}
)
for _ in _resumable_stream(out).iter_bytes():
for _ in _upload_stream(out).iter_bytes():
pass # drain without accumulating
gc.collect()
@ -384,20 +364,15 @@ class TestPathSourcedStreaming:
# Streaming from disk must not materialize the payload. Reading the whole
# file into bytes (the pre-fix path) would push peak past the file size.
assert peak < len(raw) * 0.3, (
f"peak {peak} not bounded vs payload {len(raw)} "
f"(ratio {peak / len(raw):.2f})"
)
assert peak < len(raw) * 0.3, f"peak {peak} not bounded vs payload {len(raw)} (ratio {peak / len(raw):.2f})"
def test_path_source_stream_is_reiterable(self, tmp_path):
cfg = VertexAIFilesConfig()
path, _ = self._write_jsonl(tmp_path, 50)
data = self._batch_request(path)
out = cfg.transform_create_file_request(
model="", create_file_data=data, optional_params={}, litellm_params={}
)
stream = _resumable_stream(out)
out = cfg.transform_create_file_request(model="", create_file_data=data, optional_params={}, litellm_params={})
stream = _upload_stream(out)
first = b"".join(stream.iter_bytes())
second = b"".join(stream.iter_bytes())
assert first == second and len(first) > 0
@ -436,25 +411,20 @@ def _logging_obj() -> Logging:
)
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": []}
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": []}
async def handler(request: httpx.Request) -> httpx.Response:
state["methods"].append(request.method)
state["urls"].append(str(request.url))
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)
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)
return handler, state
@ -465,8 +435,8 @@ def _async_handler_with(mock) -> AsyncHTTPHandler:
return handler
class TestResumableUploadUrl:
def test_batch_jsonl_uses_resumable_upload_type(self):
class TestUploadUrl:
def test_batch_jsonl_uses_media_upload_type(self):
cfg = VertexAIFilesConfig()
request: CreateFileRequest = {
"file": ("batch.jsonl", _make_openai_jsonl_bytes(3), "application/jsonl"),
@ -480,29 +450,12 @@ class TestResumableUploadUrl:
litellm_params={"gcs_bucket_name": "test-bucket"},
data=request,
)
assert "uploadType=resumable" in url
assert "uploadType=media" not in url
# 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
def test_batch_text_plain_uses_resumable_upload_type(self):
# Clients often label a .jsonl batch upload as text/plain; it must still
# take the streaming/resumable path, not the buffered media path.
cfg = VertexAIFilesConfig()
request: CreateFileRequest = {
"file": ("batch.jsonl", _make_openai_jsonl_bytes(3), "text/plain"),
"purpose": "batch",
}
url = cfg.get_complete_file_url(
api_base=None,
api_key=None,
model="",
optional_params={},
litellm_params={"gcs_bucket_name": "test-bucket"},
data=request,
)
assert "uploadType=resumable" in url
assert "uploadType=media" not in url
def test_binary_upload_stays_simple_media(self):
def test_binary_upload_uses_media_upload_type(self):
cfg = VertexAIFilesConfig()
request: CreateFileRequest = {
"file": ("doc.pdf", b"%PDF-1.4 binary", "application/pdf"),
@ -520,14 +473,12 @@ class TestResumableUploadUrl:
assert "uploadType=resumable" not in url
class TestResumableStreamBody:
class TestUploadStreamBody:
def test_stream_matches_legacy_pipeline(self):
cfg = VertexAIFilesConfig()
raw = _make_openai_jsonl_bytes(120)
stream = _OpenAIToVertexBatchUploadStream(raw, cfg._map_openai_to_vertex_params)
assert b"".join(stream.iter_bytes()).decode(
"utf-8"
) == _reference_vertex_jsonl_string(cfg, raw.decode("utf-8"))
assert b"".join(stream.iter_bytes()).decode("utf-8") == _reference_vertex_jsonl_string(cfg, raw.decode("utf-8"))
def test_stream_is_reiterable_for_retries(self):
# A one-shot generator would make a transport retry upload an empty body;
@ -545,59 +496,19 @@ class TestResumableStreamBody:
# an empty body silently.
cfg = VertexAIFilesConfig()
raw = _make_openai_jsonl_bytes(40)
stream = _OpenAIToVertexBatchUploadStream(
io.BytesIO(raw), cfg._map_openai_to_vertex_params
)
stream = _OpenAIToVertexBatchUploadStream(io.BytesIO(raw), cfg._map_openai_to_vertex_params)
first = b"".join(stream.iter_bytes())
second = b"".join(stream.iter_bytes())
assert first == second and len(first) > 0
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"]
def test_exact_multiple_yields_no_trailing_empty(self):
# An exactly chunk-aligned stream yields only full chunks; the upload
# finalizes on the last data chunk instead of an extra empty request.
pieces = list(BaseLLMHTTPHandler._iter_resumable_chunks(iter([b"x" * 8]), 4))
assert pieces == [b"xxxx", b"xxxx"]
def test_empty_stream_yields_nothing(self):
# A 0-byte stream yields no chunks; the caller finalizes with one empty
# request (bytes */0).
assert list(BaseLLMHTTPHandler._iter_resumable_chunks(iter([]), 4)) == []
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."""
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."""
async def _run(self, raw: bytes, chunk_size: int, final_status: int = 200):
async def _run(self, raw: bytes, status: int = 200, timeout=None):
cfg = VertexAIFilesConfig()
request: CreateFileRequest = {
"file": ("batch.jsonl", raw, "application/jsonl"),
@ -614,11 +525,8 @@ class TestResumableUploadProtocol:
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)
session_url = "https://storage.googleapis.com/upload/sess?upload_id=SID"
mock, state = _gcs_resumable_mock(session_url, final_status=final_status)
mock, state = _gcs_media_mock(status=status)
response = await BaseLLMHTTPHandler().async_create_file(
transformed_request=transformed,
litellm_params={},
@ -627,70 +535,52 @@ class TestResumableUploadProtocol:
api_base=api_base,
logging_obj=_logging_obj(),
client=_async_handler_with(mock),
timeout=None,
timeout=timeout,
)
return expected, state, response, session_url, api_base
return expected, state, response
async def test_streams_in_chunks_and_reassembles(self):
async def test_single_request_carries_whole_payload(self):
raw = _make_openai_jsonl_bytes(300)
chunk_size = 4096
expected, state, response, session_url, api_base = await self._run(
raw, chunk_size
)
expected, state, response = await self._run(raw)
# 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"
# 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]
# 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.
# 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.
assert bytes(state["received"]) == expected
assert response.object == "file"
async def test_exact_multiple_finalizes_on_last_data_chunk(self):
# A body that is an exact multiple of the chunk size finalizes on its
# last data chunk (bytes (TOTAL-chunk)-(TOTAL-1)/TOTAL), with no extra
# empty finalize request.
chunk_size = 256
total = chunk_size * 3
stream = _FixedBytesStream(b"a" * total)
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 {total - chunk_size}-{total - 1}/{total}"
assert "*" not in state["ranges"][-1]
assert bytes(state["received"]) == b"a" * total
assert response.status_code == 200
async def test_failed_chunk_raises(self):
async def test_failed_upload_raises(self):
raw = _make_openai_jsonl_bytes(80)
with pytest.raises(Exception):
await self._run(raw, chunk_size=4096, final_status=403)
await self._run(raw, status=403)
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
# 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
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 == []