mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
fix(vertex_ai): bound and pin batch output shard reads, clamp zero starting replicas
The shard fan-out now only triggers on the exact Vertex unmanaged-batch output layout (prediction-custom-unmanaged-model-<ts>/prediction.results- 00000-of-NNNNN), is capped at 512 shards, and fetches sequentially, so a crafted upload filename cannot fan out GCS requests. A scale-to-zero endpoint's minReplicaCount 0 is clamped to 1 for the batch job. ruff format applied to the touched files.
This commit is contained in:
mubashir1osmani
parent
1818c39fed
commit
5d0c552694
7 changed files with 144 additions and 73 deletions
|
|
@ -380,9 +380,7 @@ class VertexAIBatchPrediction(VertexLLM):
|
|||
model=deployed_model_resource,
|
||||
vertex_location=vertex_location,
|
||||
)
|
||||
model_fetched: Final[_FetchedResponseView] = {
|
||||
"response": safe_get(sync_handler, model_url, headers=headers)
|
||||
}
|
||||
model_fetched: Final[_FetchedResponseView] = {"response": safe_get(sync_handler, model_url, headers=headers)}
|
||||
model_response: Final = model_fetched["response"]
|
||||
if model_response.status_code != 200:
|
||||
raise VertexAIError(
|
||||
|
|
@ -410,10 +408,12 @@ class VertexAIBatchPrediction(VertexLLM):
|
|||
"spec to size the batch replicas from"
|
||||
),
|
||||
)
|
||||
# A scale-to-zero online endpoint reports minReplicaCount 0, but a batch job must start
|
||||
# at least one replica.
|
||||
batch_resources: Final[BatchDedicatedResources] = {
|
||||
"machineSpec": machine_spec,
|
||||
"startingReplicaCount": online_resources.get("minReplicaCount", 1),
|
||||
"maxReplicaCount": online_resources.get("maxReplicaCount", 1),
|
||||
"startingReplicaCount": max(online_resources.get("minReplicaCount", 1), 1),
|
||||
"maxReplicaCount": max(online_resources.get("maxReplicaCount", 1), 1),
|
||||
}
|
||||
unmanaged: Final[UnmanagedContainerModel] = {"containerSpec": container_spec}
|
||||
resolved: Final[VertexAIBatchPredictionJob] = {
|
||||
|
|
@ -471,7 +471,9 @@ class VertexAIBatchPrediction(VertexLLM):
|
|||
"""Return the base url for the vertex garden models"""
|
||||
# POST https://LOCATION-aiplatform.googleapis.com/v1/projects/PROJECT_ID/locations/LOCATION/batchPredictionJobs
|
||||
base_url: Final = get_vertex_base_url(vertex_location)
|
||||
return f"{base_url}/{vertex_api_version}/projects/{vertex_project}/locations/{vertex_location}/batchPredictionJobs"
|
||||
return (
|
||||
f"{base_url}/{vertex_api_version}/projects/{vertex_project}/locations/{vertex_location}/batchPredictionJobs"
|
||||
)
|
||||
|
||||
def retrieve_batch(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -97,7 +97,13 @@ class VertexAIFilesHandler(GCSBucketBase):
|
|||
allow_legacy_cloud_file_ids=should_allow_legacy_cloud_file_ids(litellm_params),
|
||||
)
|
||||
|
||||
_SHARDED_RESULTS_PATTERN: Final = re.compile(r"^(?P<stem>.*-)(?P<index>\d{5})(?P<sep>-of-)(?P<total>\d{5})$")
|
||||
# Only the exact directory/file layout Vertex writes for unmanaged-container batch outputs,
|
||||
# pinned to shard zero: the object path is derived from a caller-controlled file id, so a
|
||||
# looser pattern would let a crafted upload filename trigger shard fan-out.
|
||||
_SHARDED_RESULTS_PATTERN: Final = re.compile(
|
||||
r"^(?P<stem>.*/prediction-custom-unmanaged-model-[^/]+/prediction\.results-)00000(?P<sep>-of-)(?P<total>\d{5})$"
|
||||
)
|
||||
_MAX_RESULT_SHARDS: Final = 512
|
||||
|
||||
async def _download_all_result_shards(
|
||||
self,
|
||||
|
|
@ -119,14 +125,20 @@ class VertexAIFilesHandler(GCSBucketBase):
|
|||
total_shards: Final = int(shard_match["total"])
|
||||
if total_shards <= 1:
|
||||
return first_shard
|
||||
remaining: Final = await asyncio.gather(
|
||||
*(
|
||||
self.download_gcs_object(
|
||||
if total_shards > self._MAX_RESULT_SHARDS:
|
||||
raise ValueError(
|
||||
f"Vertex batch output claims {total_shards} shards, above the supported maximum "
|
||||
f"of {self._MAX_RESULT_SHARDS}"
|
||||
)
|
||||
# Sequential fetch keeps memory and connection use bounded by one shard at a time.
|
||||
remaining: Final = tuple(
|
||||
[
|
||||
await self.download_gcs_object(
|
||||
object_name=f"{shard_match['stem']}{index:05d}{shard_match['sep']}{shard_match['total']}",
|
||||
standard_callback_dynamic_params=standard_callback_dynamic_params,
|
||||
)
|
||||
for index in range(1, total_shards)
|
||||
)
|
||||
]
|
||||
)
|
||||
if any(shard is None for shard in remaining):
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -666,7 +666,9 @@ def _custom_endpoint_row_to_openai_batch_output_row(row: Mapping[str, object]) -
|
|||
"""
|
||||
key: Final = row.get("key")
|
||||
instance: Final = row.get("instance")
|
||||
tagged_custom_id: Final = instance.get(VERTEX_CUSTOM_ENDPOINT_KEY_FIELD, "") if isinstance(instance, Mapping) else ""
|
||||
tagged_custom_id: Final = (
|
||||
instance.get(VERTEX_CUSTOM_ENDPOINT_KEY_FIELD, "") if isinstance(instance, Mapping) else ""
|
||||
)
|
||||
custom_id: Final = str(key if key is not None else tagged_custom_id)
|
||||
|
||||
prediction: Final = row.get("prediction")
|
||||
|
|
@ -833,8 +835,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
|
|||
if purpose == "batch" and custom_endpoint_id is not None:
|
||||
safe_endpoint_id: Final = sanitize_cloud_object_path(custom_endpoint_id, fallback="endpoint")
|
||||
return (
|
||||
f"{VERTEX_AI_MANAGED_GCS_PREFIX}{VERTEX_CUSTOM_ENDPOINT_GCS_SEGMENT}/"
|
||||
f"{safe_endpoint_id}/{uuid.uuid4()}"
|
||||
f"{VERTEX_AI_MANAGED_GCS_PREFIX}{VERTEX_CUSTOM_ENDPOINT_GCS_SEGMENT}/{safe_endpoint_id}/{uuid.uuid4()}"
|
||||
)
|
||||
if purpose == "batch":
|
||||
## 1. If jsonl, derive the object name from the deployment model (or the first entry's)
|
||||
|
|
@ -881,9 +882,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
|
|||
configured_model: Final = litellm_params.get("model")
|
||||
deployment_api_base: Final = litellm_params.get("api_base")
|
||||
custom_endpoint_id: Final = (
|
||||
get_custom_endpoint_id_from_api_base(
|
||||
deployment_api_base if isinstance(deployment_api_base, str) else None
|
||||
)
|
||||
get_custom_endpoint_id_from_api_base(deployment_api_base if isinstance(deployment_api_base, str) else None)
|
||||
if litellm_params.get("custom_endpoint")
|
||||
else None
|
||||
)
|
||||
|
|
|
|||
|
|
@ -47,11 +47,7 @@ PROJECT = "my-project"
|
|||
LOCATION = "us-central1"
|
||||
BATCH_ID = "3814889423749775360"
|
||||
|
||||
CREATE_DATA = {
|
||||
"input_file_id": (
|
||||
"gs://bucket/publishers/google/models/gemini-1.5-flash-001/file-uuid"
|
||||
)
|
||||
}
|
||||
CREATE_DATA = {"input_file_id": ("gs://bucket/publishers/google/models/gemini-1.5-flash-001/file-uuid")}
|
||||
|
||||
|
||||
def _vertex_job_response(state: str = "JOB_STATE_SUCCEEDED") -> dict:
|
||||
|
|
@ -104,8 +100,7 @@ def test_create_vertex_batch_url():
|
|||
h = _make_handler()
|
||||
url = h.create_vertex_batch_url(vertex_location=LOCATION, vertex_project=PROJECT)
|
||||
assert url == (
|
||||
f"https://{LOCATION}-aiplatform.googleapis.com/v1/projects/{PROJECT}"
|
||||
f"/locations/{LOCATION}/batchPredictionJobs"
|
||||
f"https://{LOCATION}-aiplatform.googleapis.com/v1/projects/{PROJECT}/locations/{LOCATION}/batchPredictionJobs"
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -206,9 +201,7 @@ def test_create_batch_sync_does_not_resolve_publisher_models():
|
|||
|
||||
|
||||
ENDPOINT_ID = "7768560373388541952"
|
||||
ENDPOINT_CREATE_DATA = {
|
||||
"input_file_id": f"gs://bucket/litellm-vertex-files/endpoints/{ENDPOINT_ID}/file-uuid"
|
||||
}
|
||||
ENDPOINT_CREATE_DATA = {"input_file_id": f"gs://bucket/litellm-vertex-files/endpoints/{ENDPOINT_ID}/file-uuid"}
|
||||
TUNED_MODEL_RESOURCE = f"projects/{PROJECT}/locations/{LOCATION}/models/1234509876"
|
||||
|
||||
|
||||
|
|
@ -217,9 +210,7 @@ def _endpoint_get_response(deployed_models: list | None = None) -> MagicMock:
|
|||
resp.status_code = 200
|
||||
resp.json.return_value = {
|
||||
"name": f"projects/{PROJECT}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}",
|
||||
"deployedModels": (
|
||||
deployed_models if deployed_models is not None else [{"model": TUNED_MODEL_RESOURCE}]
|
||||
),
|
||||
"deployedModels": (deployed_models if deployed_models is not None else [{"model": TUNED_MODEL_RESOURCE}]),
|
||||
}
|
||||
return resp
|
||||
|
||||
|
|
@ -275,7 +266,7 @@ CUSTOM_ENDPOINT_API_BASE = (
|
|||
)
|
||||
|
||||
|
||||
def _custom_endpoint_get_response() -> MagicMock:
|
||||
def _custom_endpoint_get_response(min_replica_count: int = 1) -> MagicMock:
|
||||
resp = MagicMock()
|
||||
resp.status_code = 200
|
||||
resp.json.return_value = {
|
||||
|
|
@ -283,7 +274,11 @@ def _custom_endpoint_get_response() -> MagicMock:
|
|||
"deployedModels": [
|
||||
{
|
||||
"model": CONTAINER_MODEL_RESOURCE,
|
||||
"dedicatedResources": {"machineSpec": MACHINE_SPEC, "minReplicaCount": 1, "maxReplicaCount": 2},
|
||||
"dedicatedResources": {
|
||||
"machineSpec": MACHINE_SPEC,
|
||||
"minReplicaCount": min_replica_count,
|
||||
"maxReplicaCount": 2,
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
|
|
@ -378,6 +373,37 @@ def test_create_batch_sync_custom_endpoint_without_container_spec_raises_400():
|
|||
client.post.assert_not_called()
|
||||
|
||||
|
||||
def test_create_batch_sync_custom_endpoint_scale_to_zero_starts_one_replica():
|
||||
"""A scale-to-zero online endpoint reports minReplicaCount 0, which Vertex rejects as a batch
|
||||
startingReplicaCount; the job must clamp to at least one replica."""
|
||||
h = _make_handler()
|
||||
client = MagicMock()
|
||||
client.post.return_value = _http_response()
|
||||
|
||||
with (
|
||||
patch(f"{HMOD}._get_httpx_client", return_value=client),
|
||||
patch(
|
||||
f"{HMOD}.safe_get",
|
||||
side_effect=[_custom_endpoint_get_response(min_replica_count=0), _container_model_get_response()],
|
||||
),
|
||||
):
|
||||
h.create_batch(
|
||||
_is_async=False,
|
||||
create_batch_data=CUSTOM_ENDPOINT_CREATE_DATA,
|
||||
api_base=CUSTOM_ENDPOINT_API_BASE,
|
||||
vertex_credentials=None,
|
||||
vertex_project=PROJECT,
|
||||
vertex_location=LOCATION,
|
||||
timeout=600.0,
|
||||
max_retries=None,
|
||||
custom_endpoint=True,
|
||||
)
|
||||
|
||||
sent = json.loads(client.post.call_args.kwargs["data"])
|
||||
assert sent["dedicatedResources"]["startingReplicaCount"] == 1
|
||||
assert sent["dedicatedResources"]["maxReplicaCount"] == 2
|
||||
|
||||
|
||||
def test_create_batch_sync_custom_endpoint_rejects_multi_deployment_endpoint():
|
||||
"""An endpoint behind a traffic split has no single container to replicate; index-zero
|
||||
selection could run a different container than online traffic."""
|
||||
|
|
@ -628,9 +654,7 @@ def test_create_batch_sync_httpstatuserror_propagates():
|
|||
client = MagicMock()
|
||||
request = httpx.Request("POST", "https://x/batchPredictionJobs")
|
||||
err_response = httpx.Response(status_code=500, request=request, text="boom")
|
||||
client.post.side_effect = httpx.HTTPStatusError(
|
||||
"boom", request=request, response=err_response
|
||||
)
|
||||
client.post.side_effect = httpx.HTTPStatusError("boom", request=request, response=err_response)
|
||||
|
||||
with patch(f"{HMOD}._get_httpx_client", return_value=client):
|
||||
with pytest.raises(httpx.HTTPStatusError):
|
||||
|
|
@ -780,9 +804,7 @@ def test_retrieve_batch_sync_invokes_logging_pre_call():
|
|||
|
||||
logging_obj.pre_call.assert_called_once()
|
||||
_, kwargs = logging_obj.pre_call.call_args
|
||||
assert kwargs["additional_args"]["api_base"].endswith(
|
||||
f"/batchPredictionJobs/{BATCH_ID}"
|
||||
)
|
||||
assert kwargs["additional_args"]["api_base"].endswith(f"/batchPredictionJobs/{BATCH_ID}")
|
||||
|
||||
|
||||
# =========================================================================== #
|
||||
|
|
@ -853,9 +875,7 @@ def test_list_batches_sync_omits_unset_pagination_params():
|
|||
def test_list_batches_async_returns_coroutine():
|
||||
h = _make_handler()
|
||||
async_client = MagicMock()
|
||||
async_client.get = AsyncMock(
|
||||
return_value=_http_response(json_body=_list_response())
|
||||
)
|
||||
async_client.get = AsyncMock(return_value=_http_response(json_body=_list_response()))
|
||||
sync_client = MagicMock()
|
||||
|
||||
with (
|
||||
|
|
@ -910,9 +930,7 @@ def test_cancel_batch_sync_posts_cancel_then_retrieves():
|
|||
h = _make_handler()
|
||||
client = MagicMock()
|
||||
client.post.return_value = _http_response(json_body={})
|
||||
client.get.return_value = _http_response(
|
||||
json_body=_vertex_job_response(state="JOB_STATE_CANCELLED")
|
||||
)
|
||||
client.get.return_value = _http_response(json_body=_vertex_job_response(state="JOB_STATE_CANCELLED"))
|
||||
|
||||
with patch(f"{HMOD}._get_httpx_client", return_value=client):
|
||||
out = h.cancel_batch(
|
||||
|
|
@ -944,9 +962,7 @@ def test_cancel_batch_async_returns_coroutine_posts_then_retrieves():
|
|||
async_client = MagicMock()
|
||||
async_client.post = AsyncMock(return_value=_http_response(json_body={}))
|
||||
async_client.get = AsyncMock(
|
||||
return_value=_http_response(
|
||||
json_body=_vertex_job_response(state="JOB_STATE_CANCELLED")
|
||||
)
|
||||
return_value=_http_response(json_body=_vertex_job_response(state="JOB_STATE_CANCELLED"))
|
||||
)
|
||||
|
||||
with (
|
||||
|
|
@ -1004,9 +1020,7 @@ def test_cancel_batch_sync_proxy_url_without_cancel_suffix_uses_rsplit_branch():
|
|||
)
|
||||
client = MagicMock()
|
||||
client.post.return_value = _http_response(json_body={})
|
||||
client.get.return_value = _http_response(
|
||||
json_body=_vertex_job_response(state="JOB_STATE_CANCELLED")
|
||||
)
|
||||
client.get.return_value = _http_response(json_body=_vertex_job_response(state="JOB_STATE_CANCELLED"))
|
||||
|
||||
with patch(f"{HMOD}._get_httpx_client", return_value=client):
|
||||
out = h.cancel_batch(
|
||||
|
|
@ -1032,9 +1046,7 @@ def test_cancel_batch_sync_httpstatuserror_logged_and_reraised():
|
|||
client = MagicMock()
|
||||
request = httpx.Request("POST", "https://x/batchPredictionJobs/1:cancel")
|
||||
err_response = httpx.Response(status_code=502, request=request, text="bad gw")
|
||||
client.post.side_effect = httpx.HTTPStatusError(
|
||||
"boom", request=request, response=err_response
|
||||
)
|
||||
client.post.side_effect = httpx.HTTPStatusError("boom", request=request, response=err_response)
|
||||
|
||||
with patch(f"{HMOD}._get_httpx_client", return_value=client):
|
||||
with pytest.raises(httpx.HTTPStatusError):
|
||||
|
|
@ -1056,11 +1068,7 @@ def test_create_batch_async_httpstatuserror_logged_and_reraised():
|
|||
async_client = MagicMock()
|
||||
request = httpx.Request("POST", "https://x/batchPredictionJobs")
|
||||
err_response = httpx.Response(status_code=500, request=request, text="boom")
|
||||
async_client.post = AsyncMock(
|
||||
side_effect=httpx.HTTPStatusError(
|
||||
"boom", request=request, response=err_response
|
||||
)
|
||||
)
|
||||
async_client.post = AsyncMock(side_effect=httpx.HTTPStatusError("boom", request=request, response=err_response))
|
||||
|
||||
with (
|
||||
patch(f"{HMOD}._get_httpx_client", return_value=MagicMock()),
|
||||
|
|
@ -1167,9 +1175,7 @@ def test_async_cancel_batch_httpstatuserror_and_retrieve_non_200():
|
|||
async_client = MagicMock()
|
||||
request = httpx.Request("POST", "https://x/batchPredictionJobs/1:cancel")
|
||||
err_response = httpx.Response(status_code=502, request=request, text="bad")
|
||||
async_client.post = AsyncMock(
|
||||
side_effect=httpx.HTTPStatusError("boom", request=request, response=err_response)
|
||||
)
|
||||
async_client.post = AsyncMock(side_effect=httpx.HTTPStatusError("boom", request=request, response=err_response))
|
||||
async_client.get = AsyncMock()
|
||||
with (
|
||||
patch(f"{HMOD}._get_httpx_client", return_value=MagicMock()),
|
||||
|
|
|
|||
|
|
@ -36,8 +36,7 @@ INPUT_FILE = (
|
|||
|
||||
ENDPOINT_ID = "7768560373388541952"
|
||||
ENDPOINT_INPUT_FILE = (
|
||||
f"gs://litellm-testing-bucket/litellm-vertex-files/endpoints/{ENDPOINT_ID}/"
|
||||
"e9412502-2c91-42a6-8e61-f5c294cc0fc8"
|
||||
f"gs://litellm-testing-bucket/litellm-vertex-files/endpoints/{ENDPOINT_ID}/e9412502-2c91-42a6-8e61-f5c294cc0fc8"
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -149,12 +149,12 @@ class TestVertexAIFilesHandler:
|
|||
only that shard silently drops the rest of the batch (LIT-7387)."""
|
||||
file_id = (
|
||||
"gs%3A%2F%2Ftest-bucket%2Flitellm-vertex-files%2Fcustom-endpoints%2F123%2F"
|
||||
"out%2Fprediction.results-00000-of-00003"
|
||||
"prediction-custom-unmanaged-model-2026_09_09T13_00_00_000Z%2Fprediction.results-00000-of-00003"
|
||||
)
|
||||
shards = {
|
||||
"litellm-vertex-files/custom-endpoints/123/out/prediction.results-00000-of-00003": b'{"a": 1}\n',
|
||||
"litellm-vertex-files/custom-endpoints/123/out/prediction.results-00001-of-00003": b'{"b": 2}\n',
|
||||
"litellm-vertex-files/custom-endpoints/123/out/prediction.results-00002-of-00003": b'{"c": 3}',
|
||||
"litellm-vertex-files/custom-endpoints/123/prediction-custom-unmanaged-model-2026_09_09T13_00_00_000Z/prediction.results-00000-of-00003": b'{"a": 1}\n',
|
||||
"litellm-vertex-files/custom-endpoints/123/prediction-custom-unmanaged-model-2026_09_09T13_00_00_000Z/prediction.results-00001-of-00003": b'{"b": 2}\n',
|
||||
"litellm-vertex-files/custom-endpoints/123/prediction-custom-unmanaged-model-2026_09_09T13_00_00_000Z/prediction.results-00002-of-00003": b'{"c": 3}',
|
||||
}
|
||||
|
||||
async def fake_download(object_name: str, **kwargs):
|
||||
|
|
@ -184,7 +184,7 @@ class TestVertexAIFilesHandler:
|
|||
async def test_afile_content_single_shard_unchanged(self):
|
||||
file_id = (
|
||||
"gs%3A%2F%2Ftest-bucket%2Flitellm-vertex-files%2Fcustom-endpoints%2F123%2F"
|
||||
"out%2Fprediction.results-00000-of-00001"
|
||||
"prediction-custom-unmanaged-model-2026_09_09T13_00_00_000Z%2Fprediction.results-00000-of-00001"
|
||||
)
|
||||
with (
|
||||
patch.object(self.handler, "download_gcs_object", new_callable=AsyncMock) as mock_download,
|
||||
|
|
@ -208,6 +208,62 @@ class TestVertexAIFilesHandler:
|
|||
assert result.response.content == b'{"a": 1}\n'
|
||||
mock_download.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_afile_content_crafted_upload_filename_does_not_fan_out(self):
|
||||
"""The object path comes from a caller-controlled file id, so an ordinary upload whose
|
||||
name mimics the shard suffix must not trigger shard fetches (99,998 GCS requests from one
|
||||
crafted 'prediction.results-00000-of-99999' filename)."""
|
||||
file_id = "gs%3A%2F%2Ftest-bucket%2Flitellm-vertex-files%2Fuploads%2Fabc-prediction.results-00000-of-99999"
|
||||
with (
|
||||
patch.object(self.handler, "download_gcs_object", new_callable=AsyncMock) as mock_download,
|
||||
patch.object(
|
||||
self.handler,
|
||||
"get_gcs_logging_config",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_mock_gcs_logging_config(),
|
||||
),
|
||||
):
|
||||
mock_download.return_value = b"tiny"
|
||||
result = await self.handler.afile_content(
|
||||
file_content_request=FileContentRequest(file_id=file_id, extra_headers=None, extra_body=None),
|
||||
vertex_credentials=None,
|
||||
vertex_project="test-project",
|
||||
vertex_location="us-central1",
|
||||
timeout=60.0,
|
||||
max_retries=3,
|
||||
)
|
||||
|
||||
assert result.response.content == b"tiny"
|
||||
mock_download.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_afile_content_shard_count_above_cap_raises(self):
|
||||
file_id = (
|
||||
"gs%3A%2F%2Ftest-bucket%2Flitellm-vertex-files%2Fcustom-endpoints%2F123%2F"
|
||||
"prediction-custom-unmanaged-model-2026_09_09T13_00_00_000Z%2Fprediction.results-00000-of-99999"
|
||||
)
|
||||
with (
|
||||
patch.object(self.handler, "download_gcs_object", new_callable=AsyncMock) as mock_download,
|
||||
patch.object(
|
||||
self.handler,
|
||||
"get_gcs_logging_config",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_mock_gcs_logging_config(),
|
||||
),
|
||||
):
|
||||
mock_download.return_value = b"tiny"
|
||||
with pytest.raises(ValueError, match="shards"):
|
||||
await self.handler.afile_content(
|
||||
file_content_request=FileContentRequest(file_id=file_id, extra_headers=None, extra_body=None),
|
||||
vertex_credentials=None,
|
||||
vertex_project="test-project",
|
||||
vertex_location="us-central1",
|
||||
timeout=60.0,
|
||||
max_retries=3,
|
||||
)
|
||||
|
||||
mock_download.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_afile_content_missing_file_id(self):
|
||||
"""Test async file content retrieval with missing file_id"""
|
||||
|
|
@ -247,8 +303,7 @@ class TestVertexAIFilesHandler:
|
|||
with pytest.raises(
|
||||
ValueError,
|
||||
match=re.escape(
|
||||
"Failed to download file from GCS: "
|
||||
"gs://test-bucket/litellm-vertex-files/uploads/abc-test-file.txt"
|
||||
"Failed to download file from GCS: gs://test-bucket/litellm-vertex-files/uploads/abc-test-file.txt"
|
||||
),
|
||||
):
|
||||
await self.handler.afile_content(
|
||||
|
|
|
|||
|
|
@ -193,9 +193,7 @@ class TestBatchObjectNaming:
|
|||
def test_should_store_fine_tuned_endpoint_under_endpoints_path(self, config):
|
||||
"""A numeric endpoint id must not be filed under publishers/google/models/gemini/<id>,
|
||||
which the batch transformation later mangles into a nonexistent publisher model (LIT-6899)."""
|
||||
object_name = config._get_gcs_object_name_from_batch_jsonl(
|
||||
[{"body": {"model": "gemini/7768560373388541952"}}]
|
||||
)
|
||||
object_name = config._get_gcs_object_name_from_batch_jsonl([{"body": {"model": "gemini/7768560373388541952"}}])
|
||||
assert object_name.startswith("litellm-vertex-files/endpoints/7768560373388541952/")
|
||||
assert "publishers" not in object_name
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue