Restore test fixes dropped from #324; pump the agentic round in streaming tests

tests/test_agentic_round_streaming.py

@@ -132,7 +132,9 @@ def _entries(tmp_path: Path):
     return entries, {e.sample_id: decode_group(reader.read_payload(e.payload_ref)) for e in entries}
 
 
-def _wait(predicate, timeout=5.0):
+# Every wait is for something that must happen, so the timeout only bounds a
+# failure; it is generous for loaded CI hosts (2 cores running other tests).
+def _wait(predicate, timeout=30.0):
     deadline = time.monotonic() + timeout
     while not predicate():
         assert time.monotonic() < deadline, "timed out"
@@ -145,11 +147,15 @@ def test_a_fast_group_is_written_while_a_slow_group_of_the_same_round_still_runs
     harness.gates["slow"] = threading.Event()
     rnd, engine, events = _round(tmp_path, harness, Prompts({0: ["slow", "fast"]}))
     try:
-        rnd(0)  # returns once the fast group is written and a slot is free
+        rnd(0)  # returns once a slot is free
+        # The first fast trial to end frees a slot, possibly before its sibling
+        # ends; groups settle only inside a round call, so pump empty rounds.
+        _wait(lambda: _entries(tmp_path)[0] or rnd(1))
         entries, _ = _entries(tmp_path)
         assert [e.sample_id for e in entries] == ["s-r0-fast"]
-        (first,) = [f for n, f in events if n == ROUND_EVENT]
-        assert first["groups"] == 1 and first["agentic"]["in_flight"] == N and rnd.holds_weights()
+        rounds = [f for n, f in events if n == ROUND_EVENT]
+        assert sum(f["groups"] for f in rounds) == 1
+        assert rounds[-1]["agentic"]["in_flight"] == N and rnd.holds_weights()
         harness.gates["slow"].set()
         rnd.flush()
     finally:
@@ -179,6 +185,7 @@ def test_the_next_round_is_drawn_while_an_earlier_one_runs(tmp_path):
         rnd(0)  # 2 of 4 slots taken: returns at once for the next round's tasks
         assert not _entries(tmp_path)[0]
         rnd(1)
+        _wait(lambda: _entries(tmp_path)[0] or rnd(2))  # pump until the fast group settles
         assert [e.sample_id for e in _entries(tmp_path)[0]] == ["s-r1-fast"]
         harness.gates["slow"].set()
         rnd.flush()
@@ -338,7 +345,10 @@ async def close(self, trajectory_id):
     try:
         rnd(0)
         rnd(1)
-        _wait(lambda: len(prompts.calls) >= 2)
+        # Groups settle only inside a round call on this thread, so a group that
+        # ends after rnd(1) returned is reported only by a later call: pump
+        # empty rounds until fast and bad have both reported their rows.
+        _wait(lambda: rnd(2) or len(prompts.calls) >= 2)
         harness.gates["slow"].set()
         rnd.flush()
     finally:

tests/test_rl_island_leave_reason.py

@@ -242,3 +242,66 @@ def test_the_verl_pool_island_command_carries_the_lease(monkeypatch):
     run = make_miles_island_task(args, parse_gpu_spec(args.gpu)[0], 0, 1, "127.0.0.1:1").run
     assert "yeto.rl.adapters.verl.island_entry" in run
     assert " --rl-sync-island-lease-s 90.0" in run and " --rl-island-pool" in run
