Keep policy versions unique when a trainer resumes behind the store's releases

tests/test_rl_sampler_island.py

@@ -256,9 +256,84 @@ def test_release_is_committed_by_latest_and_republish_is_a_no_op(tmp_path):
     assert releases.resolve(first) == tmp_path / RELEASE_DIR / "r1-v1"
     assert releases.release(1, "h1", _write_weights) is None  # a restarted trainer republishes
     releases.release(2, "h2", _write_weights)
-    rewind = releases.release(1, "h1b", _write_weights)  # restored from an older cut
-    assert (rewind.seq, rewind.version) == (3, 1)
-    assert sorted(p.name for p in (tmp_path / RELEASE_DIR).iterdir()) == [LATEST, "r2-v2", "r3-v1"]
+    assert releases.latest().version == 2
+    # A trainer restored from an older cut used to release v1 again with other
+    # weights here (LATEST went back to v1, and v1 named two sets of weights). The
+    # trainer now continues above the store's releases and never releases a version
+    # the store holds (test_a_single_trainer_restored_from_an_older_cut_...).
+    assert sorted(p.name for p in (tmp_path / RELEASE_DIR).iterdir()) == [LATEST, "r1-v1", "r2-v2"]
+
+
+def test_a_single_trainer_restored_from_an_older_cut_continues_above_the_store(tmp_path):
+    """Intended change of semantics: the store holds v1..v3 (released after the last
+    round cut, at v1). The relaunched single trainer continues its account at v3 (as
+    the verl buffer trainer, rl_round_cut_behind_releases), releases nothing for v3,
+    and its first trained round is v4: LATEST never goes back and every version
+    names one set of weights."""
+    from yeto.rl.engine.bridges import LocalOnlySync
+
+    releases = PolicyReleaseStore(CheckpointStore(tmp_path), keep=4)
+    for version in (1, 2, 3):
+        releases.release(version, f"h{version}", _write_weights)
+    hashes = {r.version: r.policy_hash for r in releases.records()}
+    events = []
+    sync = LocalOnlySync(2)
+    sync.round_cuts = SimpleNamespace(resume=lambda d: {"next_rollout_id": 1})  # wired first
+    driver = SimpleNamespace(
+        publisher=SimpleNamespace(publish=lambda s: None), sync=sync, _max_policy_age=lambda: 2,
+        emit=lambda event, **fields: events.append((event, fields)),
+        export_local_resident=lambda *, policy_version: _State(policy_version, "cut"),
+        rollout=SimpleNamespace(_in_executor=lambda fn, *a: fn(None, *a)))
+    args = SimpleNamespace(yeto_rl_sample_buffer=sample_roles.SampleBufferConfig(
+        role="trainer", store=str(tmp_path), release_inline=True))
+    exported = []
+    actor = SimpleNamespace(export_hf=lambda rollout_id, path: exported.append(rollout_id)
+                            or _write_weights(Path(path)))
+    wired = si.wire_policy_releases(driver, miles_args=args, actor_model=actor,
+                                    runner=SimpleNamespace(run=lambda value: value))
+    br._ACTIVE.pop(br.TRAINER, None)
+    assert driver.published_policy_history() == [(1, "h1"), (2, "h2"), (3, "h3")]
+
+    start = sync.start(driver)
+    assert (start.rollout_id, start.state.policy_version) == (3, 3)  # not the cut's 1
+    assert ("rl_round_cut_behind_releases",
+            {"cut_rollout_id": 1, "latest_release_version": 3}) in events
+    wired.publish(start.state)  # the cut's weights: v3 is the store's, not released again
+    skipped = [f for e, f in events if e == "rl_policy_release_skipped"]
+    assert [(f["policy_version"], f["reason"]) for f in skipped] == [(3, "resumed_base_version")]
+    wired.publish(_State(4, "h4"))  # the first trained round
+    assert exported == [4]
+    assert [f["policy_version"] for e, f in events if e == "rl_policy_released"] == [4]
+    records = releases.records()  # the trainer's store keeps the newest two
+    assert [r.version for r in records] == [3, 4]  # no version released twice
+    assert {r.version: r.policy_hash for r in records if r.version <= 3} == {3: hashes[3]}
+    assert releases.latest().version == 4
+
+
+def test_a_single_trainer_relaunched_without_a_cut_continues_above_the_store(tmp_path):
+    """No round cut (a plain relaunch): the account starts at the store's vK, not at 0."""
+    from yeto.rl.engine.bridges import LocalOnlySync
+
+    releases = PolicyReleaseStore(CheckpointStore(tmp_path), keep=4)
+    releases.release(1, "h1", _write_weights)
+    releases.release(2, "h2", _write_weights)
+    sync = LocalOnlySync(1)
+    driver = SimpleNamespace(
+        publisher=SimpleNamespace(publish=lambda s: None), sync=sync, _max_policy_age=lambda: 1,
+        emit=lambda event, **fields: None,
+        export_local_resident=lambda *, policy_version: _State(policy_version, "base"),
+        rollout=SimpleNamespace(_in_executor=lambda fn, *a: fn(None, *a)))
+    args = SimpleNamespace(yeto_rl_sample_buffer=sample_roles.SampleBufferConfig(
+        role="trainer", store=str(tmp_path), release_inline=True))
+    wired = si.wire_policy_releases(driver, miles_args=args, actor_model=SimpleNamespace(
+        export_hf=lambda rollout_id, path: _write_weights(Path(path))),
+        runner=SimpleNamespace(run=lambda value: value))
+    br._ACTIVE.pop(br.TRAINER, None)
+    start = sync.start(driver)
+    assert start.rollout_id == 2 and sync._end() == 3
+    wired.publish(start.state)
+    wired.publish(_State(3, "h3"))
+    assert [(r.version, r.policy_hash) for r in releases.records()] == [(2, "h2"), (3, "h3")]
 
 
 def test_lora_release_records_the_file_digest_and_the_follower_checks_it(tmp_path):

