diff --git a/litellm/files/utils.py b/litellm/files/utils.py index 9639042caab..a0df7a89b0f 100644 --- a/litellm/files/utils.py +++ b/litellm/files/utils.py @@ -3,6 +3,22 @@ from typing import Optional from litellm.types.llms.openai import CreateFileRequest from litellm.types.utils import ExtractedFileData +# MIME types a .jsonl batch upload is plausibly labeled with. Clients are +# inconsistent (text/plain, application/json, octet-stream, ndjson, ...), so a +# batch file must not silently bypass the streaming path just because of its +# declared type. ``purpose == "batch"`` is the authoritative signal; non-JSONL +# content still fails loudly when the rows are parsed. +_BATCH_JSONL_CONTENT_TYPES = frozenset( + { + "application/jsonl", + "application/json", + "application/octet-stream", + "application/x-ndjson", + "application/x-jsonlines", + "text/plain", + } +) + class FilesAPIUtils: """ @@ -41,6 +57,7 @@ class FilesAPIUtils: @staticmethod def valid_content_type(content_type: Optional[str]) -> bool: """ - Check if the content type is valid + Whether the upload's MIME type is one a batch JSONL file is plausibly + sent as (see ``_BATCH_JSONL_CONTENT_TYPES``). """ - return content_type in set(["application/jsonl", "application/octet-stream"]) + return content_type in _BATCH_JSONL_CONTENT_TYPES diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index d55a715c7d3..5fa92ec544e 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1,3 +1,4 @@ +import asyncio import json import ssl from urllib.parse import parse_qs, urlencode, urlparse, urlunparse @@ -3527,10 +3528,13 @@ class BaseLLMHTTPHandler: 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. + """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. """ buf = bytearray() for piece in byte_iter: @@ -3538,7 +3542,8 @@ class BaseLLMHTTPHandler: while len(buf) >= chunk_size: yield bytes(buf[:chunk_size]) del buf[:chunk_size] - yield bytes(buf) + if buf: + yield bytes(buf) @staticmethod def _resumable_content_range(offset: int, data_len: int, is_final: bool) -> str: @@ -3676,7 +3681,15 @@ class BaseLLMHTTPHandler: offset = 0 pending: Optional[bytes] = None - for chunk in self._iter_resumable_chunks(stream.iter_bytes(), chunk_size): + # 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, 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 9d9af82c605..cd556c48b6b 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 @@ -483,6 +483,25 @@ class TestResumableUploadUrl: assert "uploadType=resumable" in url assert "uploadType=media" 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): cfg = VertexAIFilesConfig() request: CreateFileRequest = { @@ -539,9 +558,16 @@ class TestResumableChunking: pieces = list(BaseLLMHTTPHandler._iter_resumable_chunks(iter([b"x" * 10]), 4)) assert pieces == [b"xxxx", b"xxxx", b"xx"] - def test_exact_multiple_yields_trailing_empty_for_finalize(self): + 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", b""] + 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 @@ -640,11 +666,13 @@ class TestResumableUploadProtocol: assert bytes(state["received"]) == expected assert response.object == "file" - 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). + 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 - stream = _FixedBytesStream(b"a" * (chunk_size * 3)) + 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) @@ -657,8 +685,9 @@ class TestResumableUploadProtocol: timeout=None, ) - assert state["ranges"][-1] == f"bytes */{chunk_size * 3}" - assert bytes(state["received"]) == b"a" * (chunk_size * 3) + 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):