+
+
+# --- a LEAVE that does not reach the syncer is recorded as not sent -----------
+
+class _LeaveLink:
+    """The SyncerClient surface PoolIslandClient.leave uses."""
+
+    def __init__(self, *, shut_down=False, written=True):
+        import threading
+
+        self.shutdown = threading.Event()
+        if shut_down:
+            self.shutdown.set()
+        self.written = written
+        self.sent: list[int] = []
+        self.closed = False
+
+    def send_leave(self, reason, timeout=5.0):
+        self.sent.append(int(reason))
+        return self.written
+
+    def close(self):
+        self.closed = True
+
+
+@pytest.mark.parametrize(("link", "expected", "frames"), [
+    (dict(), None, [LeaveReason.PREEMPTED]),
+    (dict(shut_down=True), "syncer_shut_down", []),
+    (dict(written=False), "not_written", [LeaveReason.PREEMPTED]),
+])
+def test_the_pool_client_says_whether_its_leave_was_sent(link, expected, frames):
+    from yeto.rl.pool_client import PoolClientConfig, PoolIslandClient
+    from yeto.rl.core import CanonicalTensorSpec, build_avg_layout
+
+    layout = build_avg_layout((CanonicalTensorSpec("base_model.model.layer.lora_A.weight", (2, 2), "float32", 4),))
+    fake = _LeaveLink(**link)
+    client = PoolIslandClient(PoolClientConfig(("127.0.0.1", 1), 0), layout,
+                              client_factory=lambda *args, **kwargs: fake)
+    assert client.leave(LeaveReason.PREEMPTED) == expected
+    assert fake.sent == frames and fake.closed
+
+
+def test_a_leave_after_the_run_ended_is_taped_as_not_sent(tmp_path):
+    fake = FakePoolSyncer()
+    engine, sync, client = _pool(tmp_path, fake)
+
+    def train_step(batch):
+        client.leave = lambda reason=0: "syncer_shut_down"
+        assert sync.leave(LeaveReason.PREEMPTED)
+        raise RuntimeError("killed after the notice")
+    engine.trainer.train_step = train_step
+    with pytest.raises(RuntimeError):
+        _driver(engine, sync, tmp_path).run()
+    [leave] = [e for e in _tape(tmp_path) if e["event"] == "elastic_leave"]
+    assert (leave["reason_name"], leave["sent"], leave["not_sent"]) == ("PREEMPTED", False, "syncer_shut_down")
+
+
+def test_a_sent_leave_is_taped_as_sent(tmp_path):
+    fake = FakePoolSyncer()
+    engine, sync, _ = _pool(tmp_path, fake)
+    _driver(engine, sync, tmp_path).run()
+    [leave] = [e for e in _tape(tmp_path) if e["event"] == "elastic_leave"]
+    assert leave["sent"] is True and "not_sent" not in leave

tests/test_verl_full_parameter_launch.py

@@ -239,7 +239,7 @@ def test_two_launcher_built_verl_islands_run_async_rounds_through_the_real_synce
     process = group_popen(["bash", "-c", command], stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
     results: dict = {}
     try:
-        threads = [threading.Thread(target=_run_island, args=(i, plans[i], 20 * rounds, results))
+        threads = [threading.Thread(target=_run_island, args=(i, plans[i], None, results))
                    for i in range(2)]
         for thread in threads:
             thread.start()

tests/test_verl_full_parameter_sync.py

@@ -10,9 +10,11 @@
 
 import argparse
 import hashlib
+import itertools
 import shutil
 import subprocess
 import threading
+import time
 from dataclasses import replace
 from pathlib import Path
 from types import SimpleNamespace
@@ -394,7 +396,26 @@ def _step(learner_id: int, round_id: int) -> dict[str, torch.Tensor]:
     return {name: 0.01 * torch.randn(shape, generator=generator) for name, shape in SHAPES.items()}
 
 
-def _run_island(learner_id: int, plan: dict, max_rounds: int, results: dict) -> None:
+# An island with asynchronous rounds trains until the syncer's final cut; how
+# many local rounds that takes depends on how fast the syncer merges, which a
+# loaded host slows far more than the tiny local rounds. Bound it by time.
+ASYNC_ISLAND_SECONDS = 100.0
+
+
+def local_rounds(max_rounds: int | None):
+    """Round ids ``0..max_rounds-1``, or (``None``) until the final cut: then
+    a TimeoutError after ``ASYNC_ISLAND_SECONDS``, inside the test's join."""
+    if max_rounds is not None:
+        yield from range(max_rounds)
+        return
+    deadline = time.monotonic() + ASYNC_ISLAND_SECONDS
+    for round_id in itertools.count():
+        if time.monotonic() > deadline:
+            raise TimeoutError(f"no final cut after {round_id} local rounds")
+        yield round_id
+
+
+def _run_island(learner_id: int, plan: dict, max_rounds: int | None, results: dict) -> None:
     """Train and sync until the syncer's final cut (a sweep: exactly one round
     per sweep; asynchronous rounds: as many local rounds as the pool needs)."""
     model = _Model(seed=0)  # every island loads the same checkpoint
@@ -404,7 +425,7 @@ def _run_island(learner_id: int, plan: dict, max_rounds: int, results: dict) ->
     driver = _Driver()
     try:
         state = sync.start(driver).state
-        for round_id in range(max_rounds):
+        for round_id in local_rounds(max_rounds):
             with torch.no_grad():
                 for name, parameter in model.named_parameters():
                     parameter.add_(_step(learner_id, round_id)[name])
@@ -442,7 +463,7 @@ def _two_islands(binary, tmp_path, monkeypatch, *, rounds: int, fragments: int,
         threads = [threading.Thread(target=_run_island, args=(
             i, _plan(i, port, profile, rounds=rounds, fragments=fragments, async_rounds=async_rounds,
                      matrix_merge=matrix_merge, pipeline=pipeline, final_cut=final_cut),
-            rounds if not async_rounds else 20 * rounds, results)) for i in range(2)]
+            rounds if not async_rounds else None, results)) for i in range(2)]
         for thread in threads:
             thread.start()
         for thread in threads:

