A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 21:00 UTC
[Test] Fix flaky race condition in test_abort_final_step (#38414)
Signed-off-by: Yifan <yzong@redhat.com>
yzong-rh committed
Mar 28, 2026 at 05:06 UTC
6dad4c57226d522b1af0b12ae6a2ba66e4fc7d5d
1 file changed
+19
-11
tests/v1/engine/test_abort_final_step.py
+19
-11
@@ -259,8 +259,25 @@ async def test_abort_during_final_step(async_scheduling: bool):
259
# Wait for generation to complete
260
await gen_task
261
262
- # Give the scheduler a moment to finish cleanup
263
- await asyncio.sleep(0.1)
262
+ # Poll for the KV connector to record the finish status
263
+ timeout = 5.0
264
+ start = time.time()
265
+ captured_statuses = []
266
+ while time.time() - start < timeout:
267
+ with open(status_file) as f4:
268
+ status_lines = f4.read().strip().split("\n")
269
+ captured_statuses = [
270
+ line
271
+ for line in status_lines
272
+ if line and line.startswith("FINISHED_")
273
+ ]
274
+ if captured_statuses:
275
+ break
276
+ await asyncio.sleep(0.05)
277
+ else:
278
+ raise TimeoutError(
279
+ "Timeout waiting for KV connector to record finish status."
280
+ )
281
282
# Verify we got output
283
assert len(outputs) > 0, "Should have received at least one output"
@@ -275,15 +292,6 @@ async def test_abort_during_final_step(async_scheduling: bool):
292
f"'{final_output.outputs[0].finish_reason}'. "
293
)
294
278
- with open(status_file) as f4:
279
- status_lines = f4.read().strip().split("\n")
280
- # Filter for actual finish statuses (not INIT or empty lines)
281
- captured_statuses = [
282
- line
283
- for line in status_lines
284
- if line and line.startswith("FINISHED_")
285
- ]
286
-
295
assert len(captured_statuses) >= 1, (
296
f"Expected at least 1 captured finish status, got "
297
f"{len(captured_statuses)}. File content: {status_lines}"