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

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")