diff --git a/tests/test_agentic_round_streaming.py b/tests/test_agentic_round_streaming.py index e6c3e288..014e048f 100644 --- a/tests/test_agentic_round_streaming.py +++ b/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: diff --git a/tests/test_rl_island_leave_reason.py b/tests/test_rl_island_leave_reason.py index ffbceeb3..2c62b0b3 100644 --- a/tests/test_rl_island_leave_reason.py +++ b/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 diff --git a/tests/test_verl_full_parameter_launch.py b/tests/test_verl_full_parameter_launch.py index 46d80b39..315c64e2 100644 --- a/tests/test_verl_full_parameter_launch.py +++ b/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() diff --git a/tests/test_verl_full_parameter_sync.py b/tests/test_verl_full_parameter_sync.py index c6ee05b6..5e5b17bd 100644 --- a/tests/test_verl_full_parameter_sync.py +++ b/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: diff --git a/tests/test_verl_lazy_capture.py b/tests/test_verl_lazy_capture.py index f3927302..2a4c0a7c 100644 --- a/tests/test_verl_lazy_capture.py +++ b/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() diff --git a/tests/test_verl_megatron_sync.py b/tests/test_verl_megatron_sync.py index 1a2b02d0..8e960f9f 100644 --- a/tests/test_verl_megatron_sync.py +++ b/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() diff --git a/tests/test_verl_miles_parity.py b/tests/test_verl_miles_parity.py index 160ecc0e..db9a93e5 100644 --- a/tests/test_verl_miles_parity.py +++ b/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() diff --git a/tests/test_wandb_tape.py b/tests/test_wandb_tape.py index f1434a1a..a945c934 100644 --- a/tests/test_wandb_tape.py +++ b/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) diff --git a/yeto/rl/engine/bridges.py b/yeto/rl/engine/bridges.py index d5341190..7c673d44 100644 --- a/yeto/rl/engine/bridges.py +++ b/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: diff --git a/yeto/rl/pool_client.py b/yeto/rl/pool_client.py index 725a2a04..9c83240a 100644 --- a/yeto/rl/pool_client.py +++ b/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()