tests/test_rl_verl_buffer_trainer.py

@@ -1076,3 +1076,37 @@ def origin(state):
     assert [r["policy_version"] for r in log] == sorted(r["policy_version"] for r in log)
     assert [r["policy_version"] for r in log if r["released"] is True] == [1, 3]
     assert origin_threads and all(origin_threads)
+
+
+def test_a_round_cut_behind_the_store_continues_above_its_releases(tmp_path):
+    """The store holds v1..v3 (released after the last round cut, at v1). Retraining
+    from the cut and calling its rounds v2, v3 would give those numbers other
+    weights than the store's; the account continues at v3 as a relaunch does, and
+    the first trained round releases v4: every version names one set of weights."""
+    store_dir = tmp_path / "store"
+    store = CheckpointStore(store_dir)
+    releases = PolicyReleaseStore(store, keep=4)
+    base, _weights = _base(tmp_path, dtype=torch.float32)
+    for version in (1, 2, 3):
+        _released(releases, version)
+    hashes = {r.version: r.policy_hash for r in releases.records()}
+    writer = GroupWriter(FileSampleBuffer(store, batch_size=1), "s0")
+    for i, version in enumerate((2, 3)):
+        writer.write([_sample(version, index=0), _sample(version, index=1)], group_key=f"g{i}")
+    driver, island, _publisher = _trainer(tmp_path, store_dir, FakeBufferActor(), rounds=1, required=2,
+                                          base=base, releases=releases, resume_versions=True,
+                                          buffer=FileSampleBuffer(store, batch_size=1, create_index=False))
+    driver.sync.round_cuts = SimpleNamespace(  # the cut the run resumes from: v1
+        resume=lambda d: {"next_rollout_id": 1}, due=lambda *a, **k: False,
+        save=lambda *a, **k: None, prepare=lambda *a, **k: None)
+    state = driver.run()
+
+    assert driver.rounds_completed == 1 and state.policy_version == 4
+    events = _lines(tmp_path / "tape.jsonl")
+    behind = [e for e in events if e["event"] == "rl_round_cut_behind_releases"]
+    assert [(e["cut_rollout_id"], e["latest_release_version"]) for e in behind] == [(1, 3)]
+    release_log = _lines(tmp_path / "out" / RELEASE_LOG.format(learner=0))
+    assert [(r["policy_version"], r["released"]) for r in release_log] == [(3, False), (4, True)]
+    records = releases.records()
+    assert [r.version for r in records] == [1, 2, 3, 4]  # no version released twice
+    assert {r.version: r.policy_hash for r in records if r.version <= 3} == hashes

