feat: version rollout weight updates

Track a monotonic policy version across optimizer steps, scheduler updates, and rollout results. Serialize synchronous generation with weight acknowledgements and invalidate reusable prefix KV entries so cached samples remain attributable to the behavior policy that generated them.

astrai/inference/cache/pool.py

@@ -338,6 +338,14 @@ class TaskCacheManager:
         if state is not None:
             self._strategy.record_hashes(state, prompt_ids, start_logical_page)
 
+    def invalidate_cache(self) -> int:
+        """Drop reusable KV entries once all task-owned entries are released."""
+        if self._states:
+            raise RuntimeError("Cannot invalidate KV cache while tasks are active")
+        self._bind_state = None
+        self._bind_was_steady = False
+        return self._strategy.invalidate_cache()
+
     @staticmethod
     def task_cacheable_ids(task_id: str, prompt_ids: List[int], output_ids: List[int]):
         return list(prompt_ids) + list(output_ids[:-1])

astrai/inference/cache/strategy.py

@@ -89,6 +89,19 @@ class Allocator:
             if idx in self._lru:
                 self._lru.move_to_end(idx)
 
+    def clear_cached(self) -> int:
+        """Release every unreferenced LRU page back to the free pool."""
+        with self._lock:
+            cached = list(self._lru)
+            self._lru.clear()
+            for idx in cached:
+                if self._refs[idx] != 0:
+                    raise RuntimeError("Cannot invalidate a referenced cache page")
+                if self.on_evict:
+                    self.on_evict(idx)
+                self._free_mask |= 1 << idx
+            return len(cached)
+
 
 class RadixNode:
     """A page-aligned edge in the CPU-side prefix radix trie."""
@@ -200,6 +213,10 @@ class AllocationStrategy(ABC):
         start: int,
     ) -> None: ...
 
+    def invalidate_cache(self) -> int:
+        """Drop reusable KV entries after an inference weight update."""
+        return 0
+
 
 class ContiguousStrategy(AllocationStrategy):
     """Static contiguous allocation: slots are pre-assigned at pool init.
@@ -316,3 +333,8 @@ class PagedStrategy(AllocationStrategy):
         full = len(prompt_ids) // self._page_size
         for i in range(start, min(full, len(state.pages))):
             self._prefix.record(state.pages[i], prompt_ids, i)
+
+    def invalidate_cache(self) -> int:
+        if self._prefix is None:
+            return 0
+        return self._alloc.clear_cached()

astrai/inference/scheduler.py

@@ -2,6 +2,7 @@ import logging
 import threading
 import uuid
 from contextlib import nullcontext
+from functools import wraps
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import torch
@@ -28,6 +29,15 @@ from astrai.tokenize.tokenizer import AutoTokenizer
 logger = logging.getLogger(__name__)
 
 
+def _with_weight_lock(method):
+    @wraps(method)
+    def synchronized(self, *args, **kwargs):
+        with self._weight_lock:
+            return method(self, *args, **kwargs)
+
+    return synchronized
+
+
 class InferenceScheduler:
     """Continuous batching loop: cleanup -> refill -> prefill -> decode (all groups)."""
 
@@ -42,7 +52,14 @@ class InferenceScheduler:
         cache: Optional[PagePool] = None,
         enable_cuda_graph: bool = True,
         backend: Optional[Union[str, ATTN_BACKEND, AttentionBackend, type]] = None,
+        policy_version: int = 0,
     ):
+        if (
+            isinstance(policy_version, bool)
+            or not isinstance(policy_version, int)
+            or policy_version < 0
+        ):
+            raise ValueError("policy_version must be a non-negative integer")
         config = model.config
 
         if max_seq_len is not None:
@@ -103,6 +120,44 @@ class InferenceScheduler:
 
         self._stop_event = threading.Event()
         self._loop_thread: Optional[threading.Thread] = None
+        self._weight_lock = threading.RLock()
+        self._policy_version = policy_version
+
+    @property
+    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.
+        """
+        if (
+            isinstance(policy_version, bool)
+            or not isinstance(policy_version, int)
+            or policy_version < 0
+        ):
+            raise ValueError("policy_version must be a non-negative integer")
+        if policy_version < self._policy_version:
+            raise ValueError(
+                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 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")
+
+        self._task_cache.invalidate_cache()
+        self._policy_version = policy_version
+        return self._policy_version
 
     def add_task(self, prompt: str, **kwargs) -> str:
         return self._task_mgr.add_task(prompt, **kwargs)
@@ -125,6 +180,7 @@ class InferenceScheduler:
     def get_stats(self) -> Dict[str, Any]:
         stats = self._task_mgr.get_stats()
         stats["kv_cache_tasks"] = self._task_cache.task_count
+        stats["policy_version"] = self._policy_version
         return stats
 
     @property
@@ -343,6 +399,7 @@ class InferenceScheduler:
             )
         self._task_mgr.clear_queues()
 
