test(e2e): drop two batch waits that never finish in time (#51882)

cleanup_batch waited up to 120s for a managed batch that was still
cancelling to reach a terminal state. Across 55 saved litellm-e2e builds it
reached one in time once, so test_batch_lifecycle[openai-unified] spent 126s
and TestBatchSecondHop 140s on every run, and both then left the input file
in place anyway. A managed batch still cancelling is now left alone, the same
as a raw provider batch, and the input file cleanup still reports the
leftover.

TestBatchTerminalState listed the marker batches every 10s for up to 300s.
Each e2e stack starts with an empty database, so the list only ever holds
this run's own marker, which never completed inside the window in any of
those builds, and the test took 303s for its documented vacuous pass. It now
lists once.
This commit is contained in:
yuneng-jiang 2026-11-14 21:41:43 -08:00 • committed by GitHub
parent 26807a9925
commit 91fc6c2071
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 46 additions and 60 deletions

View file

@ -87,7 +87,6 @@ def cleanup_batch(
wait: Callable[[float], None] = sleep,
clock: Callable[[], float] = monotonic,
) -> None:
needs_terminal_state: Final = is_managed_id(batch_id)
fetched: Final = _require_cleanup_success(
cleanup_result(lambda: client.retrieve_batch(batch_id, key=key, provider=provider)),
f"Retrieve batch {batch_id} for cleanup",
@ -96,13 +95,9 @@ def cleanup_batch(
if delete_output_files:
_cleanup_batch_outputs(client, fetched, key=key, provider=provider)
return
if fetched.status == "cancelling" and not needs_terminal_state:
if fetched.status == "cancelling":
return
result: Final = (
Success(status_code=200, data=fetched)
if fetched.status == "cancelling"
else cleanup_result(lambda: client.cancel_batch(batch_id, key=key, provider=provider))
)
result: Final = cleanup_result(lambda: client.cancel_batch(batch_id, key=key, provider=provider))
conflicted: Final = isinstance(result, UnknownApiError) and result.status_code in {400, 409}
if not conflicted:
cancelled: Final = _require_cleanup_success(result, f"Cancel batch {batch_id}")
@ -113,7 +108,7 @@ def cleanup_batch(
if delete_output_files:
_cleanup_batch_outputs(client, cancelled, key=key, provider=provider)
return
if cancelled.status == "cancelling" and not needs_terminal_state:
if cancelled.status == "cancelling":
return
deadline: Final = clock() + BATCH_CANCEL_TIMEOUT_SECONDS
for current in (
@ -130,19 +125,12 @@ def cleanup_batch(
assert current.status in ({"cancelling"} if conflicted else BATCH_PENDING_STATUSES), (
f"Cancel batch {batch_id} left status {current.status}"
)
if current.status == "cancelling" and not needs_terminal_state:
return
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",
UserWarning,
stacklevel=2 + STEP_FRAMES,
)
if current.status == "cancelling":
return
assert clock() < deadline, (
f"Batch {batch_id} cancellation did not finish within {BATCH_CANCEL_TIMEOUT_SECONDS}s, "
f"last status {current.status}"
)
wait(BATCH_CANCEL_POLL_SECONDS)

View file

@ -1768,8 +1768,6 @@ class BatchOutputLine(BaseModel):
TERMINAL_MARKER_KEY = "litellm_e2e_suite"
TERMINAL_MARKER_VALUE = "batches-terminal-baton"
TERMINAL_POLL_SECONDS = 300.0
TERMINAL_POLL_INTERVAL_SECONDS = 10.0
TERMINAL_LIST_LIMIT = 100
TERMINAL_BAND_MIN_AGE_SECONDS = 25 * 3600
TERMINAL_BAND_MAX_AGE_SECONDS = 73 * 3600
@ -1786,20 +1784,12 @@ def _marker_batches(client: BatchClient, key: str) -> list[BatchObject]:
]
def _await_completed_marker(
client: BatchClient, key: str
) -> tuple[BatchObject | None, list[BatchObject]]:
deadline = time.monotonic() + TERMINAL_POLL_SECONDS
while True:
markers = _marker_batches(client, key)
completed = max(
(b for b in markers if b.status == "completed"),
key=lambda b: b.created_at or 0,
default=None,
)
if completed is not None or time.monotonic() >= deadline:
return completed, markers
time.sleep(TERMINAL_POLL_INTERVAL_SECONDS)
def _newest_completed_marker(markers: list[BatchObject]) -> BatchObject | None:
return max(
(b for b in markers if b.status == "completed"),
key=lambda b: b.created_at or 0,
default=None,
)
def _assert_aged_markers_terminal(markers: list[BatchObject]) -> None:
@ -1822,13 +1812,13 @@ class TestBatchTerminalState:
Each run submits a 1-line marker batch (stable metadata key/value plus a
per-run field) and never cancels or deletes it: the marker is the baton the
next run picks up. Polling is list-only for up to 5 minutes because a
retrieve of a non-terminal batch books a $0 spend row whose request_id then
blocks the real-cost row (skip_duplicates); the single retrieve happens only
once a completed marker exists. The assertion target is the newest completed
marker from ANY run, so on the 6h stage cadence the full assertions are
deterministic from run 2 onward. On a cold start (no marker has ever
completed within the poll budget) the test passes on the submission
next run picks up. The markers are read with one list call, never a retrieve,
because a retrieve of a non-terminal batch books a $0 spend row whose
request_id then blocks the real-cost row (skip_duplicates); the single
retrieve happens only once a completed marker exists. The assertion target
is the newest completed marker from ANY run, so on the 6h stage cadence the
full assertions are deterministic from run 2 onward. On a cold start (no
completed marker in the list) the test passes on the submission
assertions alone: that is a documented vacuous pass, not a skip, and this
run's marker becomes the next run's target. Markers aged past OpenAI's 24h
completion window (25h-73h band, within the newest list page) must be
@ -1883,7 +1873,8 @@ class TestBatchTerminalState:
f"create dropped the marker metadata: {submitted.metadata!r}"
)
completed, markers = _await_completed_marker(client, key)
markers = _marker_batches(client, key)
completed = _newest_completed_marker(markers)
_assert_aged_markers_terminal(markers)
if completed is None:
return

View file

@ -211,42 +211,41 @@ class TestCleanupRetries:
class TestBatchCancellation:
def test_cancelling_batch_is_polled_until_terminal_without_cancelling_again(self) -> None:
@pytest.mark.parametrize("batch_id", ["batch-1", MANAGED_BATCH_ID])
def test_batch_already_cancelling_is_left_without_cancelling_again_or_waiting(self, batch_id: str) -> None:
client: Final = CleanupClient(
calls=ExpectedCalls((f"retrieve None {MANAGED_BATCH_ID}",) * 3),
batches=(batch("cancelling"), batch("cancelling"), batch("cancelled")),
calls=ExpectedCalls((f"retrieve None {batch_id}",)), batches=(batch("cancelling"),)
)
delays: Final = ExpectedCalls((10.0,))
cleanup_batch(client, MANAGED_BATCH_ID, key="test-key", wait=delays)
delays: Final = ExpectedCalls[float](())
cleanup_batch(client, batch_id, key="test-key", wait=delays)
client.calls.assert_done()
delays.assert_done()
def test_batch_still_cancelling_at_the_deadline_and_its_input_file_are_left_and_reported(self) -> None:
def test_managed_batch_left_cancelling_leaves_its_input_file_and_reports_it(self) -> None:
client: Final = CleanupClient(
calls=ExpectedCalls(
(
f"retrieve None {MANAGED_BATCH_ID}",
f"retrieve None {MANAGED_BATCH_ID}",
f"cancel None {MANAGED_BATCH_ID}",
f"delete None {MANAGED_FILE_ID}",
"delete key test-key",
)
),
batches=(batch("cancelling"), batch("cancelling")),
batches=(batch("in_progress"),),
cancellations=(batch("cancelling"),),
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)
delays: Final = ExpectedCalls[float](())
manager: Final = ResourceManager(client=client, strict_cleanup=True)
key: Final = manager.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))
manager.defer(lambda: cleanup_batch(client, MANAGED_BATCH_ID, key=key, wait=delays))
with pytest.warns(UserWarning, match="^Left ") as leftovers:
manager.teardown()
client.calls.assert_done()
delays.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]
assert len(messages) == 1 and MANAGED_FILE_ID in messages[0], messages
@pytest.mark.parametrize(
"last, reported",
@ -255,11 +254,19 @@ class TestBatchCancellation:
(UnknownApiError(status_code=403, body="forbidden"), "after cancellation failed: HTTP 403"),
],
)
def test_anything_but_still_cancelling_at_the_deadline_still_fails(
def test_cancellation_that_never_takes_effect_or_cannot_be_read_back_fails(
self, last: Result[BatchObject], reported: str
) -> None:
client: Final = CleanupClient(
calls=ExpectedCalls((f"retrieve None {MANAGED_BATCH_ID}",) * 2), batches=(batch("cancelling"), last)
calls=ExpectedCalls(
(
f"retrieve None {MANAGED_BATCH_ID}",
f"cancel None {MANAGED_BATCH_ID}",
f"retrieve None {MANAGED_BATCH_ID}",
)
),
batches=(batch("in_progress"), last),
cancellations=(batch("in_progress"),),
)
times: Final = (0.0, BATCH_CANCEL_TIMEOUT_SECONDS)
ticks: Final[Callable[[], float]] = Mock(side_effect=times)