diff --git a/Dockerfile b/Dockerfile index 6862964..d8cccca 100644 --- a/Dockerfile +++ b/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/, diff --git a/inference_engine/backends/mlx/decode_worker.py b/inference_engine/backends/mlx/decode_worker.py index 741974b..2eb0d8d 100644 --- a/inference_engine/backends/mlx/decode_worker.py +++ b/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,19 +876,20 @@ 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, @@ -873,29 +897,33 @@ def append_accepted_tokens( 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, @@ -903,35 +931,42 @@ def import_snapshot( 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) diff --git a/inference_engine/bench/prefill_fleet_report.py b/inference_engine/bench/prefill_fleet_report.py index 75a91f3..304bf12 100644 --- a/inference_engine/bench/prefill_fleet_report.py +++ b/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 = { diff --git a/inference_engine/distributed/prefill_worker.py b/inference_engine/distributed/prefill_worker.py index 0a2a0e5..e7a4d53 100644 --- a/inference_engine/distributed/prefill_worker.py +++ b/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 ), ) diff --git a/scripts/start_prefill_worker_node.py b/scripts/start_prefill_worker_node.py index 64da9cd..3e84247 100644 --- a/scripts/start_prefill_worker_node.py +++ b/scripts/start_prefill_worker_node.py @@ -171,6 +171,9 @@ 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, @@ -178,7 +181,7 @@ def card() -> NodeCapability: 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(), diff --git a/tests/backends/mlx/test_decode_worker.py b/tests/backends/mlx/test_decode_worker.py index 569053b..b1d8e4f 100644 --- a/tests/backends/mlx/test_decode_worker.py +++ b/tests/backends/mlx/test_decode_worker.py @@ -17,6 +17,7 @@ PROTOCOL_VERSION, DecodeWorkerClient, DecodeWorkerConfig, + DecodeWorkerSessionClosed, _recv_frame, _send_frame, ) @@ -73,11 +74,18 @@ def build_fake_verifier(config): return FakeDecodeVerifier(str(config["model_id"])) -def _client(tmp_path, *, model_id="fake", timeout=2.0): +def _client( + tmp_path, + *, + model_id="fake", + timeout=2.0, + sink_size=2, + window_size=6, +): return DecodeWorkerClient(DecodeWorkerConfig( model_id=model_id, - sink_size=2, - window_size=6, + sink_size=sink_size, + window_size=window_size, request_timeout_s=timeout, startup_timeout_s=10.0, socket_path=str(tmp_path / "decode.sock"), @@ -134,6 +142,31 @@ def test_protocol_health_and_all_operations(tmp_path): client.close() +def test_worker_proxy_sink_window_slice_matches_child_layout(tmp_path): + client = _client(tmp_path) + try: + session = client.get("slice") + short = [1, 2, 3] + assert session._sink_window_slice(short) == short + assert session._sink_window_slice(short) is not short + + sequence = list(range(20)) + assert session._sink_window_slice(sequence) == ( + sequence[:2] + sequence[-6:] + ) + finally: + client.close() + + +def test_worker_proxy_sink_window_slice_supports_zero_window(tmp_path): + client = _client(tmp_path, window_size=0) + try: + session = client.get("sink-only") + assert session._sink_window_slice(list(range(10))) == [0, 1] + finally: + client.close() + + def test_worker_matches_in_process_greedy_parity(tmp_path): expected = FakeDecodeVerifier() expected.prefill([7, 8, 9]) @@ -207,6 +240,81 @@ def test_snapshot_plus_post_checkpoint_replay_survives_hard_kill(tmp_path): client.close() +def test_close_during_snapshot_import_waits_for_checkpoint_publication(tmp_path): + client = _client(tmp_path) + imported = threading.Event() + release_import = threading.Event() + close_done = threading.Event() + failures = [] + try: + session = client.get("import-close-race") + original_request = client._session_request + + def paused_request(session_id, operation, arguments, payload=b"", **kwargs): + state = original_request( + session_id, operation, arguments, payload, **kwargs, + ) + if operation == "ImportSnapshot": + imported.set() + assert release_import.wait(timeout=2.0) + return state + + client._session_request = paused_request + + def run_import(): + try: + session.import_snapshot( + json.dumps({"tokens": [10, 11, 12]}).encode(), + CacheCompatibility(model_id="fake"), + ) + except BaseException as exc: + failures.append(exc) + + def run_close(): + try: + client.close_session(session.session_id) + except BaseException as exc: + failures.append(exc) + finally: + close_done.set() + + import_thread = threading.Thread(target=run_import) + close_thread = threading.Thread(target=run_close) + import_thread.start() + assert imported.wait(timeout=2.0) + close_thread.start() + + # The child has acknowledged ImportSnapshot, but router checkpoint + # publication is deliberately paused. Cleanup must remain blocked. + assert not close_done.wait(timeout=0.1) + assert session.session_id in client._checkpoints + + release_import.set() + import_thread.join(timeout=2.0) + close_thread.join(timeout=2.0) + assert not import_thread.is_alive() + assert not close_thread.is_alive() + assert failures == [] + assert session.session_id not in client._checkpoints + assert session.session_id not in client._proxies + assert client.health()["session_count"] == 0 + finally: + release_import.set() + client.close() + + +def test_closed_proxy_fails_with_explicit_session_error(tmp_path): + client = _client(tmp_path) + try: + session = client.get("closed") + session.prefill([1, 2]) + client.close_session(session.session_id) + with pytest.raises(DecodeWorkerSessionClosed, match="is closed"): + session.reset() + finally: + client.close() + + def test_recycle_lazily_restores_other_sessions(tmp_path): client = _client(tmp_path) try: diff --git a/tests/inference_engine/bench/test_prefill_fleet_report.py b/tests/inference_engine/bench/test_prefill_fleet_report.py index 92b3b31..45246ce 100644 --- a/tests/inference_engine/bench/test_prefill_fleet_report.py +++ b/tests/inference_engine/bench/test_prefill_fleet_report.py @@ -63,6 +63,19 @@ def test_schema_rejects_unknown_and_private_fields(): assert normalize_stage(_stage("agent_generator", "primary_hot"))["name"] == ( "agent_generator" ) + for role in ( + "premise_auditor", + "definition_auditor", + "counterexample_worker", + "decomposer", + "formalizer", + "prover", + "adversarial_proponent", + "judge", + ): + assert normalize_stage(_stage(f"agent_{role}"))["name"] == ( + f"agent_{role}" + ) with pytest.raises(ValueError, match="unknown hit_source"): normalize_stage(_stage(source="peer:1")) with pytest.raises(ValueError, match="non-negative"): diff --git a/tests/inference_engine/distributed/test_prefill_worker.py b/tests/inference_engine/distributed/test_prefill_worker.py index 6e1cdf7..d00ee07 100644 --- a/tests/inference_engine/distributed/test_prefill_worker.py +++ b/tests/inference_engine/distributed/test_prefill_worker.py @@ -186,7 +186,8 @@ async def _rpc(address, method, request): @pytest.mark.asyncio async def test_submit_is_idempotent_and_completes(worker): - address, engine, cache, _ = worker + address, engine, cache, jobs = worker + assert jobs.measured_tokens_per_second(100.0) == 100.0 first = await _rpc(address, "SubmitPrefillJob", _submit()) second = await _rpc(address, "SubmitPrefillJob", _submit()) assert first.job_id == second.job_id @@ -204,6 +205,23 @@ async def test_submit_is_idempotent_and_completes(worker): assert status.lease_id and status.cache_address == "cache:1" assert engine.calls == 1 assert len(cache.block_hashes()) == 2 + assert jobs.measured_tokens_per_second(100.0) > 0 + + third = await _rpc(address, "SubmitPrefillJob", _submit("r2")) + for _ in range(100): + status = await _rpc( + address, + "GetPrefillJobStatus", + distributed_pb2.GetPrefillJobStatusRequest( + job_id=third.job_id, + tenant_id="tenant", + ), + ) + if status.status == int(PrefillJobState.COMPLETED): + break + await asyncio.sleep(0.01) + assert status.status == int(PrefillJobState.COMPLETED) + assert engine.calls == 2 @pytest.mark.asyncio