+    @_with_weight_lock
     def run_batch(
         self,
         prompt_ids_list: List[List[int]],

astrai/trainer/rollout.py

@@ -13,6 +13,7 @@ Provides:
   so callers do not need to rely on object identity to detect refreshes.
 """
 
+import threading
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
 from typing import Dict, List, Optional, Tuple
@@ -53,6 +54,7 @@ class RawRollout:
     responses: Tensor
     response_mask: Tensor
     logprobs_old: Tensor
+    policy_version: int = 0
     prompt_texts: List[str] = field(default_factory=list)
     response_texts: List[List[str]] = field(default_factory=list)
 
@@ -129,6 +131,16 @@ class RolloutGenerator:
         self.top_p = top_p
         self.frequency_penalty = frequency_penalty
         self.rep_window = rep_window
+        self._weight_lock = threading.RLock()
+
+    @property
+    def policy_version(self) -> int:
+        return self.scheduler.policy_version
+
+    def update_weights(self, policy_version: int) -> int:
+        """Acknowledge shared-model weights and invalidate older scheduler KV."""
+        with self._weight_lock:
+            return self.scheduler.update_weights(policy_version)
 
     @torch.no_grad()
     def generate(self, batch: Dict) -> RawRollout:
@@ -146,13 +158,14 @@ class RolloutGenerator:
         ``add_generation_prompt=True`` so rollout prompts match the
         format the policy was SFT-trained on.
         """
-        model = self.scheduler._executor.model
-        was_training = model.training
-        model.eval()
-        try:
-            return self._generate_eval(batch)
-        finally:
-            model.train(was_training)
+        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:
         prompt_texts, flat_prompt_ids = self._prepare_prompts(batch)
@@ -245,6 +258,7 @@ class RolloutGenerator:
             responses=responses,
             response_mask=response_mask,
             logprobs_old=logprobs_old,
+            policy_version=self.policy_version,
             prompt_texts=prompt_texts,
             response_texts=response_texts,
         )
@@ -370,6 +384,14 @@ class RolloutRunner:
         self._cache_key = None
         self._steps_since_rollout: int = 0
 
+    @property
+    def policy_version(self) -> int:
+        return self.generator.policy_version
+
+    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 step(self):
         """Advance the internal counter (call once per optimizer step)."""
         self._steps_since_rollout += 1
@@ -415,6 +437,7 @@ class RolloutRunner:
             response_mask=raw.response_mask,
             rewards=rewards.to(device=device),
             logprobs_old=raw.logprobs_old,
+            policy_version=raw.policy_version,
             prompt_texts=raw.prompt_texts,
             response_texts=raw.response_texts,
         )

astrai/trainer/strategy.py

@@ -239,6 +239,12 @@ class BaseStrategy(ABC):
         """Inject a :class:`RolloutRunner` to enable online rollout mode."""
         self._rollout_runner = runner
 
+    @property
+    def policy_version(self) -> Optional[int]:
+        if self._rollout_runner is None:
+            return None
+        return self._rollout_runner.policy_version
+
     def prepare_from_rollout(self, result: RolloutResult) -> Dict[str, Tensor]:
         """Map a :class:`RolloutResult` to the batch layout expected by
         :meth:`compute_loss`.
@@ -275,6 +281,7 @@ class BaseStrategy(ABC):
     def on_optimizer_step(self):
         """Advance online rollout state after a successful optimizer step."""
         if self._rollout_runner is not None:
+            self._rollout_runner.update_weights(self.policy_version + 1)
             self._rollout_runner.step()
 
     def __call__(self, batch: Dict[str, Tensor]) -> LossOutput:

astrai/trainer/train_context.py

@@ -318,6 +318,11 @@ class TrainContextBuilder:
             tokenizer=tokenizer,
             max_batch_size=group_size * max(1, cfg.batch_per_device),
             max_seq_len=getattr(context.model.config, "max_position_embeddings", None),
+            policy_version=(
+                context.checkpoint.meta.get("policy_version", context.optimizer_step)
+                if context.checkpoint is not None
+                else context.optimizer_step
+            ),
         )
         generator = RolloutGenerator(
             scheduler=scheduler,