Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 25 additions & 4 deletions tests/e2e/batches/batch_cleanup.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import warnings
from builtins import ExceptionGroup
from collections.abc import Callable
from itertools import count
Expand All @@ -12,8 +13,9 @@
CLEANUP_DELAYS: Final = (1.0, 2.0, 4.0)
BATCH_TERMINAL_STATUSES: Final = frozenset({"completed", "failed", "expired", "cancelled"})
BATCH_PENDING_STATUSES: Final = frozenset({"validating", "in_progress", "finalizing", "cancelling"})
BATCH_CANCEL_TIMEOUT_SECONDS: Final = 660.0
BATCH_CANCEL_TIMEOUT_SECONDS: Final = 120.0
BATCH_CANCEL_POLL_SECONDS: Final = 10.0
FILE_IN_USE_REFUSAL: Final = "batch(es) in non-terminal state"


class BatchCleanupClient(Protocol):
Expand All @@ -26,6 +28,10 @@ def retrieve_batch(self, batch_id: str, *, key: str, provider: str | None = None
def cancel_batch(self, batch_id: str, *, key: str, provider: str | None = None) -> Result[BatchObject]: ...


class BatchCleanupLeftover(UserWarning):
pass


def cleanup_result[R: BaseModel](
action: Callable[[], Result[R]], *, wait: Callable[[float], None] = sleep
) -> Result[R]:
Expand Down Expand Up @@ -59,6 +65,13 @@ def cleanup_file(client: BatchCleanupClient, file_id: str, *, key: str, provider
result: Final = cleanup_result(delete)
if isinstance(result, UnknownApiError) and result.status_code == 404:
return
if isinstance(result, UnknownApiError) and result.status_code == 400 and FILE_IN_USE_REFUSAL in result.body:
warnings.warn(
f"Left file {file_id} in place: LiteLLM refused to delete it while a batch still references it",
BatchCleanupLeftover,
stacklevel=2,
)
return
deleted: Final = _require_cleanup_success(result, f"Delete file {file_id}")
assert deleted.deleted is True or (
deleted.deleted is None and is_managed_id(file_id) and deleted.id == file_id and deleted.object == "file"
Expand Down Expand Up @@ -120,9 +133,17 @@ def cleanup_batch(
)
if current.status == "cancelling" and not needs_terminal_state:
return
assert clock() < deadline, (
f"Batch {batch_id} cancellation did not finish within {BATCH_CANCEL_TIMEOUT_SECONDS}s"
)
if clock() >= deadline:
assert current.status == "cancelling", (
f"Batch {batch_id} cancellation did not finish within {BATCH_CANCEL_TIMEOUT_SECONDS}s, "
f"last status {current.status}"
)
warnings.warn(
f"Left batch {batch_id} cancelling after {BATCH_CANCEL_TIMEOUT_SECONDS}s for the provider to finish",
BatchCleanupLeftover,
stacklevel=2,
)
return
wait(BATCH_CANCEL_POLL_SECONDS)


Expand Down
70 changes: 63 additions & 7 deletions tests/e2e/batches/test_batch_cleanup.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,14 @@
from unittest.mock import Mock, call

import pytest
from batch_cleanup import BATCH_CANCEL_TIMEOUT_SECONDS, CLEANUP_DELAYS, cleanup_batch, cleanup_file, cleanup_result
from batch_cleanup import (
BATCH_CANCEL_TIMEOUT_SECONDS,
CLEANUP_DELAYS,
BatchCleanupLeftover,
cleanup_batch,
cleanup_file,
cleanup_result,
)
from batch_client import AZURE_FILE_EXPIRY_SECONDS, BatchObject, FileDeleteResponse, batch_upload_form
from capabilities import CAPABILITIES, Capability
from e2e_http import NetworkError, RateLimitedError, Result, Success, UnknownApiError
Expand All @@ -13,6 +20,10 @@

MANAGED_FILE_ID: Final = "bGl0ZWxsbV9wcm94eTtmaWxlLTE="
MANAGED_BATCH_ID: Final = "bGl0ZWxsbV9wcm94eTtiYXRjaC0x"
IN_USE_REFUSAL: Final = (
f'{{"error":{{"message":"Cannot delete file {MANAGED_FILE_ID}. The file is referenced by 1 batch(es) in '
f'non-terminal state: {MANAGED_BATCH_ID}: cancelling. ","type":"invalid_request_error","code":"400"}}}}'
)


class ExpectedCalls[T]:
Expand Down Expand Up @@ -125,6 +136,29 @@ def test_success_response_must_confirm_deletion(self) -> None:
cleanup_file(client, "file-1", key="test-key")
client.calls.assert_done()

def test_delete_refused_because_a_batch_still_references_the_file_is_left_and_reported(self) -> None:
client: Final = CleanupClient(
calls=ExpectedCalls((f"delete None {MANAGED_FILE_ID}",)),
files=(UnknownApiError(status_code=400, body=IN_USE_REFUSAL),),
)
with pytest.warns(BatchCleanupLeftover, match=MANAGED_FILE_ID):
cleanup_file(client, MANAGED_FILE_ID, key="test-key")
client.calls.assert_done()

@pytest.mark.parametrize(
"failure",
[
UnknownApiError(status_code=400, body="Invalid file id"),
UnknownApiError(status_code=409, body=IN_USE_REFUSAL),
UnknownApiError(status_code=501, body=IN_USE_REFUSAL),
],
)
def test_any_other_delete_failure_still_raises(self, failure: UnknownApiError) -> None:
client: Final = CleanupClient(calls=ExpectedCalls((f"delete None {MANAGED_FILE_ID}",)), files=(failure,))
with pytest.raises(AssertionError, match=f"Delete file {MANAGED_FILE_ID} failed: HTTP {failure.status_code}"):
cleanup_file(client, MANAGED_FILE_ID, key="test-key")
client.calls.assert_done()

def test_cleanup_is_idempotent_when_file_is_already_deleted(self) -> None:
client: Final = CleanupClient(
calls=ExpectedCalls(("delete azure file-1",)),
Expand Down Expand Up @@ -188,28 +222,50 @@ def test_cancelling_batch_is_polled_until_terminal_without_cancelling_again(self
client.calls.assert_done()
delays.assert_done()

def test_cancellation_timeout_is_reported_but_file_and_key_cleanup_still_run(self) -> None:
def test_batch_still_cancelling_at_the_deadline_and_its_input_file_are_left_and_reported(self) -> None:
client: Final = CleanupClient(
calls=ExpectedCalls(
(
f"retrieve None {MANAGED_BATCH_ID}",
f"retrieve None {MANAGED_BATCH_ID}",
"delete None file-1",
f"delete None {MANAGED_FILE_ID}",
"delete key test-key",
)
),
batches=(batch("cancelling"), batch("cancelling")),
files=(deleted_file(),),
files=(UnknownApiError(status_code=400, body=IN_USE_REFUSAL),),
)
times: Final = (0.0, BATCH_CANCEL_TIMEOUT_SECONDS)
ticks: Final[Callable[[], float]] = Mock(side_effect=times)
manager: Final = ResourceManager(client=client, strict_cleanup=True)
key: Final = manager.key()
manager.defer(lambda: cleanup_file(client, "file-1", key=key))
manager.defer(lambda: cleanup_file(client, MANAGED_FILE_ID, key=key))
manager.defer(lambda: cleanup_batch(client, MANAGED_BATCH_ID, key=key, clock=ticks))
with pytest.raises(ExceptionGroup) as caught:
with pytest.warns(BatchCleanupLeftover) as leftovers:
manager.teardown()
assert "cancellation did not finish" in str(caught.value.exceptions[0])
client.calls.assert_done()
messages: Final = tuple(str(warning.message) for warning in leftovers)
assert len(messages) == 2
assert MANAGED_BATCH_ID in messages[0] and "cancelling" in messages[0]
assert MANAGED_FILE_ID in messages[1]

@pytest.mark.parametrize(
"last, reported",
[
(batch("in_progress"), f"did not finish within {BATCH_CANCEL_TIMEOUT_SECONDS}s, last status in_progress"),
(UnknownApiError(status_code=403, body="forbidden"), "after cancellation failed: HTTP 403"),
],
)
def test_anything_but_still_cancelling_at_the_deadline_still_fails(
self, last: Result[BatchObject], reported: str
) -> None:
client: Final = CleanupClient(
calls=ExpectedCalls((f"retrieve None {MANAGED_BATCH_ID}",) * 2), batches=(batch("cancelling"), last)
)
times: Final = (0.0, BATCH_CANCEL_TIMEOUT_SECONDS)
ticks: Final[Callable[[], float]] = Mock(side_effect=times)
with pytest.raises(AssertionError, match=reported):
cleanup_batch(client, MANAGED_BATCH_ID, key="test-key", clock=ticks)
client.calls.assert_done()

@pytest.mark.parametrize("status", ["completed", "failed", "expired", "cancelled"])
Expand Down
6 changes: 3 additions & 3 deletions tests/e2e/batches/test_batches_e2e.py
Original file line number Diff line number Diff line change
Expand Up @@ -93,11 +93,11 @@ class _GovCloudBedrockRecord(BaseModel):
model_input: _GovCloudBedrockInput = Field(alias="modelInput")


# Azure / Vertex cancel and the pre-cancel re-retrieve are provider-side flakes
# Vertex cancel and the pre-cancel re-retrieve are provider-side flakes
# (connection refused, brief 500s) and the registry only has one basic cell per
# provider (shared across scenarios). Create + retrieve already prove routing;
# cancel is still deferred for cleanup, just not asserted for these two.
_CANCEL_ASSERTED_PROVIDERS = frozenset({"openai", "bedrock"})
# cancel is still deferred for cleanup, just not asserted for Vertex.
_CANCEL_ASSERTED_PROVIDERS = frozenset({"openai", "azure", "bedrock"})

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Azure cancellation flakes return Adding Azure to this set makes every Azure lifecycle case retrieve and cancel the batch during the test. Azure was previously excluded because these provider calls can fail intermittently. If the connection is refused, the helpers do not retry it, so the test fails even when batch creation and retrieval worked and cleanup could proceed.



def _transient_status(status_code: int) -> bool:
Expand Down
2 changes: 1 addition & 1 deletion tests/e2e/ui/tests/mcp/mcpTools.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ test.describe("MCP Tools", () => {

// Non-empty would still pass if the proxy returned some other server's tools.
await expect(toolCard(toolList, TOOL_NAME)).toBeVisible();
await expect(toolCard(toolList, "ask_question")).toBeVisible();
await expect(toolCard(toolList, "ask_wiki_question")).toBeVisible();
await expect(toolCard(toolList, "read_wiki_contents")).toBeVisible();

// No other tool's name or description contains this string, so exactly one card survives.
Expand Down
Loading