yeto/rl/adapters/miles/sampler_island.py

@@ -30,6 +30,14 @@
   ``base_policy``). Synced trainers release it (a restarted trainer restores its
   published versions from the store's releases).
 
+  Relaunch of a single trainer (no lane, no syncer; as the verl buffer trainer,
+  ``merged_release.ResumedReleases``): the version account continues above the
+  highest release vK the store holds -- also when the round cut it resumes from
+  is behind vK (``LocalOnlySync.start``, ``rl_round_cut_behind_releases``). Its
+  first publication is vK and releases nothing (``resumed_base_version``); the
+  first trained round releases vK+1. A version number never names two sets of
+  weights and ``LATEST`` never goes back to an older version.
+
 Engine weight load (the order of the publisher, ``publish.MilesPublisher``):
 under the controller's update lock, every engine loads the release, then the
 plain policy version is stamped and read back on every engine; any failure
@@ -80,7 +88,12 @@
     copy_with_digest,
     ReleaseNotLoaded,
 )
-from yeto.rl.engine.merged_release import MergedReleaseGuard, is_publisher, release_version
+from yeto.rl.engine.merged_release import (
+    MergedReleaseGuard,
+    is_publisher,
+    release_version,
+    resumed_releases,
+)
 from yeto.rl.engine.policy_versions import engine_weight_version
 from yeto.rl.engine.release_delta import DeltaMirror, default_mirror
 from yeto.rl.engine.sample_roles import CONVERTIBLE, SAMPLER, TRAINER
@@ -673,6 +686,7 @@ def __init__(self, inner: Any, releases: PolicyReleaseStore,
                  read_position_from_thread: Callable[[], int | None] | None = None,
                  read_shards_from_thread: Callable[[], dict[str, int] | None] | None = None,
                  delta: DeltaChain | None = None,
+                 resumed_base: int | None = None,
                  timer: Callable[[], float] = time.perf_counter) -> None:
         """``delta``: staged releases go out as bf16 deltas along this chain
         (``policy_release`` module doc, delta releases). ``read_shards``: the
@@ -683,7 +697,9 @@ def __init__(self, inner: Any, releases: PolicyReleaseStore,
         policy version), ``guard`` keeps the release versions increasing and
         ``origin(state)`` is the release record's ``origin``.
         ``read_position_from_thread`` / ``read_shards_from_thread``: the same reads,
-        callable from the upload thread (default: ``read_position`` / ``read_shards``)."""
+        callable from the upload thread (default: ``read_position`` / ``read_shards``).
+        ``resumed_base``: a relaunched single trainer's vK (module doc, relaunch):
+        versions up to it are the store's already and are not released again."""
         if (staging is None) != (releaser is None):
             raise ValueError("staging and releaser go together")
         if delta is not None and staging is None:
@@ -704,6 +720,7 @@ def __init__(self, inner: Any, releases: PolicyReleaseStore,
         self.staging = None if staging is None else Path(staging)
         self.releaser = releaser
         self.skip_base = bool(skip_base)
+        self.resumed_base = None if resumed_base is None else int(resumed_base)
         self.timer = timer
 
     def _engine_checksum(self) -> str | None:
@@ -727,6 +744,13 @@ def publish(self, state: Any, **kwargs: Any) -> Any:
                           latest_version=None if latest is None else latest.version,
                           reason="not_newer_than_latest")
                 return result
