diff --git a/tests/test_sync_lead_cap.py b/tests/test_sync_lead_cap.py index cdc7d83e..764add13 100644 --- a/tests/test_sync_lead_cap.py +++ b/tests/test_sync_lead_cap.py @@ -127,6 +127,51 @@ async def test_the_wait_gives_up_at_half_the_wait_timeout_and_never_times_out_it assert len(pairing.inflight) == 1 and pairing.failure is None +@pytest.mark.asyncio +async def test_the_wait_also_covers_rounds_trained_before_the_next_pull_arrives(): + """After a merge and before the syncer's next PULL nothing is in flight, + but a round trained then lands in the next push all the same: with K = 1 + the island waits until it pushed from its current cut and that merged.""" + role = _PairedRole("actor", [_candidate(1, 1)], []) + pairing = _pairing(role, max_local_rounds_ahead=1) + pairing.local_round = 1 + pairing.submit_ready() + await pairing.commit_staged_broadcasts() + assert pairing.inflight == {} and role.fragment_versions == (1, 0) + assert await pairing.wait_for_lead() == 0.0 # pushed from this round's cut + pairing.local_round = 2 # trained one more round; the next PULL is late + assert pairing.lead_blockers() == [] and pairing.lead_unpushed() + served = [] + + async def serve(): + served.append(time.monotonic()) + if len(served) == 3: + role.candidates.append(_candidate(3, 2)) # the next PULL arrives + return 0 + + assert await pairing.wait_for_lead(serve) > 0 + assert len(served) >= 4 and ("actor", 3, True) in role.events + assert pairing.inflight == {} and role.fragment_versions == (3, 0) + assert (pairing.lead_waits, pairing.lead_wait_escapes) == (1, 0) + + +@pytest.mark.asyncio +async def test_the_wait_for_a_late_pull_gives_up_at_half_the_wait_timeout(): + role = _PairedRole("actor", [_candidate(1, 1)], []) + pairing = _pairing(role, max_local_rounds_ahead=1, wait_timeout=0.4) + pairing.local_round = 1 + pairing.submit_ready() + await pairing.commit_staged_broadcasts() + pairing.local_round = 2 + started = time.monotonic() + await pairing.wait_for_lead() + assert 0.1 <= time.monotonic() - started < 0.4 + assert (pairing.lead_waits, pairing.lead_wait_escapes) == (1, 1) + pairing.local_round = 3 + assert await pairing.wait_for_lead() > 0 # past the give-up time: one poll, then trains on + assert pairing.lead_wait_escapes == 2 + + @pytest.mark.asyncio async def test_no_wait_in_the_terminal_phase_or_without_blockers(): role = _PairedRole("actor", [_candidate(1, 1)], [], release_broadcasts=False) @@ -413,6 +458,11 @@ def record(kind, learner, value): with lock: log.append((time.monotonic(), kind, learner, value)) + # The clients dial for 10 s; a loaded machine can take longer to start the syncer. + deadline = time.monotonic() + 60 + while "syncer listening" not in log_file.read_text(): + assert process.poll() is None and time.monotonic() < deadline, "the syncer did not start" + time.sleep(0.05) clients = [harness._client(learner, port, profile_hash) for learner in (0, 1)] errors: list[BaseException] = [] roles: dict[int, _WireRole] = {} @@ -520,14 +570,25 @@ def test_real_syncer_fast_island_stays_within_the_cap_and_both_islands_join_ever syncer_binary, tmp_path, monkeypatch): # noqa: F811 - the imported fixture capped = _run_two_islands(syncer_binary, tmp_path / "capped", monkeypatch, cap=1) _check_merges(capped) - assert _max_lead(capped.log) <= 1 - fast = _pushes(capped.log, 0) - slow = _pushes(capped.log, 1) - # Bounded local progress per push: the fast island trains at most K + 1 - # rounds on a fragment before it pushes it, so its merge weight - # (c_steps * T^2) stays within (K + 1)x of the slow island's. - assert max(c for c, _ in fast.values()) <= 2 - assert max(c for c, _ in slow.values()) <= 2 + # With K = 1 a fast island starts round r + 1 only once every push it + # made from a cut <= r was merged, so the slow island pushed the same + # step (both islands join every merge). + fast_pushes: dict[int, int] = {} + slow_steps: set[int] = set() + for _, kind, who, value in sorted(capped.log, key=lambda item: item[0]): + if kind == "push" and who == 0: + fast_pushes[value[0]] = value[2] + elif kind == "push": + slow_steps.add(value[0]) + elif kind == "start" and who == 0: + unmerged = [step for step, cut in fast_pushes.items() if cut < value and step not in slow_steps] + assert not unmerged, (value, unmerged) + # ... and only after it pushed from its newest cut (round r). + assert value == 1 or value - 1 in fast_pushes.values(), (value, fast_pushes) + # So it stays close to the slow island: one round ahead, or two when + # the slow island itself waited under the cap at a boundary while its + # push from an older cut completed the fast island's merge. + assert _max_lead(capped.log) <= 2 waits = [fields for _, kind, who, (name, fields) in ((t, k, w, v) for t, k, w, v in capped.log if k == "event") if who == 0 and name == "rl_full_parameter_round"] diff --git a/yeto/rl/engine/full_parameter_sync.py b/yeto/rl/engine/full_parameter_sync.py index f542d080..f7d699a0 100644 --- a/yeto/rl/engine/full_parameter_sync.py +++ b/yeto/rl/engine/full_parameter_sync.py @@ -306,6 +306,10 @@ def __init__(self, streams: Mapping[str, RoleStream], profile: FullParameterSync # Lead cap: local rounds this island completed (set by the session at # each boundary) and what its waits did. self.local_round = 0 + # The latest local round whose cut this island retained a pair from, + # and when (lead cap: rounds trained since the island last pushed). + self.retained_round = 0 + self.retained_round_at = time.monotonic() self.lead_waits = 0 self.lead_wait_seconds = 0.0 self.lead_wait_escapes = 0 @@ -522,6 +526,9 @@ def _new_pair(self, candidates: dict[str, Any]) -> PairedSubmission: or candidate.c_steps < self.profile.local_horizon ): raise RuntimeError("role retained cuts are not one H-qualified round") + if self.local_round > self.retained_round: + self.retained_round = self.local_round + self.retained_round_at = time.monotonic() return PairedSubmission(candidates, local_round=self.local_round) def _retain(self, pair: PairedSubmission) -> None: @@ -908,8 +915,21 @@ def lead_blockers(self) -> list[PairedSubmission]: return [] return [pair for pair in self.inflight.values() if pair.local_round <= self.local_round + 1 - limit] + def lead_unpushed(self) -> bool: + """Lead cap K: the island trained K rounds since it last retained a pair. + + Between a merge and the syncer's next PULL the island has no pair in + flight, so :meth:`lead_blockers` is empty, yet every round it trains + then still lands in its next push (larger c_steps) and widens its lead + over an island the next merge waits for. So the island also waits + until it pushes from a cut at most K - 1 rounds old.""" + + limit = self.profile.max_local_rounds_ahead + return limit is not None and self.retained_round <= self.local_round - limit + async def wait_for_lead(self, serve: Any = None) -> float: - """Lead cap: block at this boundary until :meth:`lead_blockers` is empty. + """Lead cap: block at this boundary while :meth:`lead_blockers` is not + empty or :meth:`lead_unpushed` holds. The wait keeps the island's sync duties: every poll materializes and pushes PULLs it has candidates for (from this boundary's cut), applies @@ -929,14 +949,18 @@ async def wait_for_lead(self, serve: Any = None) -> float: Returns the seconds this boundary waited.""" blockers = self.lead_blockers() - if not blockers or self.any_finalizing or self.failure is not None: + unpushed = self.lead_unpushed() + if (not blockers and not unpushed) or self.any_finalizing or self.failure is not None: return 0.0 started = time.monotonic() poll = max(self.profile.poll_seconds, LEAD_WAIT_POLL_S) escaped = False - oldest = min(pair.local_round for pair in blockers) - print(f"[yeto-full-parameter] lead wait: local round {self.local_round}, " - f"{len(blockers)} pair(s) pushed at local round {oldest}.. are unmerged " + if blockers: + oldest = min(pair.local_round for pair in blockers) + reason = f"{len(blockers)} pair(s) pushed at local round {oldest}.. are unmerged" + else: + reason = f"no push since local round {self.retained_round}" + print(f"[yeto-full-parameter] lead wait: local round {self.local_round}, {reason} " f"(lead cap {self.profile.max_local_rounds_ahead})", file=sys.stderr, flush=True) while True: await asyncio.sleep(poll) @@ -946,9 +970,11 @@ async def wait_for_lead(self, serve: Any = None) -> float: if serve is not None: await serve() blockers = self.lead_blockers() - if not blockers or self.any_finalizing: + unpushed = self.lead_unpushed() + if (not blockers and not unpushed) or self.any_finalizing: break - give_up = min(pair.retained_at for pair in blockers) + self.profile.wait_timeout / 2 + since = min(pair.retained_at for pair in blockers) if blockers else self.retained_round_at + give_up = since + self.profile.wait_timeout / 2 if time.monotonic() >= give_up: escaped = True break