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.
This commit is contained in:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user