tests/test_verl_lazy_capture.py

@@ -238,7 +238,7 @@ def checking_delta(self, stored, local, local_fragment_id, **kwargs):
     try:
         threads = [threading.Thread(target=_run_island, args=(
             i, _plan(i, port, profile, rounds=rounds, fragments=fragments, async_rounds=True,
-                     matrix_merge="rda", pipeline=2, final_cut="download"), 20 * rounds, results))
+                     matrix_merge="rda", pipeline=2, final_cut="download"), None, results))
             for i in range(2)]
         for thread in threads:
             thread.start()

tests/test_verl_megatron_sync.py

@@ -32,7 +32,7 @@
 from megatron_bridge_roundtrip import _DistributedLeaf
 from test_rl_algorithm_provenance import _fake_sky, _launcher_args, _no_modal_listing  # noqa: F401
 from test_verl_full_parameter_sync import (  # noqa: F401 - syncer_binary is a fixture
-    COMPONENT, TOKENS, _Driver, _start_syncer, needs_cargo, syncer_binary,
+    COMPONENT, TOKENS, _Driver, _start_syncer, local_rounds, needs_cargo, syncer_binary,
 )
 from yeto.rl import megatron_hf_boundary as boundary
 from yeto.rl import megatron_masters
@@ -603,14 +603,14 @@ def _plan(learner_id: int, port: int, profile: str, *, rounds: int, fragments: i
                 "local_init": async_rounds, "matrix_merge": matrix_merge, "island_pool": None}}
 
 
-def _run_island(learner_id: int, plan: dict, max_rounds: int, results: dict, harness_cls=None) -> None:
+def _run_island(learner_id: int, plan: dict, max_rounds: int | None, results: dict, harness_cls=None) -> None:
     harness = (harness_cls or _MegatronHarness)()
     island = harness.island
     sync, _policy_state, _stream = vfs.build_full_parameter_sync(island, plan, backend_identity_sha256="c" * 64)
     driver = _Driver()
     try:
         state = sync.start(driver).state
-        for round_id in range(max_rounds):
+        for round_id in local_rounds(max_rounds):
             harness.train(learner_id, round_id)
             ids = (f"l{learner_id}-r{round_id}",)
             batch_hash = hashlib.sha256(ids[0].encode()).hexdigest()
@@ -649,7 +649,7 @@ def _two_islands(binary, tmp_path, monkeypatch, *, rounds, fragments, async_roun
         for plan in plans:
             plan[vfs.PLAN_KEY].update(plan_extra or {})
         threads = [threading.Thread(target=_run_island, args=(i, plans[i], rounds if not async_rounds
-                                                               else 20 * rounds, results, harnesses[i]))
+                                                               else None, results, harnesses[i]))
                    for i in range(2)]
         for thread in threads:
             thread.start()

