fix flaky multi example connector consistency (#48206)
Signed-off-by: aarushjain29 <aarushi.jain2@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Aarushi Jain committed
Jul 15, 2026 at 02:20 UTC
9dd2e72828c13bf088403ed789c4672ffef99564
1 file changed
+22
-4
tests/v1/kv_connector/unit/test_multi_connector.py
+22
-4
@@ -318,8 +318,12 @@ def test_multi_example_connector_consistency():
318
# connector (first nonzero match is chosen), so update_state_after_alloc
319
# will report those external tokens on that one. Other connectors still
320
# receive the request's real blocks but with 0 external tokens.
321
- storage1_scheduler_events = _ignore_event_collection(events["storage1-SCHEDULER"])
322
- storage2_scheduler_events = _ignore_event_collection(events["storage2-SCHEDULER"])
321
+ storage1_scheduler_events = _events_from_request(
322
+ _ignore_event_collection(events["storage1-SCHEDULER"])
323
+ )
324
+ storage2_scheduler_events = _events_from_request(
325
+ _ignore_event_collection(events["storage2-SCHEDULER"])
326
+ )
327
assert storage1_scheduler_events[:4] == [
328
"on_new_request",
329
"get_num_new_matched_tokens 0",
@@ -348,8 +352,12 @@ def test_multi_example_connector_consistency():
352
# return 0 from the first connector, while the second connector has a hit.
353
# Both connectors receive the request's real blocks, but only the chosen
354
# (second) connector reports external tokens.
351
- storage1_scheduler_events = _ignore_event_collection(events["storage1-SCHEDULER"])
352
- storage2_scheduler_events = _ignore_event_collection(events["storage2-SCHEDULER"])
355
+ storage1_scheduler_events = _events_from_request(
356
+ _ignore_event_collection(events["storage1-SCHEDULER"])
357
+ )
358
+ storage2_scheduler_events = _events_from_request(
359
+ _ignore_event_collection(events["storage2-SCHEDULER"])
360
+ )
361
assert storage1_scheduler_events[:4] == [
362
"on_new_request",
363
"get_num_new_matched_tokens 0",
@@ -375,6 +383,16 @@ def _ignore_event_collection(events: list[str]) -> list[str]:
383
return [event for event in events if event not in ignored]
384
385
386
+def _events_from_request(events: list[str]) -> list[str]:
387
+ # The async engine-core can emit a trailing build_connector_meta from the
388
+ # previous generation's final (idle) scheduler step. Depending on timing it
389
+ # may be flushed into this window ahead of on_new_request, so anchor the
390
+ # comparison on the new request's first event to avoid a flaky ordering.
391
+ if "on_new_request" in events:
392
+ return events[events.index("on_new_request") :]
393
+ return events
394
+
395
+
396
def get_connector_events() -> dict[str, list[str]]:
397
# Read in connector events and reset the files.
398
import glob