ViperEkura
8e39d9d8c9
- make decode steady state self-validating: gate the token fill and sampling reuse on the decode cache's own task signature, drop TaskCacheManager.last_task_signature_matches whose req-index signature recycles with slot reuse - extract PolicyVersionGuard (version protocol plus generation/weight mutex) and Stepper (shared one-token advancement) out of the scheduler, keeping its public API unchanged - pin the input staging buffer only on CUDA devices so CPU-only workspaces allocate
130 lines
4.7 KiB
Python
130 lines
4.7 KiB
Python
"""Policy version protocol for weights shared between trainer and server.
|
|
|
|
The guard owns the monotonic version counter, the RLock that serializes
|
|
weight publication against generation (``run_batch`` acquires the same
|
|
lock), and the validation/commit rules around both. Scheduler-specific
|
|
preconditions (no in-flight generation, no queued tasks) and side effects
|
|
(dropping stale KV entries) are injected as callables so the guard stays
|
|
free of inference-subsystem knowledge.
|
|
"""
|
|
|
|
import threading
|
|
from functools import wraps
|
|
from typing import Callable, Optional, TypeVar
|
|
|
|
T = TypeVar("T")
|
|
|
|
|
|
def _locked(method):
|
|
@wraps(method)
|
|
def synchronized(self, *args, **kwargs):
|
|
with self._lock:
|
|
return method(self, *args, **kwargs)
|
|
|
|
return synchronized
|
|
|
|
|
|
class PolicyVersionGuard:
|
|
"""Monotonic policy-version protocol over shared in-place weights.
|
|
|
|
The scheduler and the in-process trainer mutate the same model object;
|
|
every weight mutation must publish a new version atomically under the
|
|
generation lock. Versions never move backwards and every mutation
|
|
advances (``apply_weight_update``) or repeats the live version
|
|
(``update_weights`` idempotently).
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
policy_version: int,
|
|
ensure_ready: Callable[[], None],
|
|
on_commit: Callable[[], None],
|
|
):
|
|
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")
|
|
self._lock = threading.RLock()
|
|
self._policy_version = policy_version
|
|
self._ensure_ready = ensure_ready
|
|
self._on_commit = on_commit
|
|
|
|
@property
|
|
def lock(self) -> threading.RLock:
|
|
"""Generation/weight mutex; synchronous generation acquires it too."""
|
|
return self._lock
|
|
|
|
@property
|
|
def policy_version(self) -> int:
|
|
"""Version of the model weights used for subsequent generations."""
|
|
return self._policy_version
|
|
|
|
def _validate(self, policy_version: int, *, require_advance: bool = False) -> None:
|
|
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 require_advance and policy_version == self._policy_version:
|
|
raise ValueError(
|
|
f"policy_version must advance beyond {self._policy_version} "
|
|
"when model weights are mutated"
|
|
)
|
|
|
|
@_locked
|
|
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.
|
|
"""
|
|
self._validate(policy_version)
|
|
if policy_version == self._policy_version:
|
|
return self._policy_version
|
|
self._ensure_ready()
|
|
return self._commit(policy_version)
|
|
|
|
@_locked
|
|
def apply_weight_update(
|
|
self, policy_version: Optional[int], update: Callable[[], T]
|
|
) -> T:
|
|
"""Mutate shared weights and publish their version without generation.
|
|
|
|
``policy_version=None`` derives ``live + 1`` under the same lock, for
|
|
callers that only need "advance by one" (e.g. ``optimizer.step()``)
|
|
without a read-compute-write race on the current version.
|
|
"""
|
|
if not callable(update):
|
|
raise TypeError("update must be callable")
|
|
if policy_version is None:
|
|
policy_version = self._policy_version + 1
|
|
else:
|
|
self._validate(policy_version, require_advance=True)
|
|
self._ensure_ready()
|
|
|
|
result = update()
|
|
self._commit(policy_version)
|
|
return result
|
|
|
|
@_locked
|
|
def with_policy_snapshot(self, inspect: Callable[[int], T]) -> T:
|
|
"""Inspect state while the policy version remains stable."""
|
|
if not callable(inspect):
|
|
raise TypeError("inspect must be callable")
|
|
return inspect(self._policy_version)
|
|
|
|
def _commit(self, policy_version: int) -> int:
|
|
self._on_commit()
|
|
self._policy_version = policy_version
|
|
return self._policy_version
|