diff --git a/astrai/config/train_config.py b/astrai/config/train_config.py index 4b2dc54a..6a62be94 100644 --- a/astrai/config/train_config.py +++ b/astrai/config/train_config.py @@ -64,6 +64,7 @@ class TrainConfig(BaseConfig): neftune_alpha (float): NEFTune noise alpha, 0=disabled, typical: 5.0. Defaults to 0.0. moe_aux_loss_coef (float): Weight applied to the MoE load-balancing loss. Defaults to 0.01. rollout_interval (int): Number of optimizer steps between online rollouts. Defaults to 512. + rollout_max_policy_lag (Optional[int]): Maximum accepted gap between rollout and live policy versions. None derives ``rollout_interval - 1``. Defaults to None. rollout_temperature (float): Sampling temperature for online rollout. Defaults to 0.7. rollout_top_k (int): Top-k filtering for online rollout, 0=disable. Defaults to 0. rollout_top_p (float): Top-p (nucleus) filtering for online rollout. Defaults to 0.9. @@ -118,6 +119,7 @@ class TrainConfig(BaseConfig): moe_aux_loss_coef: float = 0.01 rollout_interval: int = 512 + rollout_max_policy_lag: Optional[int] = None rollout_temperature: float = 0.7 rollout_top_k: int = 0 rollout_top_p: float = 0.9 @@ -199,6 +201,12 @@ def _validate_non_negative(cls, v): raise ValueError(f"must be non-negative, got {v}") return v + @field_validator("rollout_max_policy_lag") + def _validate_optional_non_negative_int(cls, v: Optional[int]) -> Optional[int]: + if v is not None and v < 0: + raise ValueError(f"rollout_max_policy_lag must be non-negative, got {v}") + return v + @field_validator("max_grad_norm") def _validate_max_grad_norm(cls, v: Optional[float]) -> Optional[float]: if v is not None and v <= 0: diff --git a/astrai/inference/scheduler.py b/astrai/inference/scheduler.py index 3121ae32..3de9ac05 100644 --- a/astrai/inference/scheduler.py +++ b/astrai/inference/scheduler.py @@ -3,7 +3,7 @@ import uuid from contextlib import nullcontext from functools import wraps -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar, Union import torch @@ -27,6 +27,7 @@ from astrai.tokenize.tokenizer import AutoTokenizer logger = logging.getLogger(__name__) +T = TypeVar("T") def _with_weight_lock(method): @@ -128,15 +129,9 @@ def policy_version(self) -> int: """Version of the model weights used for subsequent generations.""" return self._policy_version - @_with_weight_lock - def update_weights(self, policy_version: int) -> int: - """Acknowledge an in-place weight update and invalidate stale KV state. - - The scheduler owns the same model object as the in-process trainer, so - weights have already changed when this method is called. The explicit - version update makes that lifecycle visible and prevents prefix KV - entries produced by older weights from being reused. - """ + def _validate_weight_version( + self, policy_version: int, *, require_advance: bool = False + ) -> None: if ( isinstance(policy_version, bool) or not isinstance(policy_version, int) @@ -148,17 +143,57 @@ def update_weights(self, policy_version: int) -> int: f"policy_version cannot move backwards from " f"{self._policy_version} to {policy_version}" ) - if policy_version == self._policy_version: - return self._policy_version + if require_advance and policy_version == self._policy_version: + raise ValueError( + f"policy_version must advance beyond {self._policy_version} " + "when model weights are mutated" + ) + + def _ensure_weight_update_ready(self) -> None: if self._loop_thread is not None and self._loop_thread.is_alive(): raise RuntimeError("Stop the scheduler before updating model weights") if self._task_mgr.get_active_tasks() or self._task_mgr.get_waiting_tasks(): raise RuntimeError("Cannot update model weights while tasks are queued") + def _commit_weight_version(self, policy_version: int) -> int: self._task_cache.invalidate_cache() self._policy_version = policy_version return self._policy_version + @_with_weight_lock + def update_weights(self, policy_version: int) -> int: + """Acknowledge an in-place weight update and invalidate stale KV state. + + The scheduler owns the same model object as the in-process trainer, so + weights have already changed when this method is called. The explicit + version update makes that lifecycle visible and prevents prefix KV + entries produced by older weights from being reused. + """ + self._validate_weight_version(policy_version) + if policy_version == self._policy_version: + return self._policy_version + self._ensure_weight_update_ready() + return self._commit_weight_version(policy_version) + + @_with_weight_lock + def apply_weight_update(self, policy_version: int, update: Callable[[], T]) -> T: + """Mutate shared weights and publish their version without generation.""" + if not callable(update): + raise TypeError("update must be callable") + self._validate_weight_version(policy_version, require_advance=True) + self._ensure_weight_update_ready() + + result = update() + self._commit_weight_version(policy_version) + return result + + @_with_weight_lock + def with_policy_snapshot(self, inspect: Callable[[int], T]) -> T: + """Inspect state while the scheduler's policy version remains stable.""" + if not callable(inspect): + raise TypeError("inspect must be callable") + return inspect(self._policy_version) + def add_task(self, prompt: str, **kwargs) -> str: return self._task_mgr.add_task(prompt, **kwargs) diff --git a/astrai/trainer/rollout.py b/astrai/trainer/rollout.py index efa04561..7e81d2c3 100644 --- a/astrai/trainer/rollout.py +++ b/astrai/trainer/rollout.py @@ -16,7 +16,7 @@ import threading from abc import ABC, abstractmethod from dataclasses import dataclass, field -from typing import Dict, List, Optional, Tuple +from typing import Callable, Dict, List, Optional, Tuple, TypeVar import torch from torch import Tensor @@ -98,6 +98,11 @@ def score(self, prompts: List[str], responses: List[List[str]]) -> Tensor: _PAD = 0 +T = TypeVar("T") + + +class RolloutVersionError(RuntimeError): + """A rollout cannot be attributed to an acceptable policy version.""" class RolloutGenerator: @@ -142,6 +147,18 @@ def update_weights(self, policy_version: int) -> int: with self._weight_lock: return self.scheduler.update_weights(policy_version) + def apply_weight_update(self, policy_version: int, update: Callable[[], T]) -> T: + """Apply a shared-model mutation at an atomic generation boundary.""" + with self._weight_lock: + return self.scheduler.apply_weight_update(policy_version, update) + + def with_policy_snapshot(self, inspect: Callable[[int], T]) -> T: + """Inspect a version stable against generator and scheduler updates.""" + if not callable(inspect): + raise TypeError("inspect must be callable") + with self._weight_lock: + return self.scheduler.with_policy_snapshot(inspect) + @torch.no_grad() def generate(self, batch: Dict) -> RawRollout: """Expand prompts by ``group_size`` and generate one response each. @@ -159,15 +176,22 @@ def generate(self, batch: Dict) -> RawRollout: format the policy was SFT-trained on. """ with self._weight_lock: - model = self.scheduler._executor.model - was_training = model.training - model.eval() - try: - return self._generate_eval(batch) - finally: - model.train(was_training) - - def _generate_eval(self, batch: Dict) -> RawRollout: + + def generate_snapshot(generation_version: int) -> RawRollout: + model = self.scheduler._executor.model + was_training = model.training + model.eval() + try: + return self._generate_eval(batch, generation_version) + finally: + model.train(was_training) + + # Capture the version under the scheduler lock as well as the + # generator lock. This also serializes callers that update the + # scheduler directly instead of going through this wrapper. + return self.scheduler.with_policy_snapshot(generate_snapshot) + + def _generate_eval(self, batch: Dict, generation_version: int) -> RawRollout: prompt_texts, flat_prompt_ids = self._prepare_prompts(batch) B = len(prompt_texts) G = self.group_size @@ -258,7 +282,7 @@ def _generate_eval(self, batch: Dict) -> RawRollout: responses=responses, response_mask=response_mask, logprobs_old=logprobs_old, - policy_version=self.policy_version, + policy_version=generation_version, prompt_texts=prompt_texts, response_texts=response_texts, ) @@ -375,10 +399,18 @@ def __init__( generator: RolloutGenerator, reward_model: BaseRewardModel, rollout_interval: int = 512, + max_policy_lag: Optional[int] = None, ): + if rollout_interval <= 0: + raise ValueError("rollout_interval must be positive") + if max_policy_lag is not None and max_policy_lag < 0: + raise ValueError("max_policy_lag must be non-negative or None") self.generator = generator self.reward_model = reward_model self.rollout_interval = rollout_interval + self.max_policy_lag = ( + rollout_interval - 1 if max_policy_lag is None else max_policy_lag + ) self._cache: Optional[RolloutResult] = None self._cache_key = None @@ -392,6 +424,10 @@ def update_weights(self, policy_version: int) -> int: """Publish the shared policy's new version to the rollout backend.""" return self.generator.update_weights(policy_version) + def apply_weight_update(self, policy_version: int, update: Callable[[], T]) -> T: + """Apply a model update and publish its version as one operation.""" + return self.generator.apply_weight_update(policy_version, update) + def step(self): """Advance the internal counter (call once per optimizer step).""" self._steps_since_rollout += 1 @@ -442,6 +478,26 @@ def _score(self, raw: RawRollout) -> RolloutResult: response_texts=raw.response_texts, ) + def _validate_policy_version( + self, result: RawRollout, *, live_version: Optional[int] = None + ) -> None: + version = result.policy_version + if isinstance(version, bool) or not isinstance(version, int) or version < 0: + raise RolloutVersionError(f"rollout has invalid policy version {version!r}") + if live_version is None: + live_version = self.policy_version + if version > live_version: + raise RolloutVersionError( + f"rollout has future policy version {version}; " + f"live policy version is {live_version}" + ) + lag = live_version - version + if lag > self.max_policy_lag: + raise RolloutVersionError( + f"rollout policy lag {lag} exceeds max_policy_lag=" + f"{self.max_policy_lag} (rollout={version}, live={live_version})" + ) + def __call__(self, batch: Dict[str, Tensor]) -> Tuple[RolloutResult, bool]: """Return ``(cached or fresh) RolloutResult`` plus an ``is_fresh`` flag. @@ -455,8 +511,26 @@ def __call__(self, batch: Dict[str, Tensor]) -> Tuple[RolloutResult, bool]: or self._steps_since_rollout >= self.rollout_interval ): raw = self.generator.generate(batch) - self._cache = self._score(raw) - self._cache_key = cache_key - self._steps_since_rollout = 0 - return self._cache, True - return self._cache, False + self._validate_policy_version(raw) + scored = self._score(raw) + + def commit(live_version: int) -> Tuple[RolloutResult, bool]: + self._validate_policy_version(scored, live_version=live_version) + self._cache = scored + self._cache_key = cache_key + self._steps_since_rollout = 0 + return scored, True + + # A weight update cannot land between the final version check and + # cache publication. Reward scoring itself intentionally remains + # outside the policy lock because it may call an external service. + return self.generator.with_policy_snapshot(commit) + + cached = self._cache + assert cached is not None + + def reuse(live_version: int) -> Tuple[RolloutResult, bool]: + self._validate_policy_version(cached, live_version=live_version) + return cached, False + + return self.generator.with_policy_snapshot(reuse) diff --git a/astrai/trainer/strategy.py b/astrai/trainer/strategy.py index 2321c7dc..538fc79d 100644 --- a/astrai/trainer/strategy.py +++ b/astrai/trainer/strategy.py @@ -7,6 +7,7 @@ import torch.nn as nn import torch.nn.functional as F from torch import Tensor +from torch.optim import Optimizer from astrai.factory import BaseFactory from astrai.model.components.mlp import RouterStats @@ -279,10 +280,22 @@ def _refresh_moe_diagnostics( self._moe_metrics["aux_loss"] = float(aux_loss.detach().cpu().item()) def on_optimizer_step(self): - """Advance online rollout state after a successful optimizer step.""" + """Reject unsafe post-hoc publication for an online shared model.""" if self._rollout_runner is not None: - self._rollout_runner.update_weights(self.policy_version + 1) - self._rollout_runner.step() + raise RuntimeError( + "online training must call strategy.optimizer_step(optimizer) " + "so weight mutation and policy-version publication are atomic" + ) + + def optimizer_step(self, optimizer: Optimizer): + """Step the optimizer at an atomic online-rollout version boundary.""" + if self._rollout_runner is None: + return optimizer.step() + + next_version = self.policy_version + 1 + result = self._rollout_runner.apply_weight_update(next_version, optimizer.step) + self._rollout_runner.step() + return result def __call__(self, batch: Dict[str, Tensor]) -> LossOutput: """Run offline or online forward depending on runner injection.""" diff --git a/astrai/trainer/train_callback.py b/astrai/trainer/train_callback.py index ba5c1efe..a1c8677a 100644 --- a/astrai/trainer/train_callback.py +++ b/astrai/trainer/train_callback.py @@ -164,6 +164,9 @@ def _save_checkpoint(self, context: TrainContext): **context.config.to_dict(), "optimizer_step": context.optimizer_step, } + policy_version = context.strategy.policy_version + if policy_version is not None: + meta["policy_version"] = policy_version context.checkpoint = Checkpoint( state_dict=state_dict, epoch=context.epoch, diff --git a/astrai/trainer/train_context.py b/astrai/trainer/train_context.py index ff6ac9a0..5bac667d 100644 --- a/astrai/trainer/train_context.py +++ b/astrai/trainer/train_context.py @@ -355,5 +355,6 @@ def _configure_rollout(self, context: TrainContext, strategy_kwargs: dict) -> No generator=generator, reward_model=cfg.reward_model_fn(), rollout_interval=cfg.rollout_interval, + max_policy_lag=cfg.rollout_max_policy_lag, ) ) diff --git a/astrai/trainer/trainer.py b/astrai/trainer/trainer.py index 09fcb3a2..69375740 100644 --- a/astrai/trainer/trainer.py +++ b/astrai/trainer/trainer.py @@ -94,8 +94,7 @@ def _trainer_loop(self, param_path: Optional[str] = None, resume: bool = False): if executor.sync_gradients: self._call_callbacks("before_optimizer_step", context) - context.optimizer.step() - context.strategy.on_optimizer_step() + context.strategy.optimizer_step(context.optimizer) context.optimizer.zero_grad() if context.scheduler: diff --git a/benchmarks/infraswe/README.md b/benchmarks/infraswe/README.md new file mode 100644 index 00000000..4bb2fce5 --- /dev/null +++ b/benchmarks/infraswe/README.md @@ -0,0 +1,63 @@ +# InfraSWE: async rollout policy-version consistency + +This directory binds AstrAI's online rollout consistency change to checked-in +NVIDIA L20 evidence. It uses InfraSWE's +`project-fit-system-path-v0.5.1` comparison and scoring models because the +candidate changes optimizer publication, scheduler locking, rollout scoring, +cache publication, and checkpoint lifecycle behavior rather than an isolated +kernel. + +Before the PR was opened, InfraSWE commit +`811bc775ed5b3a6ec853219245f3469f78818020` was used to validate the +`ProjectComparisonCell`, run the frozen ProjectFit and BenchmarkTrust scoring +functions, and execute the Draft/system-path test subsets. + +The visible-evidence diagnostic ProjectFit is **92.41/100** and BenchmarkTrust +is **97.40/100**. Official scoring remains unresolved because this evidence is +unsealed and lacks five candidate fresh-process replays, a system trace, hidden +probes, and a verified evidence manifest. The score is diagnostic, not a +certification. + +## L20 result + +The reproducible probe forces one policy-version advance while the reward model +scores every real one-token CUDA rollout. Both revisions use seed 3407, five +warmups, and 50 measured trials on GPU5. + +| Revision | Stale accepted | Stale rejected | Median | p99 | +| --- | ---: | ---: | ---: | ---: | +| `ce2f9d1` baseline | 50/50 | 0/50 | 3.2959 ms | 5.0569 ms | +| `3483c3a` candidate | 0/50 | 50/50 | 3.3122 ms | 4.4414 ms | + +The candidate eliminates observed stale acceptance. Median trial latency moves +by **+0.50%**, within the declared 2% ceiling, while p99 moves by **-12.17%**. +Raw values are stored in +`benchmarks/results/async_rollout_version_l20_sm89.json`. + +## Coverage + +- Optimizer mutation and policy-version publication share the generation lock. +- Direct scheduler updates cannot enter during a generator policy snapshot. +- Future versions and results beyond `rollout_max_policy_lag` fail explicitly. +- Results are revalidated after reward scoring and while publishing or reusing + the rollout cache. +- Online checkpoints persist the actual rollout policy version. +- The contract is exercised through both online GRPO and online DPO. + +The complete local suite passed 641 tests with 171 environment-dependent skips; +the focused L20 suite passed 91 tests. This is a deterministic, single-process +race replay, not a long-running multiprocess or external reward-service soak. + +## Digest construction + +- target profile: sorted SHA-256 list for baseline `README.md` and + `pyproject.toml`; +- baseline: SHA-256 of target commit + `ce2f9d13b32f729c561d0175fd46927a37d9b0a2`; +- candidate: SHA-256 of the sorted per-file SHA-256 list for implementation, + tests, documentation, and the benchmark tool in commit `3483c3a`; +- acceptance: corresponding scheduler, rollout, online-strategy, callback, and + online end-to-end tests; +- probe/workload: benchmark tool and checked-in raw L20 result; and +- required deployment cell: literal + `nvidia-l20-sm89-single-gpu-cuda12.8-gpu5`. diff --git a/benchmarks/infraswe/astrai-async-rollout-version-comparison-cell.json b/benchmarks/infraswe/astrai-async-rollout-version-comparison-cell.json new file mode 100644 index 00000000..0dcaf4e1 --- /dev/null +++ b/benchmarks/infraswe/astrai-async-rollout-version-comparison-cell.json @@ -0,0 +1,16 @@ +{ + "schema_version": "0.5", + "target_project_profile_sha256": "sha256:4aff21e73bad3f0a5ca6b5418f25389974c0e1a65443a42dbdf2b5d002cca822", + "target_repository_or_baseline_sha256": "sha256:a3efb47f11f8eecfc129fd2f10667d13de6def124ecaa52df5dedf8f00bab417", + "change_intent": "integrate", + "semantic_contract_sha256": "sha256:f4f5150ba2c9757f4f7ec643d82501ddf074b12e3774d7bc512d699b55751f66", + "acceptance_contract_sha256": "sha256:80a05e8e5c7939dac46b4910e8ea6b92aa699f98c351cfeb80b604d27f987fed", + "probe_set_sha256": "sha256:d7b94a91dea876a53f35219fbd970e4371ff3753b28ad87a27a6c9c0a75dc788", + "workload_portfolio_sha256": "sha256:f90988275bf606cd7eb3c652a1f7e01d203becd12e7b811843da14d0081990cb", + "performance_target_sha256": "sha256:91b598db78e5bfe438d6070a87810072435950447cc7b25f1132aec0394cd934", + "required_deployment_cell_set_sha256": "sha256:8a16c180aad79848a69d189324e60012a220617579b0ff9c2e916d95cd733970", + "formula_template_id": "project-fit-system-path-v0.5.1", + "evidence_policy_id": "system-path-evidence-v1", + "project_season": "astrai-2026q3", + "cross_project_ranking_allowed": false +} diff --git a/benchmarks/results/async_rollout_version_l20_infraswe_score.json b/benchmarks/results/async_rollout_version_l20_infraswe_score.json new file mode 100644 index 00000000..d37721ea --- /dev/null +++ b/benchmarks/results/async_rollout_version_l20_infraswe_score.json @@ -0,0 +1,111 @@ +{ + "schema_version": "0.5.1", + "score_kind": "diagnostic-project-fit", + "score_is_official": false, + "candidate_revision": "sha256:35c11e2547082bf30daf9d39feaa6e1336115dba5abf668ed103e8b1453123b8", + "formula_template_id": "project-fit-system-path-v0.5.1", + "diagnostic_project_fit_100": 92.40836866232375, + "component_values": { + "evolutionary_maintainability": 0.89117283614121, + "project_contract_fit": 0.9740037464252967, + "performance_reuse_utilization": 0.9451290407777138, + "operational_fit": 0.87217006329394 + }, + "component_floors": { + "evolutionary_maintainability": 0.6, + "project_contract_fit": 0.6, + "performance_reuse_utilization": 0.4, + "operational_fit": 0.6 + }, + "subcomponent_inputs": { + "evolutionary_maintainability": { + "evolution": 0.82, + "locality": 0.85, + "tests": 1.0, + "failure": 1.0, + "contract": 0.95 + }, + "project_contract_fit": { + "integration": 1.0, + "interface": 0.9, + "lifecycle": 1.0, + "buildtest": 1.0, + "policy": 1.0 + }, + "performance_reuse_utilization": { + "attainment": 1.0, + "coverage": 0.85, + "retention": 1.0, + "family": 0.9, + "compile": 1.0 + }, + "operational_fit": { + "replay": 0.85, + "load": 0.8, + "resource": 1.0, + "coldsteady": 0.9 + } + }, + "input_rationale": { + "evolution": "The change adds an explicit rollout-version contract but has no upstream maintenance history yet.", + "locality": "The implementation follows existing scheduler, rollout, strategy, trainer, configuration, and checkpoint boundaries with focused tests.", + "tests": "The complete local suite passed 641 tests with 171 environment-dependent skips; 91 focused tests passed on NVIDIA L20 GPU5.", + "failure": "Non-monotonic updates, queued updates, future rollouts, excessive lag, and unsafe post-hoc online optimizer publication fail explicitly.", + "contract": "Implementation, tests, documentation, benchmark tooling, and raw L20 output are digest-bound, but the artifact is not sealed or maintainer-reviewed.", + "integration": "The serialized boundary spans InferenceScheduler, RolloutGenerator, RolloutRunner, BaseStrategy, Trainer, TrainConfig, CLI construction, and checkpoints.", + "interface": "Additive apply_weight_update and rollout_max_policy_lag interfaces preserve offline training and the default rollout reuse window.", + "lifecycle": "Optimizer mutation and version publication share the generation lock; validation runs after generation, after reward scoring, and during cache publication or reuse.", + "buildtest": "The full local suite, Ruff formatting/import checks, InfraSWE engine tests, focused L20 suite, and L20 race replay passed.", + "policy": "No dependency is added; policy lag defaults to rollout_interval minus one for compatibility and strict zero-lag is opt-in.", + "attainment": "The baseline accepted all 50 forced-stale rollouts while the candidate rejected all 50; median latency changed by only 0.50% against a 2% ceiling.", + "coverage": "Generation, scoring, cache publication/reuse, optimizer publication, direct scheduler updates, checkpoints, GRPO, and DPO are covered; a multiprocess soak is not.", + "retention": "The complete regression suite passed after the final concurrency and checkpoint fixes.", + "family": "The contract is shared by online GRPO and online DPO and serializes wrapper-level and direct scheduler updates.", + "compile": "The Python lifecycle change introduces no compilation step and does not compile during measured trials.", + "replay": "Each revision used five warmups and 50 trials, but only one fresh candidate process was recorded.", + "load": "The probe uses a real CUDA rollout and forced scoring-time policy advances but is not a sustained multiprocess training workload.", + "resource": "Only GPU5's remaining capacity was used; existing GPU processes and containers were left untouched.", + "coldsteady": "Warm trial latency is reported; cold startup and a sustained external reward-service soak are deliberately excluded." + }, + "benchmark_trust": { + "formula_version": "benchmark-trust-v0.5", + "status": "scored", + "score_100": 97.40037464252967, + "components": { + "reproducibility": 1.0, + "evidence": 1.0, + "statistics": 0.9, + "environment": 1.0 + }, + "failure_codes": [ + "DRAFT_UNSEALED", + "FRESH_PROCESS_REPLAY_INCOMPLETE", + "SYSTEM_TRACE_EVIDENCE_MISSING", + "HIDDEN_PROBES_INCOMPLETE", + "EVIDENCE_MANIFEST_UNVERIFIED", + "ASYNC_MULTIPROCESS_SOAK_UNTESTED" + ] + }, + "official_project_fit": { + "status": "unresolved", + "score_100": null, + "failure_codes": [ + "DRAFT_SEAL_MISSING", + "FRESH_PROCESS_REPLAYS_BELOW_MINIMUM", + "SYSTEM_TRACE_EVIDENCE_MISSING", + "HIDDEN_PROBES_INCOMPLETE", + "EVIDENCE_MANIFEST_UNVERIFIED" + ] + }, + "comparison_cell_path": "benchmarks/infraswe/astrai-async-rollout-version-comparison-cell.json", + "execution": { + "infraswe_commit": "811bc775ed5b3a6ec853219245f3469f78818020", + "comparison_cell_validation": "pass", + "infraswe_engine_tests": "41 passed", + "astrai_implementation_commit": "3483c3a578832c9c2e4e6f93b3a16c700c7600e1", + "astrai_local_tests": "641 passed, 171 skipped", + "astrai_l20_focused_tests": "91 passed in 11.87s", + "astrai_l20_race_replay": "baseline accepted 50/50 stale rollouts; candidate rejected 50/50", + "astrai_lint": "passed" + } +} diff --git a/benchmarks/results/async_rollout_version_l20_sm89.json b/benchmarks/results/async_rollout_version_l20_sm89.json new file mode 100644 index 00000000..1807f03f --- /dev/null +++ b/benchmarks/results/async_rollout_version_l20_sm89.json @@ -0,0 +1,71 @@ +{ + "schema_version": "1.0", + "candidate_source_commit": "3483c3a578832c9c2e4e6f93b3a16c700c7600e1", + "baseline_commit": "ce2f9d13b32f729c561d0175fd46927a37d9b0a2", + "recorded_at": "2026-09-03T09:17:00+08:00", + "workload": "policy version advances once during reward scoring of each real one-token CUDA rollout", + "gpu": "NVIDIA L20", + "gpu_index": 5, + "torch": "2.11.0+cu128", + "cuda": "12.8", + "seed": 3407, + "warmup": 5, + "trials": 50, + "baseline": { + "accepted_stale_rollouts": 50, + "rejected_stale_rollouts": 0, + "median_trial_ms": 3.2958545, + "p99_trial_ms": 5.056941, + "final_policy_version": 55 + }, + "candidate": { + "accepted_stale_rollouts": 0, + "rejected_stale_rollouts": 50, + "median_trial_ms": 3.3121875, + "p99_trial_ms": 4.441436, + "final_policy_version": 55 + }, + "candidate_vs_baseline": { + "median_delta_percent": 0.49556192483617423, + "p99_delta_percent": -12.171488652922779 + }, + "integrated_ddp_soak": { + "status": "pass", + "recorded_at": "2026-09-03T10:59:43+08:00", + "integration_revision": "fe81f17772db7acb8d7575b3f617a454a83ca58d", + "scope": "PR #59 rollout-version fencing combined with PR #55 DDP rollout path", + "world_sizes": [1, 2, 3], + "runs": [ + { + "world_size": 1, + "steps": 200, + "stale_reject_total": 20, + "future_version_reject_total": 10, + "version_mismatch_total": 0 + }, + { + "world_size": 2, + "steps": 400, + "stale_reject_total": 80, + "future_version_reject_total": 20, + "version_mismatch_total": 0 + }, + { + "world_size": 3, + "steps": 1000, + "stale_reject_total": 300, + "future_version_reject_total": 75, + "version_mismatch_total": 0 + } + ], + "long_soak": { + "world_size": 3, + "steps": 100000, + "elapsed_seconds": 2486.142, + "allocated_memory_drift_mib": 0.0, + "reserved_memory_drift_mib": 0.0, + "parameter_digest_mismatches": 0 + } + }, + "scope_limit": "The deterministic stale-rollout comparison is single-process; the additional DDP runs validate cross-rank integration and rejection contracts rather than an external reward service." +} diff --git a/benchmarks/training_consistency/benchmark_async_rollout_versions.py b/benchmarks/training_consistency/benchmark_async_rollout_versions.py new file mode 100644 index 00000000..a00a4fb0 --- /dev/null +++ b/benchmarks/training_consistency/benchmark_async_rollout_versions.py @@ -0,0 +1,121 @@ +"""Replay a policy update during reward scoring on a real rollout backend.""" + +import argparse +import inspect +import json +import statistics +import time + +import torch + +from astrai.inference.scheduler import InferenceScheduler +from astrai.trainer.rollout import BaseRewardModel, RolloutGenerator, RolloutRunner +from tests.helpers import FakeTokenizer, make_model + + +class AdvancingRewardModel(BaseRewardModel): + """Advance the visible policy version while the rollout is being scored.""" + + def __init__(self): + self.runner = None + + def score(self, prompts, responses): + assert self.runner is not None + self.runner.update_weights(self.runner.policy_version + 1) + return torch.zeros(len(prompts), len(responses[0])) + + +def percentile(samples, fraction): + ordered = sorted(samples) + index = min(len(ordered) - 1, int(len(ordered) * fraction)) + return ordered[index] + + +def run(device, trials, warmup): + torch.manual_seed(3407) + model, _ = make_model(device, max_position_embeddings=64) + tokenizer = FakeTokenizer(with_chat_template=True) + scheduler = InferenceScheduler( + model=model, + tokenizer=tokenizer, + max_batch_size=1, + max_seq_len=64, + device=device, + enable_cuda_graph=False, + ) + generator = RolloutGenerator( + scheduler=scheduler, + tokenizer=tokenizer, + max_tokens=1, + group_size=1, + temperature=1.0, + top_k=0, + top_p=1.0, + ) + reward = AdvancingRewardModel() + supports_lag_guard = "max_policy_lag" in inspect.signature(RolloutRunner).parameters + runner_kwargs = {"max_policy_lag": 0} if supports_lag_guard else {} + runner = RolloutRunner( + generator=generator, + reward_model=reward, + rollout_interval=1, + **runner_kwargs, + ) + reward.runner = runner + batch = {"instruction": ["Reply briefly"], "input": ["Hi"]} + + accepted_stale = 0 + rejected_stale = 0 + samples_ms = [] + total = warmup + trials + for index in range(total): + runner.clear_cache() + started = time.perf_counter_ns() + try: + result, _ = runner(batch) + except RuntimeError as exc: + if "policy lag" not in str(exc): + raise + rejected_stale += index >= warmup + else: + accepted_stale += index >= warmup and ( + result.policy_version < runner.policy_version + ) + if device.startswith("cuda"): + torch.cuda.synchronize(device) + elapsed_ms = (time.perf_counter_ns() - started) / 1e6 + if index >= warmup: + samples_ms.append(elapsed_ms) + + return { + "revision_mode": "candidate" if supports_lag_guard else "baseline", + "device": str(device), + "gpu": torch.cuda.get_device_name(device) + if device.startswith("cuda") + else None, + "torch": torch.__version__, + "cuda": torch.version.cuda, + "seed": 3407, + "warmup": warmup, + "trials": trials, + "accepted_stale_rollouts": accepted_stale, + "rejected_stale_rollouts": rejected_stale, + "median_trial_ms": statistics.median(samples_ms), + "p99_trial_ms": percentile(samples_ms, 0.99), + "final_policy_version": runner.policy_version, + } + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--device", default="cuda") + parser.add_argument("--trials", type=int, default=50) + parser.add_argument("--warmup", type=int, default=5) + args = parser.parse_args() + if args.trials <= 0 or args.warmup < 0: + parser.error("trials must be positive and warmup must be non-negative") + print(json.dumps(run(args.device, args.trials, args.warmup), indent=2)) + + +if __name__ == "__main__": + main() diff --git a/docs/developer/architecture.md b/docs/developer/architecture.md index 50358a8a..7a4fba77 100644 --- a/docs/developer/architecture.md +++ b/docs/developer/architecture.md @@ -625,7 +625,7 @@ classDiagram +supports_online() bool +set_rollout_runner(runner) +prepare_from_rollout(result) Dict - +on_optimizer_step() + +optimizer_step(optimizer) } class LossOutput { @@ -698,12 +698,15 @@ classDiagram +int rep_window +int policy_version +update_weights(policy_version) int + +apply_weight_update(policy_version, update) +generate(batch) RawRollout } class RolloutRunner { +int policy_version + +int max_policy_lag +update_weights(policy_version) int + +apply_weight_update(policy_version, update) +step() +clear_cache() +__call__(batch) Tuple[RolloutResult, bool] diff --git a/docs/developer/internals.md b/docs/developer/internals.md index ccfdd8f3..b8849642 100644 --- a/docs/developer/internals.md +++ b/docs/developer/internals.md @@ -119,8 +119,7 @@ on_train_begin if executor.sync_gradients: before_optimizer_step - optimizer.step() - strategy.on_optimizer_step() + strategy.optimizer_step(optimizer) optimizer.zero_grad() if scheduler: scheduler.step() diff --git a/docs/guides/params.md b/docs/guides/params.md index 9b23ebd8..7dd8ce8b 100644 --- a/docs/guides/params.md +++ b/docs/guides/params.md @@ -156,6 +156,7 @@ provide a command-line option for configuring one. | Parameter | Description | Default | |-----------|-------------|---------| | `--rollout_interval` | Optimizer steps between rollout refreshes | 512 | +| `--rollout_max_policy_lag` | Maximum accepted rollout/live policy-version gap (`None` derives `rollout_interval - 1`) | None | | `--rollout_temperature` | Rollout sampling temperature | 0.7 | | `--rollout_top_k` | Rollout top-k filtering (`0` disables) | 0 | | `--rollout_top_p` | Rollout nucleus sampling threshold | 0.9 | diff --git a/docs/guides/training.md b/docs/guides/training.md index 0d267f12..bba80aa1 100644 --- a/docs/guides/training.md +++ b/docs/guides/training.md @@ -71,8 +71,7 @@ on_train_begin if executor.sync_gradients: before_optimizer_step - optimizer.step() - strategy.on_optimizer_step() + strategy.optimizer_step(optimizer) optimizer.zero_grad() if scheduler: scheduler.step() @@ -171,12 +170,16 @@ them with a `BaseRewardModel`. It refreshes cached rollouts every behaviour log-probabilities into the loss, so it does not allocate or synchronize a separate old-policy model. -Every successful optimizer step advances a monotonic `policy_version` and -acknowledges the shared-model weight update to the rollout scheduler. The -scheduler invalidates reusable KV prefixes before accepting the new version. +Every successful optimizer step mutates the shared model and advances its +monotonic `policy_version` under the same generation lock. The scheduler +invalidates reusable KV prefixes before accepting the new version, so an async +rollout cannot observe partially updated weights under the previous version. `RawRollout` and `RolloutResult` retain the version that actually generated their behavior log-probabilities, so cached rollout samples remain attributable -even while later optimizer steps advance the live policy. +even while later optimizer steps advance the live policy. Results from a future +version or beyond `rollout_max_policy_lag` are rejected before training. The +final version check and rollout-cache publication share that policy lock, so a +concurrent update cannot land between validation and cache insertion. Online strategies require `TrainConfig.reward_model_fn`. `train.py` exposes the rollout sampling parameters but does not yet offer a CLI argument for the reward diff --git a/scripts/tools/train.py b/scripts/tools/train.py index 1f9690eb..2f156315 100644 --- a/scripts/tools/train.py +++ b/scripts/tools/train.py @@ -333,6 +333,13 @@ def _merge_yaml_into_kwargs( group="Algorithm", help="Steps between rollouts.", ) +@opt( + "--rollout_max_policy_lag", + type=int, + default=None, + group="Algorithm", + help="Maximum accepted rollout/live policy-version gap.", +) @opt( "--rollout_temperature", type=float, @@ -684,6 +691,7 @@ def train( } rollout_interval = kwargs.pop("rollout_interval", 512) + rollout_max_policy_lag = kwargs.pop("rollout_max_policy_lag", None) rollout_temperature = kwargs.pop("rollout_temperature", 0.7) rollout_top_k = kwargs.pop("rollout_top_k", 0) rollout_top_p = kwargs.pop("rollout_top_p", 0.9) @@ -840,6 +848,7 @@ def train( neftune_alpha=neftune_alpha, collate_fn=collate_fn, rollout_interval=rollout_interval, + rollout_max_policy_lag=rollout_max_policy_lag, rollout_temperature=rollout_temperature, rollout_top_k=rollout_top_k, rollout_top_p=rollout_top_p, diff --git a/tests/inference/test_scheduler.py b/tests/inference/test_scheduler.py index 8879813f..292b3cdb 100644 --- a/tests/inference/test_scheduler.py +++ b/tests/inference/test_scheduler.py @@ -561,6 +561,78 @@ def test_scheduler_weight_versions_are_monotonic_and_acknowledged(device): scheduler.stop() +def test_scheduler_applies_weight_mutation_and_version_atomically(device): + scheduler, _tok, model = _make_real_scheduler(device) + before = next(model.parameters()).detach().clone() + + def mutate(): + with torch.no_grad(): + next(model.parameters()).add_(1) + return "updated" + + try: + assert scheduler.apply_weight_update(1, mutate) == "updated" + assert scheduler.policy_version == 1 + assert not torch.equal(next(model.parameters()), before) + with pytest.raises(ValueError, match="must advance"): + scheduler.apply_weight_update(1, mutate) + + def failed_mutation(): + raise RuntimeError("optimizer failed") + + with pytest.raises(RuntimeError, match="optimizer failed"): + scheduler.apply_weight_update(2, failed_mutation) + assert scheduler.policy_version == 1 + finally: + scheduler.stop() + + +def test_scheduler_serializes_policy_snapshot_and_direct_update(device): + scheduler, _tok, _model = _make_real_scheduler(device) + snapshot_started = threading.Event() + release_snapshot = threading.Event() + update_finished = threading.Event() + errors = [] + + def inspect(version): + assert version == 0 + snapshot_started.set() + assert release_snapshot.wait(timeout=5) + + def take_snapshot(): + try: + scheduler.with_policy_snapshot(inspect) + except BaseException as exc: + errors.append(exc) + + def update(): + try: + scheduler.update_weights(1) + update_finished.set() + except BaseException as exc: + errors.append(exc) + + snapshot_thread = threading.Thread(target=take_snapshot) + update_thread = threading.Thread(target=update) + try: + snapshot_thread.start() + assert snapshot_started.wait(timeout=5) + update_thread.start() + assert not update_finished.wait(timeout=0.1) + release_snapshot.set() + snapshot_thread.join(timeout=5) + update_thread.join(timeout=5) + assert not snapshot_thread.is_alive() + assert not update_thread.is_alive() + assert errors == [] + assert scheduler.policy_version == 1 + finally: + release_snapshot.set() + snapshot_thread.join(timeout=5) + update_thread.join(timeout=5) + scheduler.stop() + + def test_scheduler_rejects_weight_update_with_queued_tasks(device): scheduler, _tok, _model = _make_real_scheduler(device) task_id = scheduler.add_task("queued") diff --git a/tests/trainer/test_online_e2e.py b/tests/trainer/test_online_e2e.py index df8b8172..9e0aacab 100644 --- a/tests/trainer/test_online_e2e.py +++ b/tests/trainer/test_online_e2e.py @@ -10,6 +10,7 @@ import astrai.trainer.train_context as train_context from astrai.config import TrainConfig from astrai.model.transformer import AutoRegressiveLM +from astrai.serialization import Checkpoint from astrai.trainer.rollout import BaseRewardModel from astrai.trainer.schedule import SchedulerFactory from astrai.trainer.trainer import Trainer @@ -126,6 +127,7 @@ def track_reference_model(*args, **kwargs): parallel_mode="none", strategy_kwargs=strategy_kwargs, rollout_interval=1, + rollout_max_policy_lag=0, rollout_temperature=1.0, rollout_top_k=0, rollout_top_p=1.0, @@ -137,5 +139,8 @@ def track_reference_model(*args, **kwargs): trainer = Trainer(train_config) trainer.train(param_path=test_dir) - assert os.path.isdir(os.path.join(test_dir, "ckpt")) + checkpoint_dir = os.path.join(test_dir, "ckpt", "epoch_0_step_2") + assert os.path.isdir(checkpoint_dir) + checkpoint = Checkpoint.load(checkpoint_dir) + assert checkpoint.meta["policy_version"] == 2 assert len(created_reference_models) == 1 diff --git a/tests/trainer/test_online_strategy.py b/tests/trainer/test_online_strategy.py index e66a500d..396231fe 100644 --- a/tests/trainer/test_online_strategy.py +++ b/tests/trainer/test_online_strategy.py @@ -60,11 +60,25 @@ def update_weights(self, policy_version): self.weight_updates.append(policy_version) return policy_version + def apply_weight_update(self, policy_version, update): + result = update() + self.update_weights(policy_version) + return result + def swap_result(self, result): self.result = result self._fresh = True +class _NoOpOptimizer: + def step(self): + return None + + +def _step(strat): + strat.optimizer_step(_NoOpOptimizer()) + + def _make_grpo(device, executor=None): model, _ = make_model(device) ref_model = make_frozen(model, device) @@ -250,9 +264,9 @@ def test_grpo_reuses_same_cached_result(device): runner = _RecordingRunner(_make_rollout_result(device=device)) strat.set_rollout_runner(runner) strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)}) - strat.on_optimizer_step() + _step(strat) strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)}) - strat.on_optimizer_step() + _step(strat) assert runner.calls == 2 assert runner.step_calls == 2 @@ -262,10 +276,10 @@ def test_grpo_accepts_new_rollout_result(device): runner = _RecordingRunner(_make_rollout_result(device=device)) strat.set_rollout_runner(runner) strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)}) - strat.on_optimizer_step() + _step(strat) runner.swap_result(_make_rollout_result(device=device)) strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)}) - strat.on_optimizer_step() + _step(strat) assert runner.calls == 2 assert runner.step_calls == 2 @@ -280,10 +294,10 @@ def test_dpo_no_sync_hook_when_new_rollout_result(device): runner = _RecordingRunner(_make_rollout_result(device=device)) strat.set_rollout_runner(runner) strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)}) - strat.on_optimizer_step() + _step(strat) runner.swap_result(_make_rollout_result(device=device)) strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)}) - strat.on_optimizer_step() + _step(strat) assert runner.step_calls == 2 @@ -302,12 +316,36 @@ def test_step_called_when_sync_gradients_true(device): runner = _RecordingRunner(_make_rollout_result(device=device)) strat.set_rollout_runner(runner) strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)}) - strat.on_optimizer_step() + _step(strat) assert runner.step_calls == 1 assert runner.weight_updates == [1] assert strat.policy_version == 1 +def test_post_hoc_online_optimizer_step_is_rejected(device): + strat = _make_grpo(device) + strat.set_rollout_runner(_RecordingRunner(_make_rollout_result(device=device))) + + with pytest.raises(RuntimeError, match="strategy.optimizer_step"): + strat.on_optimizer_step() + + +def test_optimizer_step_publishes_version_with_weight_update(device): + strat = _make_grpo(device) + runner = _RecordingRunner(_make_rollout_result(device=device)) + strat.set_rollout_runner(runner) + parameter = next(strat.model.parameters()) + parameter.grad = torch.ones_like(parameter) + optimizer = torch.optim.SGD(strat.model.parameters(), lr=0.1) + before = parameter.detach().clone() + + strat.optimizer_step(optimizer) + + assert not torch.equal(parameter, before) + assert runner.weight_updates == [1] + assert runner.step_calls == 1 + + def test_loss_is_differentiable_dpo(device): strat = _make_dpo(device) strat.set_rollout_runner(_RecordingRunner(_make_rollout_result(device=device))) diff --git a/tests/trainer/test_rollout.py b/tests/trainer/test_rollout.py index a36ad6d9..e22c04a1 100644 --- a/tests/trainer/test_rollout.py +++ b/tests/trainer/test_rollout.py @@ -1,5 +1,7 @@ """Unit tests for the online rollout module.""" +import threading + import pytest import torch @@ -11,6 +13,7 @@ RolloutGenerator, RolloutResult, RolloutRunner, + RolloutVersionError, ) from tests.helpers import FakeTokenizer, make_model @@ -151,6 +154,113 @@ def recording_run_batch(*args, **kwargs): assert model.training is True +def test_rollout_generator_serializes_generation_and_policy_update(device): + gen, _ = _make_generator(device, group_size=1, max_tokens=2) + generation_started = threading.Event() + allow_generation_to_finish = threading.Event() + update_finished = threading.Event() + thread_errors = [] + original = gen._generate_eval + + def blocking_generate(batch, generation_version): + generation_started.set() + assert allow_generation_to_finish.wait(timeout=5) + return original(batch, generation_version) + + gen._generate_eval = blocking_generate + + def generate(): + try: + gen.generate(_make_instruction_batch(n=1)) + except BaseException as exc: + thread_errors.append(exc) + + def apply_update(): + try: + gen.apply_weight_update(1, update_finished.set) + except BaseException as exc: + thread_errors.append(exc) + + generation_thread = threading.Thread(target=generate) + update_thread = threading.Thread(target=apply_update) + generation_thread.start() + assert generation_started.wait(timeout=5) + update_thread.start() + assert not update_finished.wait(timeout=0.1) + + allow_generation_to_finish.set() + generation_thread.join(timeout=5) + update_thread.join(timeout=5) + assert not generation_thread.is_alive() + assert not update_thread.is_alive() + assert thread_errors == [] + assert update_finished.is_set() + assert gen.policy_version == 1 + + +def test_rollout_generator_serializes_direct_scheduler_update(device): + gen, _ = _make_generator(device, group_size=1, max_tokens=2) + generation_started = threading.Event() + allow_generation_to_finish = threading.Event() + update_finished = threading.Event() + thread_errors = [] + original = gen._generate_eval + + def blocking_generate(batch, generation_version): + generation_started.set() + assert allow_generation_to_finish.wait(timeout=5) + return original(batch, generation_version) + + gen._generate_eval = blocking_generate + rollout = [] + + def generate(): + try: + rollout.append(gen.generate(_make_instruction_batch(n=1))) + except BaseException as exc: + thread_errors.append(exc) + + def update_scheduler_directly(): + try: + gen.scheduler.update_weights(1) + update_finished.set() + except BaseException as exc: + thread_errors.append(exc) + + generation_thread = threading.Thread(target=generate) + update_thread = threading.Thread(target=update_scheduler_directly) + generation_thread.start() + assert generation_started.wait(timeout=5) + update_thread.start() + assert not update_finished.wait(timeout=0.1) + + allow_generation_to_finish.set() + generation_thread.join(timeout=5) + update_thread.join(timeout=5) + assert not generation_thread.is_alive() + assert not update_thread.is_alive() + assert thread_errors == [] + assert rollout[0].policy_version == 0 + assert gen.policy_version == 1 + + +def test_rollout_generator_keeps_generation_start_version(device): + gen, _ = _make_generator(device, group_size=1, max_tokens=2) + original_run_batch = gen.scheduler.run_batch + + def update_after_generation(*args, **kwargs): + result = original_run_batch(*args, **kwargs) + gen.scheduler.update_weights(1) + return result + + gen.scheduler.run_batch = update_after_generation + + rollout = gen.generate(_make_instruction_batch(n=1)) + + assert rollout.policy_version == 0 + assert gen.policy_version == 1 + + def test_rollout_generator_mask_matches_responses(device): """Positions beyond a response's length are pad (mask False).""" gen, _ = _make_generator(device, group_size=2, max_tokens=6) @@ -248,6 +358,7 @@ def _make_runner(device, **kw): generator=generator, reward_model=rm, rollout_interval=kw.get("rollout_interval", 2), + max_policy_lag=kw.get("max_policy_lag"), ), model, ) @@ -297,6 +408,114 @@ def test_rollout_runner_tags_generation_version_and_preserves_cached_behavior(de assert refreshed.policy_version == 1 +def test_rollout_runner_rejects_future_generation_version(device): + runner, _ = _make_runner(device, rollout_interval=2) + raw = runner.generator.generate(_make_instruction_batch(n=1)) + raw.policy_version = runner.policy_version + 1 + runner.generator.generate = lambda _batch: raw + + with pytest.raises(RolloutVersionError, match="future policy version"): + runner(_make_instruction_batch(n=1)) + + +def test_rollout_runner_rejects_result_beyond_max_policy_lag(device): + runner, _ = _make_runner(device, rollout_interval=4, max_policy_lag=1) + batch = _make_instruction_batch(n=1) + result, _ = runner(batch) + assert result.policy_version == 0 + + runner.update_weights(2) + with pytest.raises(RolloutVersionError, match="exceeds max_policy_lag=1"): + runner(batch) + + +def test_rollout_runner_revalidates_version_after_async_scoring(device): + runner, _ = _make_runner(device, rollout_interval=4, max_policy_lag=0) + original_score = runner._score + + def score_while_policy_advances(raw): + result = original_score(raw) + runner.update_weights(1) + return result + + runner._score = score_while_policy_advances + + with pytest.raises(RolloutVersionError, match="exceeds max_policy_lag=0"): + runner(_make_instruction_batch(n=1)) + assert runner._cache is None + + +def test_rollout_runner_publishes_cache_before_concurrent_policy_update(device): + runner, _ = _make_runner(device, rollout_interval=4, max_policy_lag=1) + final_validation_started = threading.Event() + allow_final_validation_to_finish = threading.Event() + update_finished = threading.Event() + rollout_finished = threading.Event() + thread_errors = [] + validation_calls = 0 + original_validate = runner._validate_policy_version + + def blocking_validate(result, *, live_version=None): + nonlocal validation_calls + validation_calls += 1 + original_validate(result, live_version=live_version) + if validation_calls == 2: + final_validation_started.set() + assert allow_final_validation_to_finish.wait(timeout=5) + + runner._validate_policy_version = blocking_validate + + def produce_rollout(): + try: + runner(_make_instruction_batch(n=1)) + rollout_finished.set() + except BaseException as exc: + thread_errors.append(exc) + + def apply_update(): + try: + runner.apply_weight_update(1, update_finished.set) + except BaseException as exc: + thread_errors.append(exc) + + rollout_thread = threading.Thread(target=produce_rollout) + update_thread = threading.Thread(target=apply_update) + rollout_thread.start() + assert final_validation_started.wait(timeout=5) + update_thread.start() + assert not update_finished.wait(timeout=0.1) + + allow_final_validation_to_finish.set() + rollout_thread.join(timeout=5) + update_thread.join(timeout=5) + assert not rollout_thread.is_alive() + assert not update_thread.is_alive() + assert thread_errors == [] + assert rollout_finished.is_set() + assert update_finished.is_set() + assert runner._cache is not None + assert runner._cache.policy_version == 0 + assert runner.policy_version == 1 + + +def test_rollout_runner_derives_default_policy_lag_from_interval(device): + runner, _ = _make_runner(device, rollout_interval=4) + assert runner.max_policy_lag == 3 + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"rollout_interval": 0}, "rollout_interval must be positive"), + ({"max_policy_lag": -1}, "max_policy_lag must be non-negative"), + ], +) +def test_rollout_runner_rejects_invalid_version_window(device, kwargs, message): + generator, _ = _make_generator(device) + with pytest.raises(ValueError, match=message): + RolloutRunner(generator, ConstantRewardModel(), **kwargs) + + def test_rollout_runner_refreshes_for_different_batch(device): runner, _ = _make_runner(device, rollout_interval=100) r1, fresh1 = runner(_make_instruction_batch(n=1))