diff --git a/astrai/inference/cache/pool.py b/astrai/inference/cache/pool.py index c6302d5..90761b4 100644 --- a/astrai/inference/cache/pool.py +++ b/astrai/inference/cache/pool.py @@ -406,22 +406,6 @@ class TaskCacheManager: """True if the last bind was a steady-state increment (same tasks, +1 seq_lens).""" return self._bind_was_steady - def last_task_signature_matches(self, task_ids: List[str]) -> bool: - """Check if task_ids match the previous bind's signature. - - Used by Executor to detect steady-state decode for device-to-device - token copy optimization. - """ - if self._bind_state is None: - return False - prev_sig = self._bind_state.sig - # sig is tuple of req_indices, need to map task_ids to req_indices - try: - current_sig = tuple(self._states[tid].req_idx for tid in task_ids) - return prev_sig == current_sig - except KeyError: - return False - # -- internals -- def _rollback(self, state: TaskCacheState, task_id: str): diff --git a/astrai/inference/runtime/executor.py b/astrai/inference/runtime/executor.py index 54f3585..d68d6ab 100644 --- a/astrai/inference/runtime/executor.py +++ b/astrai/inference/runtime/executor.py @@ -415,14 +415,14 @@ class Executor: # inference-mode context. # # ``cache_valid`` checks the decode cache's own task signature: - # req-index signatures in the cache manager are recycled when freed - # slots are reallocated to new tasks, so a fresh batch whose prefill - # re-bind coincides with a stale signature would otherwise replay a - # previous generation's tokens into ``input_ids``. - task_sig_match = self.task_cache.last_task_signature_matches(task_ids) + # task ids are globally unique, so equality alone proves the cached + # tokens were sampled for exactly this ordered batch. Req-index + # signatures in the cache manager are deliberately NOT consulted — + # they are recycled when freed slots are reallocated, which once + # let a fresh batch replay a previous generation's tokens. cached = self._decode_cache cache_valid = cached is not None and cached.task_sig == task_sig - if task_sig_match and cache_valid and cached.last_tokens is not None: + if cache_valid and cached.last_tokens is not None: with torch.inference_mode(): input_ids = ws.fill_input_ids_from_device(cached.last_tokens) else: @@ -433,12 +433,9 @@ class Executor: kv_cache = self.task_cache.bind(task_ids, ws) # Reuse sampling state only if all conditions hold: - # 1. KV bind detected steady increment (same req_indices, seq_lens +1) - # 2. Task signature matches (same task_ids in same order) - # 3. We have a valid cached decode state for THIS task set - reuse_decode_state = ( - cache_valid and self.task_cache.bind_was_steady and task_sig_match - ) + # 1. The cached decode state belongs to THIS task set (task_sig) + # 2. KV bind detected steady increment (same req_indices, seq_lens +1) + reuse_decode_state = cache_valid and self.task_cache.bind_was_steady if reuse_decode_state: info = cached.sampling_info ws.position_ids[:b] += 1 diff --git a/astrai/inference/runtime/stepper.py b/astrai/inference/runtime/stepper.py new file mode 100644 index 0000000..56f8b13 --- /dev/null +++ b/astrai/inference/runtime/stepper.py @@ -0,0 +1,127 @@ +"""One-token advancement primitive shared by every scheduling mode.""" + +from contextlib import nullcontext +from typing import Dict, List, Optional, Tuple + +from astrai.extension import AttentionBackend, attn_backend +from astrai.inference.cache import PagePool, TaskCacheManager +from astrai.inference.metrics import MetricsCollector +from astrai.inference.runtime.executor import Executor +from astrai.inference.task import Task, TaskStatus + + +class Stepper: + """Advance every active task by one token (prefill + decode). + + Single shared primitive for both the continuous-batching loop and the + synchronous ``run_batch`` path, so the two cannot drift. + + Tasks must already be allocated in the KV cache. Tasks without output + are prefilled first and sample their first token from the final prompt + position. Tasks with output extend the cache by one position and decode + from their latest generated token. + """ + + def __init__( + self, + pool: PagePool, + task_cache: TaskCacheManager, + executor: Executor, + metrics: MetricsCollector, + ): + self._pool = pool + self._task_cache = task_cache + self._executor = executor + self._metrics = metrics + + @staticmethod + def _task_backend_groups(tasks: List[Task]): + groups = {} + for task in tasks: + groups.setdefault(task.backend, (task.backend, []))[1].append(task) + return groups.values() + + def step( + self, tasks: List[Task], return_logprobs: bool = False + ) -> Tuple[List[Task], List[Task]]: + """Advance ``tasks`` by one token. + + Args: + tasks: Active tasks to advance by one token. + return_logprobs: Forwarded to the executor; per-token logprobs + are recorded on each task's ``output_logprobs``. + + Returns: + ``(decoded, aborted)``: tasks that produced a new token (its ID + already appended to ``output_ids``) and tasks that hit the + sequence cap and were marked ``ABORTED``. + """ + to_prefill = [t for t in tasks if not t.prefill_done and t.prompt_ids] + prefilled_ids = set() + produced: List[Task] = [] + if to_prefill: + for t in to_prefill: + t.input_tokens = len(t.prompt_ids) + + groups: Dict[Tuple[int, Optional[AttentionBackend]], List[Task]] = {} + for t in to_prefill: + start_pos = min( + self._task_cache.task_cached(t.task_id), len(t.prompt_ids) - 1 + ) + groups.setdefault((start_pos, t.backend), []).append(t) + + for (start_pos, _), group in groups.items(): + backend = group[0].backend + backend_context = ( + attn_backend(backend) if backend is not None else nullcontext() + ) + with ( + backend_context, + self._metrics.record([t.task_id for t in group], "prefill"), + ): + prefilled, step_out = self._executor.execute_prefill( + group, start_pos=start_pos, return_logprobs=return_logprobs + ) + + for t, out in zip(prefilled, step_out): + t.output_ids.append(out[0] if return_logprobs else out) + t.output_tokens += 1 + t.mark_prefill_done() + prefilled_ids.add(t.task_id) + produced.append(t) + + start_logical_page = start_pos // self._pool.page_size + for t in group: + self._task_cache.task_record_hashes( + t.task_id, t.prompt_ids, start_logical_page + ) + + decoded: List[Task] = [] + aborted: List[Task] = [] + for t in tasks: + if t.task_id in prefilled_ids: + continue + if self._task_cache.task_extend(t.task_id, t.next_pos): + decoded.append(t) + else: + t.status = TaskStatus.ABORTED + aborted.append(t) + + for backend, group in self._task_backend_groups(decoded): + backend_context = ( + attn_backend(backend) if backend is not None else nullcontext() + ) + with ( + backend_context, + self._metrics.record([t.task_id for t in group], "decode"), + ): + step_out = self._executor.execute_decode( + group, return_logprobs=return_logprobs + ) + for t, out in zip(group, step_out): + t.output_ids.append(out[0] if return_logprobs else out) + t.output_tokens += 1 + t.advance_kv() + produced.append(t) + + return produced, aborted diff --git a/astrai/inference/scheduler.py b/astrai/inference/scheduler.py index 879e7ac..f22d8bc 100644 --- a/astrai/inference/scheduler.py +++ b/astrai/inference/scheduler.py @@ -16,6 +16,7 @@ from astrai.extension import ( from astrai.inference.cache import PagePool, TaskCacheManager from astrai.inference.metrics import MetricsCollector from astrai.inference.runtime.executor import Executor +from astrai.inference.runtime.stepper import Stepper from astrai.inference.task import ( STOP, GenerationResult, @@ -23,6 +24,7 @@ from astrai.inference.task import ( TaskManager, TaskStatus, ) +from astrai.inference.versioning import PolicyVersionGuard from astrai.model.automodel import AutoModel from astrai.tokenize.tokenizer import AutoTokenizer @@ -119,49 +121,32 @@ class InferenceScheduler: enable_cuda_graph=enable_cuda_graph, ) + self._stepper = Stepper( + self._cache, self._task_cache, self._executor, self._metrics + ) + self._stop_event = threading.Event() self._loop_thread: Optional[threading.Thread] = None - self._weight_lock = threading.RLock() - self._policy_version = policy_version + self._policy_guard = PolicyVersionGuard( + policy_version, + ensure_ready=self._ensure_weight_update_ready, + on_commit=self._task_cache.invalidate_cache, + ) + # Synchronous generation shares the guard's generation/weight mutex. + self._weight_lock = self._policy_guard.lock @property def policy_version(self) -> int: """Version of the model weights used for subsequent generations.""" - return self._policy_version - - def _validate_weight_version( - 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" - ) + return self._policy_guard.policy_version def _ensure_weight_update_ready(self) -> None: - """Check weight update preconditions. Must be called under _weight_lock.""" + """Check weight update preconditions. Must be called under the lock.""" if self._loop_thread is not None and self._loop_thread.is_alive(): raise RuntimeError("Stop the scheduler before updating model weights") if self._task_mgr.get_active_tasks() or self._task_mgr.get_waiting_tasks(): raise RuntimeError("Cannot update model weights while tasks are queued") - def _commit_weight_version(self, policy_version: int) -> int: - self._task_cache.invalidate_cache() - self._policy_version = policy_version - return self._policy_version - - @_with_weight_lock def update_weights(self, policy_version: int) -> int: """Acknowledge an in-place weight update and invalidate stale KV state. @@ -170,13 +155,8 @@ class InferenceScheduler: version update makes that lifecycle visible and prevents prefix KV entries produced by older weights from being reused. """ - self._validate_weight_version(policy_version) - if policy_version == self._policy_version: - return self._policy_version - self._ensure_weight_update_ready() - return self._commit_weight_version(policy_version) + return self._policy_guard.update_weights(policy_version) - @_with_weight_lock def apply_weight_update( self, policy_version: Optional[int], update: Callable[[], T] ) -> T: @@ -186,24 +166,11 @@ class InferenceScheduler: 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_weight_version(policy_version, require_advance=True) - self._ensure_weight_update_ready() + return self._policy_guard.apply_weight_update(policy_version, update) - result = update() - self._commit_weight_version(policy_version) - return result - - @_with_weight_lock def with_policy_snapshot(self, inspect: Callable[[int], T]) -> T: """Inspect state while the scheduler's policy version remains stable.""" - if not callable(inspect): - raise TypeError("inspect must be callable") - return inspect(self._policy_version) + return self._policy_guard.with_policy_snapshot(inspect) def add_task(self, prompt: str, **kwargs) -> str: return self._task_mgr.add_task(prompt, **kwargs) @@ -226,7 +193,7 @@ class InferenceScheduler: def get_stats(self) -> Dict[str, Any]: stats = self._task_mgr.get_stats() stats["kv_cache_tasks"] = self._task_cache.task_count - stats["policy_version"] = self._policy_version + stats["policy_version"] = self._policy_guard.policy_version return stats @property @@ -242,105 +209,11 @@ class InferenceScheduler: return nullcontext() return attn_backend(self._backend) - @staticmethod - def _task_backend_groups(tasks: List[Task]): - groups = {} - for task in tasks: - groups.setdefault(task.backend, (task.backend, []))[1].append(task) - return groups.values() - def _step( self, tasks: List[Task], return_logprobs: bool = False ) -> Tuple[List[Task], List[Task]]: - """Advance every active task by one token (prefill + decode). - - Single shared primitive for both the continuous-batching loop and - the synchronous ``run_batch`` path, so the two cannot drift. - - Tasks must already be allocated in the KV cache. Tasks without output - are prefilled first and sample their first token from the final prompt - position. Tasks with output extend the cache by one position and decode - from their latest generated token. - - Args: - tasks: Active tasks to advance by one token. - return_logprobs: Forwarded to ``execute_decode``; per-token - logprobs are recorded on each task's ``output_logprobs``. - - Returns: - ``(decoded, aborted)``: tasks that produced a new token (its ID - already appended to ``output_ids``) and tasks that hit the - sequence cap and were marked ``ABORTED``. - """ - to_prefill = [t for t in tasks if not t.prefill_done and t.prompt_ids] - prefilled_ids = set() - produced: List[Task] = [] - if to_prefill: - for t in to_prefill: - t.input_tokens = len(t.prompt_ids) - - groups: Dict[Tuple[int, Optional[AttentionBackend]], List[Task]] = {} - for t in to_prefill: - start_pos = min( - self._task_cache.task_cached(t.task_id), len(t.prompt_ids) - 1 - ) - groups.setdefault((start_pos, t.backend), []).append(t) - - for (start_pos, _), group in groups.items(): - backend = group[0].backend - backend_context = ( - attn_backend(backend) if backend is not None else nullcontext() - ) - with ( - backend_context, - self._metrics.record([t.task_id for t in group], "prefill"), - ): - prefilled, step_out = self._executor.execute_prefill( - group, start_pos=start_pos, return_logprobs=return_logprobs - ) - - for t, out in zip(prefilled, step_out): - t.output_ids.append(out[0] if return_logprobs else out) - t.output_tokens += 1 - t.mark_prefill_done() - prefilled_ids.add(t.task_id) - produced.append(t) - - start_logical_page = start_pos // self._cache.page_size - for t in group: - self._task_cache.task_record_hashes( - t.task_id, t.prompt_ids, start_logical_page - ) - - decoded: List[Task] = [] - aborted: List[Task] = [] - for t in tasks: - if t.task_id in prefilled_ids: - continue - if self._task_cache.task_extend(t.task_id, t.next_pos): - decoded.append(t) - else: - t.status = TaskStatus.ABORTED - aborted.append(t) - - for backend, group in self._task_backend_groups(decoded): - backend_context = ( - attn_backend(backend) if backend is not None else nullcontext() - ) - with ( - backend_context, - self._metrics.record([t.task_id for t in group], "decode"), - ): - step_out = self._executor.execute_decode( - group, return_logprobs=return_logprobs - ) - for t, out in zip(group, step_out): - t.output_ids.append(out[0] if return_logprobs else out) - t.output_tokens += 1 - t.advance_kv() - produced.append(t) - - return produced, aborted + """Advance every active task by one token; see :class:`Stepper`.""" + return self._stepper.step(tasks, return_logprobs=return_logprobs) def _run_generation_loop(self): stop_ids = self._task_mgr.tokenizer.stop_ids @@ -389,7 +262,7 @@ class InferenceScheduler: if task.status != TaskStatus.ABORTED ] - decoded, aborted = self._step(active) + decoded, aborted = self._stepper.step(active) for t in aborted: self._task_mgr.invoke_callback(t.task_id, STOP) @@ -533,7 +406,9 @@ class InferenceScheduler: with self._backend_context(): while live: - decoded, aborted = self._step(live, return_logprobs=return_logprobs) + decoded, aborted = self._stepper.step( + live, return_logprobs=return_logprobs + ) for task in aborted: runtime_errors[task.task_id] = "kv_cache_extension_failed" live = [t for t in decoded if not t.is_finished(stop_ids)] diff --git a/astrai/inference/versioning.py b/astrai/inference/versioning.py new file mode 100644 index 0000000..98685c7 --- /dev/null +++ b/astrai/inference/versioning.py @@ -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 diff --git a/astrai/inference/workspace.py b/astrai/inference/workspace.py index a70d27a..85ef0f1 100644 --- a/astrai/inference/workspace.py +++ b/astrai/inference/workspace.py @@ -80,7 +80,10 @@ class InferenceWorkspace: (max_batch_size,), dtype=torch.long, device=device ) self._pin = torch.empty( - (max_batch_size,), dtype=torch.long, pin_memory=True + (max_batch_size,), + dtype=torch.long, + pin_memory=torch.cuda.is_available() + and torch.device(device).type == "cuda", ) # KV-cache bind metadata (fixed shape, written by diff --git a/tests/inference/test_scheduler.py b/tests/inference/test_scheduler.py index d867381..add9725 100644 --- a/tests/inference/test_scheduler.py +++ b/tests/inference/test_scheduler.py @@ -12,6 +12,7 @@ from astrai.extension import CudaBackend, TorchNativeBackend, get_backend from astrai.inference import GenerationResult, InferenceScheduler from astrai.inference.metrics import MetricsCollector from astrai.inference.runtime.executor import DecodeSteadyState, Executor +from astrai.inference.runtime.stepper import Stepper from astrai.inference.task import Task from astrai.model.transformer import AutoRegressiveLM from tests.helpers import FakeTokenizer, make_rollout_config @@ -119,10 +120,14 @@ def test_generation_loop_activates_backend_in_worker_thread(): def test_step_splits_decode_batch_by_request_backend(): scheduler = object.__new__(InferenceScheduler) + scheduler._cache = SimpleNamespace(page_size=1) scheduler._task_cache = MagicMock() scheduler._task_cache.task_extend.return_value = True scheduler._metrics = MetricsCollector() scheduler._executor = MagicMock() + scheduler._stepper = Stepper( + scheduler._cache, scheduler._task_cache, scheduler._executor, scheduler._metrics + ) observed = [] @@ -157,6 +162,9 @@ def test_step_batches_ragged_prefill_with_shared_cache_start(): scheduler._task_cache.task_cached.return_value = 0 scheduler._metrics = MetricsCollector() scheduler._executor = MagicMock() + scheduler._stepper = Stepper( + scheduler._cache, scheduler._task_cache, scheduler._executor, scheduler._metrics + ) short = Task("short", [1, 2, 3]) long = Task("long", [4, 5, 6, 7, 8]) @@ -701,8 +709,8 @@ def test_run_batch_details_report_extension_failure_and_cleanup(device): scheduler, _tok, _model = _make_real_scheduler(device) try: with patch.object( - scheduler, - "_step", + scheduler._stepper, + "step", side_effect=lambda tasks, **_kwargs: ([], list(tasks)), ): result = scheduler.run_batch([[10, 20]], max_tokens=2, return_details=True)[ @@ -722,9 +730,6 @@ def test_decode_does_not_reuse_previous_batch_state(): executor.device = torch.device("cpu") executor.task_cache = MagicMock() executor.task_cache.bind_was_steady = True - executor.task_cache.last_task_signature_matches.return_value = ( - False # Different task - ) executor.task_cache.bind.return_value = MagicMock() executor._graph_supported = False executor._graph_ctx = SimpleNamespace(enabled=False) @@ -771,7 +776,6 @@ def test_decode_fills_input_ids_from_device_on_matching_signature(): executor.device = torch.device("cpu") executor.task_cache = MagicMock() executor.task_cache.bind_was_steady = True - executor.task_cache.last_task_signature_matches.return_value = True # Same task executor.task_cache.bind.return_value = MagicMock() executor._graph_supported = False executor._graph_ctx = SimpleNamespace(enabled=False)