refactor: centralize batch state and split scheduler duties

- 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
2026-09-04 14:44:22 +08:00
parent ae7fc3059a
commit 8e39d9d8c9
7 changed files with 304 additions and 185 deletions
@@ -0,0 +1,129 @@
+"""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