From 49db27528a966d001df467c7017c8ca28c40fead Mon Sep 17 00:00:00 2001 From: datdo-msft <131494842+datdo-msft@users.noreply.github.com> Date: Tue, 1 Sep 2026 03:14:55 -0700 Subject: [PATCH] fix(test): deflake zmq load-snapshot round-trip tests (#35787) Co-authored-by: Shangming Cai --- .../managers/test_load_snapshot_backends.py | 51 +++++++++++++++---- 1 file changed, 41 insertions(+), 10 deletions(-) diff --git a/test/registered/unit/managers/test_load_snapshot_backends.py b/test/registered/unit/managers/test_load_snapshot_backends.py index 3e18f8e37..bb507bf7e 100644 --- a/test/registered/unit/managers/test_load_snapshot_backends.py +++ b/test/registered/unit/managers/test_load_snapshot_backends.py @@ -60,6 +60,24 @@ def _warmup_zmq(writers, reader, attempts=20, interval=0.05): raise RuntimeError(f"warmup failed: expected {expected}, received {received}") +def _read_until(read_fn, predicate, attempts=100, interval=0.05): + """Poll ``read_fn`` until ``predicate`` holds on its result, or give up. + + zmq delivery is asynchronous: once ``write()`` returns, the background IO + thread may not have handed the message to the reader yet, so a fixed sleep + races with delivery and flakes under load. Polling keeps the reads + deterministic without weakening the assertions -- a backend that never + surfaces the expected snapshot still fails, only after the full timeout. + """ + result = read_fn() + for _ in range(attempts): + if predicate(result): + return result + time.sleep(interval) + result = read_fn() + raise AssertionError(f"predicate never held; last result: {result!r}") + + class TestShmRoundTrip(CustomTestCase): def test_single_rank_write_read(self): path = _temp_path() @@ -140,9 +158,11 @@ class TestZmqRoundTrip(CustomTestCase): _warmup_zmq([writer], reader) writer.write(LoadSnapshot(dp_rank=0, num_running_reqs=7, timestamp=2.0)) - time.sleep(0.05) - load = reader.read(0) + load = _read_until( + lambda: reader.read(0), + lambda snap: snap is not None and snap.timestamp == 2.0, + ) self.assertIsNotNone(load) self.assertEqual(load.num_running_reqs, 7) self.assertEqual(load.timestamp, 2.0) @@ -169,9 +189,12 @@ class TestZmqRoundTrip(CustomTestCase): w.write( LoadSnapshot(dp_rank=rank, num_running_reqs=rank + 1, timestamp=3.0) ) - time.sleep(0.05) - loads = reader.read_all() + loads = _read_until( + lambda: reader.read_all(), + lambda snaps: len(snaps) == dp_size + and all(snap.timestamp == 3.0 for snap in snaps), + ) self.assertEqual(len(loads), dp_size) for load in loads: self.assertEqual(load.num_running_reqs, load.dp_rank + 1) @@ -194,9 +217,11 @@ class TestZmqRoundTrip(CustomTestCase): writer.write( LoadSnapshot(dp_rank=0, num_running_reqs=i, timestamp=float(i)) ) - time.sleep(0.05) - load = reader.read(0) + load = _read_until( + lambda: reader.read(0), + lambda snap: snap is not None and snap.num_running_reqs == 9, + ) self.assertIsNotNone(load) self.assertEqual(load.num_running_reqs, 9) self.assertEqual(load.timestamp, 9.0) @@ -414,9 +439,12 @@ class TestEndToEndZmqSimulation(CustomTestCase): num_total_tokens=100 + rank * 50, ) ) - time.sleep(0.05) - loads = reader.read_all() + loads = _read_until( + lambda: reader.read_all(), + lambda snaps: len(snaps) == dp_size + and all(snap.timestamp == 1.0 for snap in snaps), + ) self.assertEqual(len(loads), dp_size) self.assertEqual(loads[0].num_running_reqs, 10) self.assertEqual(loads[1].num_running_reqs, 11) @@ -433,9 +461,12 @@ class TestEndToEndZmqSimulation(CustomTestCase): num_total_tokens=200 + rank * 50, ) ) - time.sleep(0.05) - loads = reader.read_all() + loads = _read_until( + lambda: reader.read_all(), + lambda snaps: len(snaps) == dp_size + and all(snap.timestamp == 2.0 for snap in snaps), + ) self.assertEqual(loads[0].num_running_reqs, 20) self.assertEqual(loads[1].num_running_reqs, 21) self.assertEqual(loads[0].num_total_tokens, 200)