[CI][Bugfix] Fix flaky `test_store_orders_after_compute_write` (#50926)
Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>
Nick Hill committed
Aug 3, 2026 at 17:54 UTC
74295e3bd47b64964fe6a51825079e75c84cc172
1 file changed
+7
tests/v1/simple_kv_offload/test_worker.py
+7
@@ -100,6 +100,13 @@ def _drive_store(
100
101
if int((cpu[:, 0].to(torch.int32) != val).sum().item()):
102
corrupt += 1
103
+
104
+ # Drain the compute stream before returning: in the no-barrier control
105
+ # phase the store never waits on compute, so the host loop runs far ahead
106
+ # and leaves a backlog of sleep+fill kernels in flight. Without this, the
107
+ # leftover control-phase fills race the barrier phase's fill->copy window
108
+ # on the shared gpu tensor and flakily corrupt one iteration.
109
+ compute_stream.synchronize()
110
return corrupt
111
112