Merge pull request #229 from FluffyAIcode/AgentMemory/decode-worker-sink-window-fix-0722

fix(runtime): stabilize isolated Prefill sessions

Author: FluffyAIcode

Dockerfile

@@ -44,6 +44,7 @@ WORKDIR /app
 # pins all the runtime packages we ship.
 COPY requirements.txt ./
 RUN pip install --upgrade pip \
+    && pip install --index-url https://download.pytorch.org/whl/cpu "torch>=2.4,<3.0" \
     && pip install -r requirements.txt
 
 # Application source. `.dockerignore` prunes tests/, results/, .git/,

inference_engine/backends/mlx/decode_worker.py

@@ -56,6 +56,10 @@ def __init__(self, error_type: str, message: str) -> None:
         self.message = message
 
 
+class DecodeWorkerSessionClosed(DecodeWorkerError):
+    """A verifier proxy was used after its router session was closed."""
+
+
 @dataclass(frozen=True)
 class DecodeWorkerConfig:
     model_id: str
@@ -824,6 +828,15 @@ def _mark_current(self, session_id: str) -> None:
         with self._lock:
             self._restored_generation[session_id] = self._generation
 
+    def _checkpoint_locked(self, session_id: str) -> ProofCheckpoint:
+        """Return a live checkpoint while the caller holds ``_lock``."""
+        try:
+            return self._checkpoints[session_id]
+        except KeyError as exc:
+            raise DecodeWorkerSessionClosed(
+                f"decode session {session_id!r} is closed"
+            ) from exc
+
 
 class DecodeWorkerSession:
     """Verifier-shaped proxy used by existing append/generate coordinators."""
@@ -837,6 +850,16 @@ def __init__(self, client: DecodeWorkerClient, session_id: str) -> None:
         self.next_global_position = 0
         self._kv_live_bytes = 0
 
