diff --git a/tests/e2e/batches/batch_cleanup.py b/tests/e2e/batches/batch_cleanup.py index f8fc1bbf01f..054fa84a24e 100644 --- a/tests/e2e/batches/batch_cleanup.py +++ b/tests/e2e/batches/batch_cleanup.py @@ -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) diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index c0933a7d8e0..8044d7aa93b 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/tests/e2e/batches/test_batches_e2e.py @@ -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 diff --git a/tests/e2e_harness/batches/test_batch_cleanup.py b/tests/e2e_harness/batches/test_batch_cleanup.py index 218e79f37ac..2155c16ec8c 100644 --- a/tests/e2e_harness/batches/test_batch_cleanup.py +++ b/tests/e2e_harness/batches/test_batch_cleanup.py @@ -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)