Pass CI

rust/exo_pyo3_bindings/tests/test_python.py

@@ -1,7 +1,12 @@
 import asyncio
 
 import pytest
-from exo_pyo3_bindings import Keypair, NetworkingHandle, NoPeersSubscribedToTopicError
+from exo_pyo3_bindings import (
+    Keypair,
+    NetworkingHandle,
+    NoPeersSubscribedToTopicError,
+    PyFromSwarm,
+)
 
 
 @pytest.mark.asyncio
@@ -9,8 +14,7 @@ async def test_sleep_on_multiple_items() -> None:
     print("PYTHON: starting handle")
     h = NetworkingHandle(Keypair.generate())
 
-    ct = asyncio.create_task(_await_cons(h))
-    mt = asyncio.create_task(_await_msg(h))
+    rt = asyncio.create_task(_await_recv(h))
 
     # sleep for 4 ticks
     for i in range(4):
@@ -22,13 +26,11 @@ async def test_sleep_on_multiple_items() -> None:
             print("caught it", e)
 
 
-async def _await_cons(h: NetworkingHandle):
+async def _await_recv(h: NetworkingHandle):
     while True:
-        c = await h.connection_update_recv()
-        print(f"PYTHON: connection update: {c}")
-
-
-async def _await_msg(h: NetworkingHandle):
-    while True:
-        m = await h.gossipsub_recv()
-        print(f"PYTHON: message: {m}")
+        event = await h.recv()
+        match event:
+            case PyFromSwarm.Connection() as c:
+                print(f"PYTHON: connection update: {c}")
+            case PyFromSwarm.Message() as m:
+                print(f"PYTHON: message: {m}")

src/exo/worker/tests/unittests/test_runner/test_event_ordering.py

@@ -119,6 +119,12 @@ def patch_out_mlx(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setattr(mlx_batch_generator, "warmup_inference", make_nothin(1))
     monkeypatch.setattr(mlx_batch_generator, "_check_for_debug_prompts", nothin)
     monkeypatch.setattr(mlx_batch_generator, "mx_any", make_nothin(False))
+    def fake_all_gather(
+        tasks: list[TextGeneration], group: object
+    ) -> tuple[list[TextGeneration], list[TextGeneration]]:
+        return (tasks, [])
+
+    monkeypatch.setattr(mlx_batch_generator, "mx_all_gather_tasks", fake_all_gather)
     # Mock apply_chat_template since we're using a fake tokenizer (integer 1).
     # Returns a prompt without thinking tag so detect_thinking_prompt_suffix returns None.
     monkeypatch.setattr(