+    def _sink_window_slice(self, sequence: list[int]) -> list[int]:
+        """Mirror the child verifier's bounded token-id cache layout."""
+        sink_size = int(self.client.config.sink_size)
+        window_size = int(self.client.config.window_size)
+        budget = sink_size + window_size
+        if len(sequence) <= budget:
+            return list(sequence)
+        tail = list(sequence[-window_size:]) if window_size else []
+        return list(sequence[:sink_size]) + tail
+
     def _apply(self, state: dict[str, Any]) -> dict[str, Any]:
         self.cached_token_sequence = [
             int(token) for token in state.get("cached_token_ids", ())
@@ -853,85 +876,97 @@ def prefill(
         cancel_event: threading.Event | None = None,
     ) -> None:
         tokens = [int(token) for token in prompt_ids]
-        state = self.client._session_request(
-            self.session_id,
-            "Init",
-            {"token_ids": tokens},
-            cancel_event=cancel_event,
-        )
-        checkpoint = self.client._checkpoints[self.session_id]
-        checkpoint.snapshot = None
-        checkpoint.compatibility = None
-        checkpoint.replay_token_ids = list(tokens)
-        checkpoint.initialized = True
-        self.client._mark_current(self.session_id)
-        self._apply(state)
+        with self.client._lock:
+            checkpoint = self.client._checkpoint_locked(self.session_id)
+            state = self.client._session_request(
+                self.session_id,
+                "Init",
+                {"token_ids": tokens},
+                cancel_event=cancel_event,
+            )
+            checkpoint.snapshot = None
+            checkpoint.compatibility = None
+            checkpoint.replay_token_ids = list(tokens)
+            checkpoint.initialized = True
+            self.client._mark_current(self.session_id)
+            self._apply(state)
 
     def append_accepted_tokens(
         self,
         tokens: list[int],
         cancel_event: threading.Event | None = None,
     ) -> None:
         committed = [int(token) for token in tokens]
-        state = self.client._session_request(
-            self.session_id,
-            "Append",
-            {"token_ids": committed},
-            cancel_event=cancel_event,
-        )
-        self.client._checkpoints[self.session_id].replay_token_ids.extend(committed)
-        self._apply(state)
+        with self.client._lock:
+            checkpoint = self.client._checkpoint_locked(self.session_id)
+            state = self.client._session_request(
+                self.session_id,
+                "Append",
+                {"token_ids": committed},
+                cancel_event=cancel_event,
+            )
+            checkpoint.replay_token_ids.extend(committed)
+            self._apply(state)
 
     def generate_step(
         self,
         cancel_event: threading.Event | None = None,
     ) -> int:
-        state = self.client._session_request(
-            self.session_id,
-            "GenerateStep",
-            {},
-            cancel_event=cancel_event,
-        )
-        token_id = int(state["token_id"])
-        self.client._checkpoints[self.session_id].replay_token_ids.append(token_id)
-        self._apply(state)
-        return token_id
+        with self.client._lock:
+            checkpoint = self.client._checkpoint_locked(self.session_id)
+            state = self.client._session_request(
+                self.session_id,
+                "GenerateStep",
+                {},
+                cancel_event=cancel_event,
+            )
+            token_id = int(state["token_id"])
+            checkpoint.replay_token_ids.append(token_id)
+            self._apply(state)
+            return token_id
 
     def import_snapshot(
         self,
         payload: bytes,
         compatibility: Any,
     ) -> dict[str, Any]:
         compat = asdict(compatibility)
-        # Ensure a worker-side session exists before importing its cache.
-        self.client._session_request(
-            self.session_id, "Init", {"token_ids": []}
-        )
-        state = self.client._session_request(
-            self.session_id,
-            "ImportSnapshot",
-            {"compatibility": compat},
-            bytes(payload),
-        )
-        checkpoint = self.client._checkpoints[self.session_id]
-        checkpoint.snapshot = bytes(payload)
-        checkpoint.compatibility = compat
-        checkpoint.replay_token_ids = []
-        checkpoint.initialized = True
-        self.client._mark_current(self.session_id)
-        return self._apply(state)
+        snapshot = bytes(payload)
+        # Pin the router checkpoint and proxy lifecycle across the full
+        # Init -> ImportSnapshot -> checkpoint-publication transaction.
+        # Close/cancel cleanup uses the same RLock and therefore cannot remove
+        # restart state after the child accepted the import but before it is
+        # made durable on the router.
+        with self.client._lock:
+            checkpoint = self.client._checkpoint_locked(self.session_id)
+            self.client._session_request(
+                self.session_id, "Init", {"token_ids": []}
+            )
+            state = self.client._session_request(
+                self.session_id,
+                "ImportSnapshot",
+                {"compatibility": compat},
+                snapshot,
+            )
+            checkpoint.snapshot = snapshot
+            checkpoint.compatibility = compat
+            checkpoint.replay_token_ids = []
+            checkpoint.initialized = True
+            self.client._mark_current(self.session_id)
+            return self._apply(state)
 
     def reset(self) -> None:
-        state = self.client._session_request(
-            self.session_id, "Init", {"token_ids": []}
-        )
-        checkpoint = self.client._checkpoints[self.session_id]
-        checkpoint.snapshot = None
-        checkpoint.compatibility = None
-        checkpoint.replay_token_ids = []
-        checkpoint.initialized = True
-        self.client._mark_current(self.session_id)
-        self._apply(state)
+        with self.client._lock:
+            checkpoint = self.client._checkpoint_locked(self.session_id)
+            state = self.client._session_request(
+                self.session_id, "Init", {"token_ids": []}
+            )
+            checkpoint.snapshot = None
+            checkpoint.compatibility = None
+            checkpoint.replay_token_ids = []
+            checkpoint.initialized = True
+            self.client._mark_current(self.session_id)
+            self._apply(state)
 
     def k_seq_length(self, _session: Any) -> int:
         return len(self.cached_token_sequence)

inference_engine/bench/prefill_fleet_report.py

@@ -10,6 +10,14 @@
     "allens_cold_restore",
     "agent_generator",
     "agent_critic",
+    "agent_premise_auditor",
+    "agent_definition_auditor",
+    "agent_counterexample_worker",
+    "agent_decomposer",
+    "agent_formalizer",
+    "agent_prover",
+    "agent_adversarial_proponent",
+    "agent_judge",
 )
 HIT_SOURCES = ("remote_worker", "primary_hot", "allens_offload", "unknown")
 _PRIVATE_KEYS = {

inference_engine/distributed/prefill_worker.py

@@ -147,6 +147,7 @@ def __init__(
         self._requests: dict[tuple[str, str], str] = {}
         self._lock = threading.RLock()
         self._thread_local = threading.local()
+        self._measured_tokens_per_second = 0.0
         self._executor = ThreadPoolExecutor(
             max_workers=self.max_concurrent_jobs,
             thread_name_prefix="kakeya-prefill-worker",
@@ -275,6 +276,11 @@ def stats(self) -> tuple[int, int, float, int]:
             load = min(1.0, (running + queued) / self.max_concurrent_jobs)
             return running, queued, load, queued_tokens
 
+    def measured_tokens_per_second(self, fallback: float = 0.0) -> float:
+        with self._lock:
+            measured = self._measured_tokens_per_second
+        return measured if measured > 0 else max(0.0, float(fallback))
+
     def close(self) -> None:
         self._executor.shutdown(wait=False, cancel_futures=True)
 
@@ -364,8 +370,22 @@ def _run(self, job_id: str) -> None:
             if timer is not None:
                 timer.cancel()
             with self._lock:
-                job.compute_ms = (time.perf_counter() - started) * 1000.0
+                elapsed_s = time.perf_counter() - started
+                job.compute_ms = elapsed_s * 1000.0
                 job.finished_at = time.time()
+                if (
+                    job.state == PrefillJobState.COMPLETED
+                    and elapsed_s > 0
+                    and job.token_ids
+                ):
+                    sample = len(job.token_ids) / elapsed_s
+                    if self._measured_tokens_per_second <= 0:
+                        self._measured_tokens_per_second = sample
+                    else:
+                        self._measured_tokens_per_second = (
+                            0.25 * sample
+                            + 0.75 * self._measured_tokens_per_second
+                        )
 
     def _update_job_progress(self, job_id: str, token_count: int) -> None:
         with self._lock:
@@ -460,8 +480,14 @@ async def SubmitPrefillJob(self, request, context):  # noqa: N802
             status=int(job.state),
             worker_node_id=self.node_id,
             queue_eta_ms=(
-                self.jobs.stats()[3] / self.tokens_per_second_prefill * 1000.0
-                if self.tokens_per_second_prefill > 0 else 0.0
+                self.jobs.stats()[3]
+                / self.jobs.measured_tokens_per_second(
+                    self.tokens_per_second_prefill,
+                )
+                * 1000.0
+                if self.jobs.measured_tokens_per_second(
+                    self.tokens_per_second_prefill,
+                ) > 0 else 0.0
             ),
         )
 

scripts/start_prefill_worker_node.py

@@ -171,14 +171,17 @@ def refresh_cache_budget() -> tuple[int, int]:
     def card() -> NodeCapability:
         active_model_bytes, _ = refresh_cache_budget()
         inflight, queued, load, queued_tokens = jobs.stats()
+        measured_prefill_tps = jobs.measured_tokens_per_second(
+            args.prefill_tps,
+        )
         worker = PrefillWorkerCapability(
             compatibility=compatibility,
             worker_address=args.advertise,
             max_concurrent_jobs=args.max_concurrent_jobs,
             inflight_jobs=inflight,
             queued_jobs=queued,
             load=load,
-            tokens_per_second_prefill=args.prefill_tps,
+            tokens_per_second_prefill=measured_prefill_tps,
             ram_bytes_free=max(
                 0,
                 physical_memory_bytes()
@@ -197,7 +200,7 @@ def card() -> NodeCapability:
                     args.cache_model_id or args.model_id,
                     CapabilityRole.PREFILL_COMPUTE,
                     args.quantization,
-                    args.prefill_tps,
+                    measured_prefill_tps,
                 ),
             ),
             announced_at_unix=time.time(),