+        if self.resumed_base is not None and version <= self.resumed_base:
+            # the store holds vK already and nothing newer is trained yet: releasing it
+            # would give vK a second set of weights (module doc, relaunch)
+            self.emit("rl_policy_release_skipped", policy_version=version,
+                      local_version=int(state.policy_version), latest_version=self.resumed_base,
+                      reason="resumed_base_version")
+            return result
         if self.skip_base and version == BASE_POLICY_VERSION and self.releases.latest() is None:
             self.emit("rl_policy_release_skipped", policy_version=version,
                       local_version=int(state.policy_version), latest_version=None,
@@ -907,6 +931,8 @@ def on_gap(previous: int, version: int) -> None:
         if not publisher_island:
             return None
         synced = {"version_of": version_of, "guard": MergedReleaseGuard(on_gap), "origin": origin}
+    else:
+        synced = _resume_single_trainer(driver, release_store(config, environ))
 
     lora = lora_mode(miles_args)
     # module doc: a trainer island stages its full releases unless --rl-release-inline
@@ -946,6 +972,25 @@ def export(version: int, directory: Path) -> None:
     return publisher
 
 
+def _resume_single_trainer(driver: Any, releases: PolicyReleaseStore) -> dict[str, Any]:
+    """A single trainer relaunched over a store that holds releases (module doc,
+    relaunch): its account continues at the highest release vK, and the versions
+    below it are restored with the hashes the store recorded. Returns the
+    ``ReleasingPublisher`` keywords ({} for a store without releases)."""
+    limit = getattr(driver, "_max_policy_age", None)
+    resumed = resumed_releases(releases, int(limit()) if callable(limit) else 0)
+    if resumed is None:
+        return {}
+    resume_at = getattr(getattr(driver, "sync", None), "resume_at", None)
+    if callable(resume_at):
+        resume_at(resumed.start_version, behind_round_cuts=True)
+    driver.published_policy_history = resumed.history
+    driver.emit("rl_buffer_resume_versions", base_version=resumed.base_version,
+                start_version=resumed.start_version, account_resumed=callable(resume_at),
+                known_versions=list(resumed.known_versions()))
+    return {"resumed_base": resumed.base_version}
+
+
 def start_trainer_releaser(miles_args: Any, config: Any, environ: Any = None) -> Any:
     """FP8 rollout (tier 2), trainer side: a thread of this driver converts every BF16 release
     into an FP8 release (``yeto.rl.quant.release.scoped_releaser``) for the sampler islands."""

yeto/rl/engine/bridges.py

@@ -177,24 +177,38 @@ def __init__(self, num_rollout: int) -> None:
         # (:meth:`resume_at`); 0 = this process starts the account at v0.
         self.resumed_at = 0
 
-    def resume_at(self, version: int) -> None:
+    def resume_at(self, version: int, *, behind_round_cuts: bool = False) -> None:
         """Continue the version account at ``version`` K instead
         of at 0, for a buffer trainer that found releases up to vK in its store
         (``merged_release.ResumedReleases``): its first publication is vK, its rounds
         are K, K+1, ..., and ``num_rollout`` stays the rounds *this process* runs,
         counted from K.
 
-        Refused with a round-cut resume, which brings its own start round."""
+        Refused when the round cuts are wired first. Wired before them (a buffer
+        trainer binds its publisher first), a round cut behind ``version`` continues
+        at ``version`` (:meth:`start`). ``behind_round_cuts``: wired after them (a
+        Miles single trainer wires its round cuts before its releases), the same
+        holds -- ``version`` is a floor under the round cut's start round, not a
+        second start round."""
         if int(version) < 0:
             raise ValueError(f"the version account cannot start below 0, got {version!r}")
-        if self.round_cuts is not None:
+        if self.round_cuts is not None and not behind_round_cuts:
             raise ValueError("a resumed version account and a round-cut resume both set the "
                              "start round; this island has round cuts configured")
         self.resumed_at = int(version)
 
     def start(self, driver: IslandDriver) -> SyncStart:
         resumed = self.round_cuts.resume(driver) if self.round_cuts is not None else None
         rollout_id = self.resumed_at if resumed is None else int(resumed["next_rollout_id"])
+        if resumed is not None and self.resumed_at > rollout_id:
+            # The store already holds releases up to vK above this round cut (they were
+            # released after the cut). Retraining from the cut and naming its rounds
+            # v(cut)..vK would give those numbers other weights than the store's; the
+            # account continues at vK instead (as a relaunch without a cut does), so
+            # the first trained round releases vK+1 and every version names one set of weights.
+            driver.emit("rl_round_cut_behind_releases", cut_rollout_id=rollout_id,
+                        latest_release_version=self.resumed_at)
+            rollout_id = self.resumed_at
         return SyncStart(_local_state(driver, rollout_id), rollout_id,
                          finished=rollout_id >= self._end())