tests/test_verl_miles_parity.py

@@ -92,7 +92,7 @@ def _run_islands(binary, tmp_path, monkeypatch, *, rounds, fragments, async_roun
                        matrix_merge=matrix_merge, pipeline=pipeline, final_cut=final_cut) for i in range(2)]
         for plan in plans:
             plan[vfs.PLAN_KEY].update(plan_extra)
-        threads = [threading.Thread(target=_run_island, args=(i, plans[i], 20 * rounds, results))
+        threads = [threading.Thread(target=_run_island, args=(i, plans[i], None, results))
                    for i in range(2)]
         for thread in threads:
             thread.start()

tests/test_wandb_tape.py

@@ -327,7 +327,9 @@ def run():
     return t
 
 
-def _wait_for(predicate, timeout=5.0):
+# Waits only for things that must happen: the timeout bounds a failure and is
+# generous for loaded CI hosts.
+def _wait_for(predicate, timeout=30.0):
     deadline = time.monotonic() + timeout
     while not predicate() and time.monotonic() < deadline:
         time.sleep(0.01)

yeto/rl/engine/bridges.py

@@ -785,14 +785,18 @@ def leave(self, reason: int) -> bool:
         code = sendable(reason)
         driver = getattr(self, "_driver", None)
         try:
-            self.client.leave(code)
+            not_sent = self.client.leave(code)
         except Exception as error:  # noqa: BLE001 - the syncer may already be gone
             if driver is not None:
                 driver.emit("elastic_leave_failed", reason=int(code), reason_name=code.name,
                             error_type=type(error).__name__, error=str(error))
             return False
         if driver is not None:
-            driver.emit("elastic_leave", reason=int(code), reason_name=code.name)
+            # sent=False: no LEAVE reached the syncer (e.g. the run had already
+            # ended), so its tape has no pool_leave for this island.
+            extra = {} if not_sent is None else {"not_sent": not_sent}
+            driver.emit("elastic_leave", reason=int(code), reason_name=code.name,
+                        sent=not_sent is None, **extra)
         return True
 
     def close(self, reason: int = LeaveReason.FINISHED) -> None:

yeto/rl/pool_client.py

@@ -34,6 +34,10 @@
 from ..tensor_io import pack_tensor, unpack_fragment
 
 
+# PoolClient.leave: why a LEAVE frame was not sent.
+LEAVE_NOT_SENT_SHUT_DOWN = "syncer_shut_down"
+LEAVE_NOT_SENT_NOT_WRITTEN = "not_written"
+
 class LeaveReason(IntEnum):
     """Why an island left the pool (one byte of MSG_LEAVE; the syncer's
     ``pool_leave`` record names it, see ``protocol.rs`` LEAVE_REASON_NAMES).
@@ -231,12 +235,22 @@ def join(self) -> JoinAck:
         version = 0 if self.latest_base is None else self.latest_base.outer_version
         return JoinAck(base_version=version, catch_up=version > 0)
 
-    def leave(self, reason: int = LeaveReason.FINISHED) -> None:
-        """Leave the pool: LEAVE with the reason, then close the connection."""
+    def leave(self, reason: int = LeaveReason.FINISHED) -> str | None:
+        """Leave the pool: LEAVE with the reason, then close the connection.
+
+        Returns None when the LEAVE frame was written, else why it was not:
+        ``LEAVE_NOT_SENT_SHUT_DOWN`` (the run already ended, e.g. the final cut
+        arrived first; the syncer records no leave) or
+        ``LEAVE_NOT_SENT_NOT_WRITTEN`` (the frame was not written in time)."""
         self._stop.set()
-        if not self.client.shutdown.is_set():
-            self.client.send_leave(int(reason))
-        self.close()
+        try:
+            if self.client.shutdown.is_set():
+                return LEAVE_NOT_SENT_SHUT_DOWN
+            if not self.client.send_leave(int(reason)):
+                return LEAVE_NOT_SENT_NOT_WRITTEN
+            return None
+        finally:
+            self.close()
 
     def close(self) -> None:
         self._stop.set()