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.

tests/inference/test_cache.py

@@ -1,5 +1,6 @@
 """Unit tests for inference cache components."""
 
+import pytest
 import torch
 
 from astrai.inference.cache import (
@@ -435,6 +436,26 @@ def test_page_pool_prefix_hit_populates_request_mapping():
     )
 
 
+def test_task_cache_invalidation_drops_cross_version_prefix_hits():
+    pool = _make_paged_pool_ps64(page_size=2, max_seq_len=8, n_tokens=16)
+    task_cache = _make_task_cache(pool)
+    prompt = [11, 12, 13, 14]
+
+    assert task_cache.task_alloc("first", prompt)
+    task_cache.task_record_hashes("first", prompt)
+    task_cache.task_free("first")
+    assert task_cache.task_alloc("cached", prompt)
+    assert task_cache.task_cached("cached") == len(prompt)
+
+    with pytest.raises(RuntimeError, match="while tasks are active"):
+        task_cache.invalidate_cache()
+
+    task_cache.task_free("cached")
+    assert task_cache.invalidate_cache() == 2
+    assert task_cache.task_alloc("after_update", prompt)
+    assert task_cache.task_cached("after_update") == 0
+
+
 def test_page_pool_paged_ps64_bind_roundtrip():
     pool = _make_paged_pool_ps64(n_layers=1, n_kv_heads=2, head_dim=4)
     task_cache = _make_task_cache(pool)

tests/inference/test_scheduler.py

@@ -545,6 +545,34 @@ def test_run_batch_empty_prompts(device):
         scheduler.stop()
 
 
+def test_scheduler_weight_versions_are_monotonic_and_acknowledged(device):
+    scheduler, _tok, _model = _make_real_scheduler(device)
+    try:
+        assert scheduler.policy_version == 0
+        assert scheduler.update_weights(1) == 1
+        assert scheduler.policy_version == 1
+        assert scheduler.get_stats()["policy_version"] == 1
+        assert scheduler.update_weights(1) == 1
+        with pytest.raises(ValueError, match="cannot move backwards"):
+            scheduler.update_weights(0)
+        with pytest.raises(ValueError, match="non-negative integer"):
+            scheduler.update_weights(True)
+    finally:
+        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")
+    try:
+        with pytest.raises(RuntimeError, match="while tasks are queued"):
+            scheduler.update_weights(1)
+        scheduler.remove_task(task_id)
+        assert scheduler.update_weights(1) == 1
+    finally:
+        scheduler.stop()
+
+
 def test_run_batch_too_long_prompt_skipped(device):
     """A prompt longer than max_seq_len yields an empty result slot."""
     scheduler, _tok, _model = _make_real_scheduler(device)

tests/trainer/test_online_strategy.py

@@ -43,6 +43,8 @@ class _RecordingRunner:
         self.calls = 0
         self.step_calls = 0
         self._fresh = True
+        self.policy_version = result.policy_version
+        self.weight_updates = []
 
     def __call__(self, batch):
         self.calls += 1
@@ -53,6 +55,11 @@ class _RecordingRunner:
     def step(self):
         self.step_calls += 1
 
+    def update_weights(self, policy_version):
+        self.policy_version = policy_version
+        self.weight_updates.append(policy_version)
+        return policy_version
+
     def swap_result(self, result):
         self.result = result
         self._fresh = True
@@ -302,6 +309,8 @@ def test_step_called_when_sync_gradients_true(device):
     strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})
     strat.on_optimizer_step()
     assert runner.step_calls == 1
+    assert runner.weight_updates == [1]
+    assert strat.policy_version == 1
 
 
 def test_loss_is_differentiable_dpo(device):

tests/trainer/test_rollout.py

@@ -65,6 +65,7 @@ def test_raw_rollout_fields():
     )
     assert r.prompts.shape == (2, 4)
     assert r.responses.shape == (2, 3, 5)
+    assert r.policy_version == 0
     assert r.prompt_texts == []
     assert r.response_texts == []
 
@@ -131,6 +132,7 @@ def test_rollout_generator_shapes(device):
     assert len(r.prompt_texts) == 2
     assert len(r.response_texts) == 2
     assert len(r.response_texts[0]) == 3
+    assert r.policy_version == 0
 
 
 def test_rollout_generator_uses_eval_and_restores_mode(device):
@@ -275,6 +277,26 @@ def test_rollout_runner_cache_returns_stale_flag(device):
     assert fresh2 is False
 
 
+def test_rollout_runner_tags_generation_version_and_preserves_cached_behavior(device):
+    runner, _ = _make_runner(device, rollout_interval=100)
+    batch = _make_instruction_batch(n=1)
+
+    first, first_fresh = runner(batch)
+    assert first_fresh is True
+    assert first.policy_version == 0
+
+    assert runner.update_weights(1) == 1
+    cached, cached_fresh = runner(batch)
+    assert cached is first
+    assert cached_fresh is False
+    assert cached.policy_version == 0
+
+    runner.clear_cache()
+    refreshed, refreshed_fresh = runner(batch)
+    assert refreshed_fresh is True
+    assert refreshed.policy_version == 1
+
+
 def test_rollout_runner_refreshes_for_different_batch(device):
     runner, _ = _make_runner(device, rollout_interval=100)
     r1, fresh1 = runner(_make_instruction_batch(n=1))