fix: keep async rollouts version-consistent

- serialize shared-model optimizer updates with generation
- reject future or over-lagged rollout results after asynchronous scoring
- close cache publication races
- persist policy versions in online checkpoints

astrai/trainer/rollout.py

@@ -16,7 +16,7 @@ Provides:
 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 @@ class BaseRewardModel(ABC):
 
 
 _PAD = 0
+T = TypeVar("T")
+
+
+class RolloutVersionError(RuntimeError):
+    """A rollout cannot be attributed to an acceptable policy version."""
 
 
 class RolloutGenerator:
@@ -142,6 +147,18 @@ class RolloutGenerator:
         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 @@ class RolloutGenerator:
         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 @@ class RolloutGenerator:
             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 @@ class RolloutRunner:
         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 @@ class RolloutRunner:
         """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 @@ class RolloutRunner:
             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 @@ class RolloutRunner:
             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)

astrai/trainer/strategy.py

@@ -7,6 +7,7 @@ import torch
 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 @@ class BaseStrategy(ABC):
         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."""

astrai/trainer/train_callback.py

@@ -164,6 +164,9 @@ class CheckpointCallback(TrainCallback):
                     **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,

astrai/trainer/train_context.py

@@ -355,5 +355,6 @@ class TrainContextBuilder:
                 generator=generator,
                 reward_model=cfg.reward_model_fn(),
                 rollout_interval=cfg.rollout_interval,
+                max_policy_lag=cfg.rollout_max_policy_lag,
             )
         )

astrai/trainer/trainer.py

@@ -94,8 +94,7 @@ class Trainer:
 
                         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: