diff --git a/examples/fully_async/README.md b/examples/fully_async/README.md index e69d9a80aa..07590e93c7 100644 --- a/examples/fully_async/README.md +++ b/examples/fully_async/README.md @@ -60,6 +60,25 @@ work unchanged under fully-async: See `examples/coding_agent_rl/` for a non-trivial example that plugs in a multi-turn agent (Claude Code in a Docker-Proxy sandbox) this way. +To bound off-policy staleness, set `--max-policy-version-lag N`. A completed +group is accepted when `current_policy_version - group_policy_version <= N`; +older groups are rejected and their admission-time prompt snapshots are +requeued for regeneration. The default is unset (no stale rejection), while +`--max-policy-version-lag 0` accepts only the current policy version. + +## Metrics + +Fully-async rollouts report fixed keys only—policy versions and sample IDs +never become labels—so metric cardinality stays bounded: + +* `fully_async/version_lag/count_{0,1,2_to_3,4_plus,unknown}` +* `fully_async/count/{stale_rejected,stale_requeued,aborted_requeued}` +* `fully_async/queue/completed_groups` +* `fully_async/policy/current_version` + +Counters cover the interval since the previous completed rollout; queue size +and current version are gauges captured when that interval is emitted. + ## Worker Internals (Very Short) * First call: create a process-wide `AsyncRolloutWorker` (thread + asyncio @@ -70,6 +89,8 @@ multi-turn agent (Claude Code in a Docker-Proxy sandbox) this way. * Completed groups land on an output queue; each `generate_rollout` call drains until it has `rollout_batch_size` groups and returns them sorted by `sample.index`. +* When a policy-version lag budget is configured, stale completed groups are + replaced by fresh attempts from their admission-time prompt snapshots. * Groups containing an `ABORTED` sample are pushed back into `data_buffer.add_samples` instead of being shipped to training. * Worker is stopped automatically at process exit via `atexit`. diff --git a/slime/ray/rollout.py b/slime/ray/rollout.py index abd8e35d70..e64a317587 100644 --- a/slime/ray/rollout.py +++ b/slime/ray/rollout.py @@ -1,4 +1,5 @@ import dataclasses +import importlib import itertools import logging import multiprocessing @@ -504,6 +505,8 @@ def __init__(self, args, pg): runtime_env={"env_vars": add_default_ray_env_vars()}, ).remote() self.rollout_id = -1 + self.policy_version = 0 + self.args.policy_version = self.policy_version self._health_monitors = [] if not self.args.debug_train_only and self.args.use_fault_tolerance: @@ -587,6 +590,37 @@ def get_num_rollout_per_epoch(self): assert self.args.rollout_global_dataset return len(self.data_source) // self.args.rollout_batch_size + def _call_generate_rollout_hook(self, hook_name: str, **kwargs): + """Call an optional lifecycle hook defined beside the rollout function.""" + + module = importlib.import_module(self.generate_rollout.__module__) + hook = getattr(module, hook_name, None) + if hook is None: + return None + return hook(**kwargs) + + def before_weight_update(self) -> int: + """Notify the rollout implementation before actor weights are updated.""" + + self._call_generate_rollout_hook( + "before_weight_update", + policy_version=self.policy_version, + ) + return self.policy_version + + def after_weight_update(self, succeeded: bool) -> int: + """Publish a new policy version only after all weight updates succeed.""" + + if succeeded: + self.policy_version += 1 + self.args.policy_version = self.policy_version + self._call_generate_rollout_hook( + "after_weight_update", + policy_version=self.policy_version, + succeeded=succeeded, + ) + return self.policy_version + def generate(self, rollout_id): start_time = time.time() self.rollout_id = rollout_id diff --git a/slime/rollout/fully_async_rollout.py b/slime/rollout/fully_async_rollout.py index 6b94fe7acd..f37191b1eb 100644 --- a/slime/rollout/fully_async_rollout.py +++ b/slime/rollout/fully_async_rollout.py @@ -14,24 +14,28 @@ the number of sglang engines to match the per-sample semaphore cap in :mod:`slime.rollout.sglang_rollout`. -The worker is intentionally oblivious to slime's higher-level pause / -weight-update signalling (e.g. ``GenerateState.aborted``). Each in-flight -generation short-circuits on those signals on its own and surfaces -:data:`Sample.Status.ABORTED`; the only piece the worker owns is -**redirecting ABORTED groups back to ``data_buffer``** instead of shipping -them to training, so the next rollout (with refreshed weights) can pick -them up. +The worker snapshots the current training-side ``policy_version`` when each +group is admitted and preserves that version on completion, even if a weight +update finishes while generation is still running. When +``args.max_policy_version_lag`` is set, queue consumers reject results outside +that budget and requeue their admission-time inputs for regeneration. It +remains oblivious to pause / abort policy: each in-flight generation surfaces +:data:`Sample.Status.ABORTED` on its own, and the worker redirects those groups +back to ``data_buffer`` instead of shipping them to training. """ from __future__ import annotations import asyncio import atexit +import copy import logging import queue import threading import time +from dataclasses import dataclass, field +from slime.rollout.base_types import RolloutFnTrainOutput from slime.rollout.sglang_rollout import GenerateState, generate_and_rm_group from slime.utils.async_utils import run from slime.utils.http_utils import get_rollout_num_engines @@ -39,6 +43,7 @@ __all__ = [ "AsyncRolloutWorker", + "CompletedSampleRecord", "generate_rollout_fully_async", ] @@ -48,6 +53,34 @@ # Global worker, shared across rollout calls so the queue stays warm. _global_worker: AsyncRolloutWorker | None = None _worker_lock = threading.Lock() +_published_policy_version = 0 + + +def _new_metrics_state() -> dict[str, int]: + return { + "lag_zero": 0, + "lag_one": 0, + "lag_two_to_three": 0, + "lag_four_plus": 0, + "lag_unknown": 0, + "stale_rejected": 0, + "stale_requeued": 0, + "aborted_requeued": 0, + } + + +@dataclass(frozen=True) +class CompletedSampleRecord: + """A completed group and the immutable policy version captured at admission.""" + + gid: int + group: list[Sample] + policy_version: int + # Generation mutates Sample objects in place and custom generators may + # replace them entirely. Keep the admission-time input so a stale result + # can be retried instead of accidentally treating a completed response as + # a fresh prompt. + retry_group: list[Sample] | None = field(default=None, repr=False, compare=False) def _get_global_worker(args, data_buffer) -> AsyncRolloutWorker: @@ -56,7 +89,10 @@ def _get_global_worker(args, data_buffer) -> AsyncRolloutWorker: if _global_worker is None or not _global_worker.worker_thread.is_alive(): logger.info("starting fully-async rollout worker") _global_worker = AsyncRolloutWorker( - args, data_buffer, concurrency=args.sglang_server_concurrency * get_rollout_num_engines(args) + args, + data_buffer, + concurrency=args.sglang_server_concurrency * get_rollout_num_engines(args), + policy_version=max(getattr(args, "policy_version", 0), _published_policy_version), ) _global_worker.start() return _global_worker @@ -73,11 +109,35 @@ def _stop_global_worker() -> None: atexit.register(_stop_global_worker) +def before_weight_update(*, policy_version: int) -> None: + """Lifecycle hook called before training publishes a new actor policy.""" + + logger.debug("fully-async: preparing to update policy version %d", policy_version) + + +def after_weight_update(*, policy_version: int, succeeded: bool) -> None: + """Publish ``policy_version`` after a successful actor weight update.""" + + if not succeeded: + logger.warning("fully-async: weight update failed; policy remains unpublished") + return + + global _published_policy_version + with _worker_lock: + if policy_version < _published_policy_version: + raise ValueError( + f"policy version must be monotonic: current={_published_policy_version}, new={policy_version}" + ) + _published_policy_version = policy_version + if _global_worker is not None: + _global_worker.publish_policy_version(policy_version) + + class AsyncRolloutWorker: """Background thread + asyncio loop that continuously consumes groups from ``data_buffer`` and runs :func:`generate_and_rm_group` on each.""" - def __init__(self, args, data_buffer, concurrency: int = 10): + def __init__(self, args, data_buffer, concurrency: int = 10, policy_version: int = 0): self.args = args self.data_buffer = data_buffer self.concurrency = concurrency @@ -87,10 +147,14 @@ def __init__(self, args, data_buffer, concurrency: int = 10): # and freeze every in-flight generation. Backpressure lives in _loop() # instead, which stops topping up while a full pool of completed groups # is already waiting to be consumed. - self.output_queue: queue.Queue[tuple[int, list[Sample]]] = queue.Queue() + self.output_queue: queue.Queue[CompletedSampleRecord] = queue.Queue() self.poll_interval = 1.0 self.worker_thread: threading.Thread | None = None self.state = GenerateState(args) + self._policy_version = policy_version + self._policy_version_lock = threading.Lock() + self._metrics_lock = threading.Lock() + self._metrics_state = _new_metrics_state() # -- public -------------------------------------------------------------- @@ -104,25 +168,121 @@ def stop(self) -> None: if self.worker_thread and self.worker_thread.is_alive(): self.worker_thread.join(timeout=5) - def get_completed_groups(self, limit: int | None = None) -> list[tuple[int, list[Sample]]]: - """Pop up to ``limit`` completed groups (all of them when ``None``). + def get_completed_groups( + self, + limit: int | None = None, + *, + max_policy_version_lag: int | None = None, + ) -> list[CompletedSampleRecord]: + """Pop up to ``limit`` trainable completed groups. Callers that only need a fixed number of groups must pass ``limit`` — anything popped beyond it would otherwise have to be thrown away, and these groups are fully generated and reward-scored, with their prompts - already consumed from ``data_buffer``. + already consumed from ``data_buffer``. When a lag budget is configured, + records beyond it are deterministically rejected and their admission- + time prompt snapshots are requeued for regeneration. """ - completed: list[tuple[int, list[Sample]]] = [] + if max_policy_version_lag is not None and max_policy_version_lag < 0: + raise ValueError(f"max_policy_version_lag must be non-negative, got {max_policy_version_lag}") + + completed: list[CompletedSampleRecord] = [] + stale_count = 0 + current_policy_version = self.policy_version while limit is None or len(completed) < limit: try: - completed.append(self.output_queue.get_nowait()) + record = self.output_queue.get_nowait() except queue.Empty: break + version_lag = current_policy_version - record.policy_version + if version_lag < 0: + raise ValueError( + "completed record is newer than the worker policy: " + f"current={current_policy_version}, record={record.policy_version}" + ) + if max_policy_version_lag is not None and version_lag > max_policy_version_lag: + if record.retry_group is None: + raise RuntimeError(f"stale completed record {record.gid} has no admission snapshot to requeue") + self.data_buffer.add_samples([record.retry_group]) + stale_count += 1 + continue + completed.append(record) + if stale_count: + with self._metrics_lock: + self._metrics_state["stale_rejected"] += stale_count + self._metrics_state["stale_requeued"] += stale_count + logger.info( + "fully-async: rejected and requeued %d stale group(s) at policy version %d (max lag=%d)", + stale_count, + current_policy_version, + max_policy_version_lag, + ) return completed def queue_size(self) -> int: return self.output_queue.qsize() + def _record_version_lag_locked(self, version_lag: int | None) -> None: + if version_lag is None or version_lag < 0: + self._metrics_state["lag_unknown"] += 1 + elif version_lag == 0: + self._metrics_state["lag_zero"] += 1 + elif version_lag == 1: + self._metrics_state["lag_one"] += 1 + elif version_lag <= 3: + self._metrics_state["lag_two_to_three"] += 1 + else: + self._metrics_state["lag_four_plus"] += 1 + + def record_processed_groups(self, groups: list[list[Sample]]) -> None: + """Accumulate a fixed-bucket lag histogram for shipped samples.""" + + current_policy_version = self.policy_version + with self._metrics_lock: + for group in groups: + for sample in group: + sample_version = sample.policy_version + version_lag = current_policy_version - sample_version if isinstance(sample_version, int) else None + self._record_version_lag_locked(version_lag) + + def snapshot_metrics(self, *, reset: bool) -> dict[str, int]: + """Return only bounded-cardinality keys and optionally start a new interval.""" + + with self._metrics_lock: + state = dict(self._metrics_state) + if reset: + self._metrics_state = _new_metrics_state() + + return { + "fully_async/version_lag/count_0": state["lag_zero"], + "fully_async/version_lag/count_1": state["lag_one"], + "fully_async/version_lag/count_2_to_3": state["lag_two_to_three"], + "fully_async/version_lag/count_4_plus": state["lag_four_plus"], + "fully_async/version_lag/count_unknown": state["lag_unknown"], + "fully_async/count/stale_rejected": state["stale_rejected"], + "fully_async/count/stale_requeued": state["stale_requeued"], + "fully_async/count/aborted_requeued": state["aborted_requeued"], + "fully_async/queue/completed_groups": self.queue_size(), + "fully_async/policy/current_version": self.policy_version, + } + + @property + def policy_version(self) -> int: + """Return the latest successfully published policy version.""" + + with self._policy_version_lock: + return self._policy_version + + def publish_policy_version(self, policy_version: int) -> None: + """Publish a monotonic policy version for future admissions.""" + + with self._policy_version_lock: + if policy_version < self._policy_version: + raise ValueError( + f"policy version must be monotonic: current={self._policy_version}, new={policy_version}" + ) + self._policy_version = policy_version + # -- internals ----------------------------------------------------------- def _thread_main(self) -> None: @@ -157,6 +317,14 @@ async def _loop(self) -> None: for group in groups: gid = gid_counter gid_counter += 1 + retry_group = ( + copy.deepcopy(group) + if getattr(self.args, "max_policy_version_lag", None) is not None + else None + ) + policy_version = self.policy_version + for sample in group: + sample.policy_version = policy_version task = asyncio.create_task( generate_and_rm_group( self.args, @@ -165,7 +333,7 @@ async def _loop(self) -> None: evaluation=False, ) ) - task.add_done_callback(self._make_done_cb(gid)) + task.add_done_callback(self._make_done_cb(gid, policy_version, retry_group)) active_tasks.add(task) await asyncio.sleep(self.poll_interval) @@ -183,7 +351,12 @@ async def _loop(self) -> None: except Exception: # noqa: BLE001 pass - def _make_done_cb(self, gid: int): + def _make_done_cb( + self, + gid: int, + policy_version: int, + retry_group: list[Sample] | None = None, + ): def _cb(done_task: asyncio.Task) -> None: try: result = done_task.result() @@ -196,14 +369,28 @@ def _cb(done_task: asyncio.Task) -> None: type(result).__name__, ) return + # Custom generators may replace the admitted Sample objects. Stamp + # their outputs from the admission snapshot, never from mutable + # worker state observed at completion time. + for sample in result: + sample.policy_version = policy_version # Aborted group → requeue, don't ship to training. if any(getattr(s, "status", None) == Sample.Status.ABORTED for s in result): try: self.data_buffer.add_samples([result]) + with self._metrics_lock: + self._metrics_state["aborted_requeued"] += 1 except Exception: # noqa: BLE001 logger.exception("fully-async: failed to requeue aborted group") return - self.output_queue.put((gid, result)) + self.output_queue.put( + CompletedSampleRecord( + gid=gid, + group=result, + policy_version=policy_version, + retry_group=retry_group, + ) + ) return _cb @@ -229,8 +416,11 @@ async def _generate_rollout_async(args, rollout_id: int, data_buffer) -> list[li # Pull only what this rollout still needs; the surplus stays queued for # the next rollout (that is the "queue stays warm" contract). drained = 0 - for gid, group in worker.get_completed_groups(limit=target - len(collected)): - collected[gid] = group + for record in worker.get_completed_groups( + limit=target - len(collected), + max_policy_version_lag=getattr(args, "max_policy_version_lag", None), + ): + collected[record.gid] = record.group drained += 1 if not drained: @@ -263,6 +453,7 @@ def _key(group: list[Sample]) -> int: time.time() - started, worker.queue_size(), ) + worker.record_processed_groups(out) return out @@ -271,4 +462,8 @@ def generate_rollout_fully_async(args, rollout_id, data_buffer, evaluation: bool if evaluation: raise ValueError("fully-async rollout doesn't support evaluation mode") - return run(_generate_rollout_async(args, rollout_id, data_buffer)) + samples = run(_generate_rollout_async(args, rollout_id, data_buffer)) + with _worker_lock: + worker = _global_worker + metrics = worker.snapshot_metrics(reset=True) if worker is not None else None + return RolloutFnTrainOutput(samples=samples, metrics=metrics) diff --git a/slime/utils/arguments.py b/slime/utils/arguments.py index 2ef7dad59b..8fb26a034c 100644 --- a/slime/utils/arguments.py +++ b/slime/utils/arguments.py @@ -464,6 +464,16 @@ def add_rollout_arguments(parser): "This is useful for long responses." ), ) + parser.add_argument( + "--max-policy-version-lag", + type=int, + default=None, + help=( + "Maximum policy-version lag accepted for fully-async rollout samples. " + "Completed groups beyond this lag are rejected and requeued for regeneration. " + "Set to 0 for current-policy-only samples; the default disables stale rejection." + ), + ) parser.add_argument( "--mask-offpolicy-in-partial-rollout", action="store_true", @@ -1980,6 +1990,11 @@ def slime_validate_args(args): if args.over_sampling_batch_size is None: args.over_sampling_batch_size = args.rollout_batch_size + if getattr(args, "max_policy_version_lag", None) is not None: + assert ( + args.max_policy_version_lag >= 0 + ), f"max_policy_version_lag must be non-negative, got {args.max_policy_version_lag}" + assert args.over_sampling_batch_size >= args.rollout_batch_size, ( f"over_sampling_batch_size {args.over_sampling_batch_size} should be greater than or equal to " f"rollout_batch_size {args.rollout_batch_size}" diff --git a/slime/utils/types.py b/slime/utils/types.py index df4b10062c..98e7e5b0ca 100644 --- a/slime/utils/types.py +++ b/slime/utils/types.py @@ -104,6 +104,10 @@ class Sample: # sibling, so loss aggregation averages within the rollout instead of # over-counting it. rollout_id: int | None = None + # Monotonic version of the actor policy used when this sample was admitted + # for generation. This is separate from ``weight_versions``, which contains + # backend-reported weight metadata collected while tokens are generated. + policy_version: int | None = None # prompt prompt: str | list[dict[str, str]] = "" tokens: list[int] = field(default_factory=list) diff --git a/tests/test_fully_async_rollout.py b/tests/test_fully_async_rollout.py index 7c5a595b28..807d75319d 100644 --- a/tests/test_fully_async_rollout.py +++ b/tests/test_fully_async_rollout.py @@ -43,9 +43,9 @@ import pytest import slime.rollout.fully_async_rollout as fa +from slime.rollout.base_types import RolloutFnTrainOutput from slime.utils.types import Sample - NUM_GPUS = 0 @@ -77,17 +77,32 @@ def _make_group(index: int) -> list[Sample]: return [sample] -def _make_worker(monkeypatch, data_buffer=None, concurrency=4) -> fa.AsyncRolloutWorker: +def _make_worker( + monkeypatch, + data_buffer=None, + concurrency=4, + policy_version=0, + max_policy_version_lag=None, +) -> fa.AsyncRolloutWorker: monkeypatch.setattr(fa, "GenerateState", _FakeGenerateState) - args = SimpleNamespace(rollout_global_dataset=True, rollout_batch_size=4) - return fa.AsyncRolloutWorker(args, data_buffer or _FakeDataBuffer([]), concurrency=concurrency) + args = SimpleNamespace( + rollout_global_dataset=True, + rollout_batch_size=4, + max_policy_version_lag=max_policy_version_lag, + ) + return fa.AsyncRolloutWorker( + args, + data_buffer or _FakeDataBuffer([]), + concurrency=concurrency, + policy_version=policy_version, + ) @pytest.mark.unit def test_rollout_takes_target_groups_and_leaves_surplus_queued(monkeypatch): worker = _make_worker(monkeypatch) for gid in range(10): - worker.output_queue.put((gid, _make_group(gid))) + worker.output_queue.put(fa.CompletedSampleRecord(gid, _make_group(gid), policy_version=0)) monkeypatch.setattr(fa, "_get_global_worker", lambda args, data_buffer: worker) args = SimpleNamespace(rollout_global_dataset=True, rollout_batch_size=4) @@ -98,20 +113,148 @@ def test_rollout_takes_target_groups_and_leaves_surplus_queued(monkeypatch): assert [group[0].index for group in out] == [0, 1, 2, 3] # The other six are still queued for the next rollout, not thrown away. assert worker.queue_size() == 6 - assert [gid for gid, _ in worker.get_completed_groups()] == [4, 5, 6, 7, 8, 9] + assert [record.gid for record in worker.get_completed_groups()] == [4, 5, 6, 7, 8, 9] @pytest.mark.unit def test_get_completed_groups_limit(monkeypatch): worker = _make_worker(monkeypatch) for gid in range(5): - worker.output_queue.put((gid, _make_group(gid))) + worker.output_queue.put(fa.CompletedSampleRecord(gid, _make_group(gid), policy_version=0)) - assert [gid for gid, _ in worker.get_completed_groups(limit=2)] == [0, 1] - assert [gid for gid, _ in worker.get_completed_groups()] == [2, 3, 4] + assert [record.gid for record in worker.get_completed_groups(limit=2)] == [0, 1] + assert [record.gid for record in worker.get_completed_groups()] == [2, 3, 4] assert worker.get_completed_groups(limit=3) == [] +@pytest.mark.unit +def test_stale_groups_are_requeued_and_exact_lag_boundary_is_accepted(monkeypatch): + data_buffer = _FakeDataBuffer([]) + worker = _make_worker(monkeypatch, data_buffer=data_buffer, policy_version=5) + + stale_retry = [Sample(index=30, prompt="retry-30")] + worker.output_queue.put( + fa.CompletedSampleRecord( + gid=3, + group=_make_group(3), + policy_version=3, + retry_group=stale_retry, + ) + ) + worker.output_queue.put( + fa.CompletedSampleRecord( + gid=4, + group=_make_group(4), + policy_version=4, + retry_group=[Sample(index=40, prompt="retry-40")], + ) + ) + + records = worker.get_completed_groups(limit=1, max_policy_version_lag=1) + + assert [(record.gid, record.policy_version) for record in records] == [(4, 4)] + assert data_buffer.requeued == [stale_retry] + assert data_buffer.requeued[0][0].status is Sample.Status.PENDING + metrics = worker.snapshot_metrics(reset=False) + assert metrics["fully_async/count/stale_rejected"] == 1 + assert metrics["fully_async/count/stale_requeued"] == 1 + + +@pytest.mark.unit +def test_unset_staleness_budget_preserves_legacy_queue_behavior(monkeypatch): + worker = _make_worker(monkeypatch, policy_version=5) + worker.output_queue.put(fa.CompletedSampleRecord(gid=3, group=_make_group(3), policy_version=1)) + + records = worker.get_completed_groups(limit=1) + + assert [(record.gid, record.policy_version) for record in records] == [(3, 1)] + + +@pytest.mark.unit +def test_rollout_never_returns_stale_group_when_budget_is_enabled(monkeypatch): + data_buffer = _FakeDataBuffer([]) + worker = _make_worker(monkeypatch, data_buffer=data_buffer, policy_version=2) + stale_retry = [Sample(index=10, prompt="retry-10")] + worker.output_queue.put( + fa.CompletedSampleRecord( + gid=1, + group=_make_group(1), + policy_version=1, + retry_group=stale_retry, + ) + ) + worker.output_queue.put( + fa.CompletedSampleRecord( + gid=2, + group=_make_group(2), + policy_version=2, + retry_group=[Sample(index=20, prompt="retry-20")], + ) + ) + monkeypatch.setattr(fa, "_get_global_worker", lambda args, data_buffer: worker) + args = SimpleNamespace( + rollout_global_dataset=True, + rollout_batch_size=1, + max_policy_version_lag=0, + ) + + out = asyncio.run(fa._generate_rollout_async(args, rollout_id=0, data_buffer=data_buffer)) + + assert [group[0].index for group in out] == [2] + assert data_buffer.requeued == [stale_retry] + + +@pytest.mark.unit +def test_stale_retry_uses_pristine_admission_snapshot(monkeypatch): + admitted_group = [Sample(index=6, prompt="p6")] + data_buffer = _FakeDataBuffer([admitted_group]) + + async def _complete_generation(args, group, sampling_params, evaluation): + group[0].response = "generated" + group[0].response_length = 1 + group[0].reward = 1.0 + group[0].status = Sample.Status.COMPLETED + return group + + monkeypatch.setattr(fa, "generate_and_rm_group", _complete_generation) + worker = _make_worker( + monkeypatch, + data_buffer=data_buffer, + concurrency=1, + policy_version=2, + max_policy_version_lag=0, + ) + worker.poll_interval = 0.01 + worker.start() + try: + deadline = time.time() + 3.0 + while worker.queue_size() == 0 and time.time() < deadline: + time.sleep(0.01) + assert worker.queue_size() == 1 + + worker.publish_policy_version(3) + assert worker.get_completed_groups(limit=1, max_policy_version_lag=0) == [] + finally: + worker.stop() + + assert len(data_buffer.requeued) == 1 + retry = data_buffer.requeued[0][0] + assert retry.prompt == "p6" + assert retry.response == "" + assert retry.response_length == 0 + assert retry.reward is None + assert retry.status is Sample.Status.PENDING + assert retry.policy_version is None + + +@pytest.mark.unit +def test_negative_staleness_budget_is_rejected(monkeypatch): + worker = _make_worker(monkeypatch) + + with pytest.raises(ValueError, match="must be non-negative"): + worker.get_completed_groups(max_policy_version_lag=-1) + + @pytest.mark.unit def test_done_callback_never_blocks_event_loop_thread(monkeypatch): """The callback runs on the loop thread; blocking there freezes every @@ -128,7 +271,7 @@ def result(self): def _push_all(): for gid in range(1001): - worker._make_done_cb(gid)(_DoneTask(gid)) + worker._make_done_cb(gid, worker.policy_version)(_DoneTask(gid)) pusher = threading.Thread(target=_push_all, daemon=True) pusher.start() @@ -138,6 +281,68 @@ def _push_all(): assert worker.queue_size() == 1001 +@pytest.mark.unit +def test_metrics_use_fixed_lag_buckets_and_reset_counters(monkeypatch): + worker = _make_worker(monkeypatch, policy_version=5) + groups = [ + [Sample(index=0, policy_version=5)], + [Sample(index=1, policy_version=4)], + [Sample(index=2, policy_version=2)], + [Sample(index=3, policy_version=0)], + [Sample(index=4, policy_version=None)], + ] + + worker.record_processed_groups(groups) + metrics = worker.snapshot_metrics(reset=True) + + assert metrics == { + "fully_async/version_lag/count_0": 1, + "fully_async/version_lag/count_1": 1, + "fully_async/version_lag/count_2_to_3": 1, + "fully_async/version_lag/count_4_plus": 1, + "fully_async/version_lag/count_unknown": 1, + "fully_async/count/stale_rejected": 0, + "fully_async/count/stale_requeued": 0, + "fully_async/count/aborted_requeued": 0, + "fully_async/queue/completed_groups": 0, + "fully_async/policy/current_version": 5, + } + assert worker.snapshot_metrics(reset=False)["fully_async/version_lag/count_0"] == 0 + + +@pytest.mark.unit +def test_aborted_group_increments_requeue_counter(monkeypatch): + worker = _make_worker(monkeypatch) + aborted = _make_group(7) + aborted[0].status = Sample.Status.ABORTED + + class _DoneTask: + def result(self): + return aborted + + worker._make_done_cb(gid=7, policy_version=0)(_DoneTask()) + + assert worker.queue_size() == 0 + assert worker.data_buffer.requeued == [aborted] + assert worker.snapshot_metrics(reset=False)["fully_async/count/aborted_requeued"] == 1 + + +@pytest.mark.unit +def test_generate_entrypoint_returns_samples_with_metrics(monkeypatch): + worker = _make_worker(monkeypatch, policy_version=3) + samples = [_make_group(1)] + samples[0][0].policy_version = 3 + worker.record_processed_groups(samples) + monkeypatch.setattr(fa, "run", lambda awaitable: (awaitable.close(), samples)[1]) + monkeypatch.setattr(fa, "_global_worker", worker) + + output = fa.generate_rollout_fully_async(SimpleNamespace(), 0, None) + + assert isinstance(output, RolloutFnTrainOutput) + assert output.samples == samples + assert output.metrics["fully_async/version_lag/count_0"] == 1 + + @pytest.mark.unit def test_loop_backpressure_stops_topping_up_when_queue_is_full(monkeypatch): """With instantly-completing generations and plenty of fuel, the queue must @@ -169,3 +374,52 @@ async def _instant_generate(args, group, sampling_params, evaluation): # In-flight tasks may still land after the gate check, so allow one pool # beyond the gate — but nothing near the unthrottled fuel size. assert 0 < max_seen <= 2 * concurrency, f"queue grew to {max_seen} with concurrency={concurrency}" + + +@pytest.mark.unit +def test_out_of_order_completion_keeps_each_admission_policy_version(monkeypatch): + worker = _make_worker(monkeypatch, policy_version=3) + + class _DoneTask: + def __init__(self, index): + self.index = index + + def result(self): + return _make_group(self.index) + + old_admission_version = worker.policy_version + worker.publish_policy_version(4) + new_admission_version = worker.policy_version + + # The newer-policy request finishes first. Neither record may consult the + # worker's mutable current version when its callback eventually runs. + worker._make_done_cb(gid=12, policy_version=new_admission_version)(_DoneTask(8)) + worker._make_done_cb(gid=11, policy_version=old_admission_version)(_DoneTask(7)) + + records = worker.get_completed_groups(limit=2) + assert [(record.gid, record.policy_version) for record in records] == [(12, 4), (11, 3)] + assert [record.group[0].policy_version for record in records] == [4, 3] + assert worker.policy_version == 4 + + +@pytest.mark.unit +def test_failed_weight_update_does_not_publish_policy_version(monkeypatch): + worker = _make_worker(monkeypatch, policy_version=5) + monkeypatch.setattr(fa, "_global_worker", worker) + monkeypatch.setattr(fa, "_published_policy_version", 5) + + fa.after_weight_update(policy_version=5, succeeded=False) + assert fa._published_policy_version == 5 + assert worker.policy_version == 5 + + fa.after_weight_update(policy_version=6, succeeded=True) + assert fa._published_policy_version == 6 + assert worker.policy_version == 6 + + +@pytest.mark.unit +def test_policy_version_rejects_regression(monkeypatch): + worker = _make_worker(monkeypatch, policy_version=2) + + with pytest.raises(ValueError, match="policy version must be monotonic"): + worker.publish_policy_version(1) diff --git a/tests/test_megatron_argument_validation.py b/tests/test_megatron_argument_validation.py index ba3304d79c..c44f8f98e0 100644 --- a/tests/test_megatron_argument_validation.py +++ b/tests/test_megatron_argument_validation.py @@ -299,6 +299,15 @@ def test_slime_validate_args_rejects_equal_debug_data_paths(monkeypatch): module.slime_validate_args(args) +@pytest.mark.unit +def test_slime_validate_args_rejects_negative_policy_version_lag(monkeypatch): + module = load_slime_arguments_module(monkeypatch) + args = make_slime_validate_args(max_policy_version_lag=-1) + + with pytest.raises(AssertionError, match="max_policy_version_lag must be non-negative"): + module.slime_validate_args(args) + + @pytest.mark.unit def test_slime_validate_args_preserves_zero_rollout_gpus_under_colocate(monkeypatch): module = load_slime_arguments_module(monkeypatch) diff --git a/tests/test_sample.py b/tests/test_sample.py index 6ff16daaf9..646de2d898 100644 --- a/tests/test_sample.py +++ b/tests/test_sample.py @@ -23,7 +23,6 @@ from slime.utils.types import Sample - NUM_GPUS = 0 @@ -40,6 +39,7 @@ def _make_sample(**overrides) -> Sample: group_index=0, index=42, rollout_id=7, + policy_version=3, prompt="hello", tokens=[1, 2, 3], multimodal_inputs={"images": ["fake_url"]}, @@ -123,6 +123,7 @@ def test_round_trip_preserves_every_field(): "group_index", "index", "rollout_id", + "policy_version", "prompt", "tokens", "multimodal_inputs", diff --git a/tests/test_train_async_policy_version.py b/tests/test_train_async_policy_version.py new file mode 100644 index 0000000000..2837263c05 --- /dev/null +++ b/tests/test_train_async_policy_version.py @@ -0,0 +1,106 @@ +"""CPU tests for async-training policy-version publication.""" + +from __future__ import annotations + +import sys +import types + +import pytest + +# Keep this lifecycle test independent of GPU/server packages imported by the +# real training bootstrap. Restore the real module entries after importing the +# helper so these stubs cannot leak into other test modules. +_stubs = { + "slime.ray.placement_group": { + "create_placement_groups": None, + "create_rollout_manager": None, + "create_training_models": None, + }, + "slime.utils.arguments": {"parse_args": None}, + "slime.utils.logging_utils": { + "configure_logger": None, + "finish_tracking": None, + "init_tracking": None, + }, +} +_original_modules = {name: sys.modules.get(name) for name in _stubs} +try: + for name, attributes in _stubs.items(): + module = types.ModuleType(name) + for attribute, value in attributes.items(): + setattr(module, attribute, value) + sys.modules[name] = module + + import train_async +finally: + for name, original_module in _original_modules.items(): + if original_module is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = original_module + +NUM_GPUS = 0 + + +class _RemoteMethod: + def __init__(self, fn): + self._fn = fn + + def remote(self, *args, **kwargs): + return self._fn(*args, **kwargs) + + +class _FakeRolloutManager: + def __init__(self): + self.policy_version = 0 + self.events = [] + self.before_weight_update = _RemoteMethod(self._before_weight_update) + self.after_weight_update = _RemoteMethod(self._after_weight_update) + + def _before_weight_update(self): + self.events.append(("before", self.policy_version)) + return self.policy_version + + def _after_weight_update(self, *, succeeded): + if succeeded: + self.policy_version += 1 + self.events.append(("after", self.policy_version, succeeded)) + return self.policy_version + + +class _FakeActorModel: + def __init__(self, *, fail=False): + self.fail = fail + self.update_calls = 0 + + def update_weights(self): + self.update_calls += 1 + if self.fail: + raise RuntimeError("weight update failed") + + +@pytest.mark.unit +def test_policy_version_advances_after_successful_weight_update(monkeypatch): + monkeypatch.setattr(train_async.ray, "get", lambda value: value) + actor_model = _FakeActorModel() + rollout_manager = _FakeRolloutManager() + + train_async._update_actor_weights(actor_model, rollout_manager) + + assert actor_model.update_calls == 1 + assert rollout_manager.policy_version == 1 + assert rollout_manager.events == [("before", 0), ("after", 1, True)] + + +@pytest.mark.unit +def test_failed_weight_update_keeps_policy_version(monkeypatch): + monkeypatch.setattr(train_async.ray, "get", lambda value: value) + actor_model = _FakeActorModel(fail=True) + rollout_manager = _FakeRolloutManager() + + with pytest.raises(RuntimeError, match="weight update failed"): + train_async._update_actor_weights(actor_model, rollout_manager) + + assert actor_model.update_calls == 1 + assert rollout_manager.policy_version == 0 + assert rollout_manager.events == [("before", 0), ("after", 0, False)] diff --git a/train_async.py b/train_async.py index 9141b612fe..a2f2f2c93c 100644 --- a/train_async.py +++ b/train_async.py @@ -6,6 +6,18 @@ from slime.utils.misc import should_run_periodic_action +def _update_actor_weights(actor_model, rollout_manager) -> None: + """Publish actor weights and advance the rollout policy version only on success.""" + + ray.get(rollout_manager.before_weight_update.remote()) + succeeded = False + try: + actor_model.update_weights() + succeeded = True + finally: + ray.get(rollout_manager.after_weight_update.remote(succeeded=succeeded)) + + # The framework supports other asynchronous approaches such as fully async (which is shown in examples/full_async). def train(args): assert not args.colocate, "Colocation is not supported for async training." @@ -23,7 +35,7 @@ def train(args): actor_model, critic_model = create_training_models(args, pgs, rollout_manager) # Always push actor weights to rollout once weights are loaded. - actor_model.update_weights() + _update_actor_weights(actor_model, rollout_manager) if args.check_weight_update_equal: ray.get(rollout_manager.check_weights.remote(action="compare")) @@ -67,7 +79,7 @@ def train(args): # sync generate before update weights to prevent update weight in the middle of generation rollout_data_curr_ref = ray.get(x) if (x := rollout_data_next_future) is not None else None rollout_data_next_future = None - actor_model.update_weights() + _update_actor_weights(actor_model, rollout_manager) if should_run_periodic_action(rollout_id, args.eval_interval, num_rollout_per_epoch): ray.get(rollout_manager.eval.remote(rollout_id))