From fb5d7856c985fb90e1a0f8a429e14d86fef49310 Mon Sep 17 00:00:00 2001 From: fluffy314 Date: Fri, 17 Jul 2026 16:23:22 +0800 Subject: [PATCH 1/2] perf(prefill): segment long jobs for autoresearch Split full-context model compute into observable sub-five-minute segments and add a Karpathy-style fixed harness for iterating on Prefill strategy without weakening semantic gates. Co-authored-by: Cursor --- autoresearch/prefill/candidate.py | 7 ++ autoresearch/prefill/prepare.py | 108 ++++++++++++++++++ autoresearch/prefill/program.md | 41 +++++++ deploy/install_prefill_worker_launchd.sh | 2 + .../ai.kakeya.prefill-worker-peer.plist | 1 + docs/ops/distributed-prefill-kv-network.md | 11 ++ .../backends/mlx/prefill_worker.py | 36 +++++- .../distributed/prefill_cache_runtime.py | 8 ++ .../distributed/prefill_worker.py | 21 +++- scripts/agent_gan_repl.py | 29 ++++- scripts/start_prefill_worker_node.py | 11 +- .../bench/test_prefill_autoresearch.py | 50 ++++++++ .../bridge/test_agent_gan_repl.py | 13 ++- .../bridge/test_prefill_worker_launchd.py | 6 + .../distributed/test_prefill_worker.py | 52 ++++++++- 15 files changed, 382 insertions(+), 14 deletions(-) create mode 100644 autoresearch/prefill/candidate.py create mode 100644 autoresearch/prefill/prepare.py create mode 100644 autoresearch/prefill/program.md create mode 100644 tests/inference_engine/bench/test_prefill_autoresearch.py diff --git a/autoresearch/prefill/candidate.py b/autoresearch/prefill/candidate.py new file mode 100644 index 00000000..572a6945 --- /dev/null +++ b/autoresearch/prefill/candidate.py @@ -0,0 +1,7 @@ +"""The only Prefill strategy file the autoresearch agent may edit.""" + +PREFILL_COMPUTE_CHUNK_TOKENS = 256 +SNAPSHOT_MODE = "final_only" +MAX_SEGMENT_SECONDS = 300.0 +REQUIRE_FULL_CONTEXT = True +ALLOW_FALLBACK = False diff --git a/autoresearch/prefill/prepare.py b/autoresearch/prefill/prepare.py new file mode 100644 index 00000000..b273d60e --- /dev/null +++ b/autoresearch/prefill/prepare.py @@ -0,0 +1,108 @@ +#!/usr/bin/env python3 +"""Fixed evaluation harness for Karpathy-style Prefill autoresearch.""" +from __future__ import annotations + +import argparse +import csv +import importlib.util +import json +import time +from pathlib import Path + + +def _load_candidate(path: Path): + spec = importlib.util.spec_from_file_location("prefill_candidate", path) + if spec is None or spec.loader is None: + raise RuntimeError("cannot load candidate") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def evaluate(report: dict, candidate) -> dict: + stages = report.get("stages", []) + critic = next( + (stage for stage in stages if stage.get("name") == "agent_critic"), + None, + ) + if critic is None: + raise ValueError("report has no Critic stage") + prefix_tokens = int(critic.get("prefix_tokens", 0)) + warmup_s = float(critic.get("warmup_wall_s", 0)) + measured_tps = prefix_tokens / warmup_s if warmup_s > 0 else 0.0 + estimated_max_segment_s = ( + candidate.PREFILL_COMPUTE_CHUNK_TOKENS / measured_tps + if measured_tps > 0 else float("inf") + ) + delta = critic.get("delta", {}) + constraints = { + "stage_ok": bool(critic.get("ok")), + "complete": bool(critic.get("complete")), + "full_context": ( + critic.get("review_scope") == "full" + and int(critic.get("critic_omitted_tokens", -1)) == 0 + and int(critic.get("critic_context_tokens", -1)) + == int(critic.get("generator_full_tokens", -2)) + ), + "recursive_protocol": ( + critic.get("critic_protocol") + == "recursive_proof_decomposition_v2" + ), + "no_fallback": int(delta.get("fallbacks", 0)) == 0, + "no_job_failure": int(delta.get("remote_job_failures", 0)) == 0, + "segment_under_budget": ( + estimated_max_segment_s <= candidate.MAX_SEGMENT_SECONDS + ), + "final_only_snapshot": candidate.SNAPSHOT_MODE == "final_only", + } + return { + "accepted": all(constraints.values()), + "metric_cold_critic_prefill_s": warmup_s, + "measured_prefill_tps": measured_tps, + "estimated_max_segment_s": estimated_max_segment_s, + "compute_chunk_tokens": candidate.PREFILL_COMPUTE_CHUNK_TOKENS, + "constraints": constraints, + } + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--report", type=Path, required=True) + parser.add_argument( + "--candidate", + type=Path, + default=Path(__file__).with_name("candidate.py"), + ) + parser.add_argument( + "--results", + type=Path, + default=Path(__file__).with_name("results.tsv"), + ) + args = parser.parse_args() + candidate = _load_candidate(args.candidate) + result = evaluate(json.loads(args.report.read_text()), candidate) + write_header = not args.results.exists() + with args.results.open("a", newline="") as handle: + writer = csv.DictWriter( + handle, + fieldnames=( + "timestamp", + "accepted", + "metric_cold_critic_prefill_s", + "measured_prefill_tps", + "estimated_max_segment_s", + "compute_chunk_tokens", + ), + delimiter="\t", + ) + if write_header: + writer.writeheader() + writer.writerow({"timestamp": time.time(), **{ + key: result[key] for key in writer.fieldnames if key != "timestamp" + }}) + print(json.dumps(result, indent=2, sort_keys=True)) + return 0 if result["accepted"] else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/autoresearch/prefill/program.md b/autoresearch/prefill/program.md new file mode 100644 index 00000000..82e10bd6 --- /dev/null +++ b/autoresearch/prefill/program.md @@ -0,0 +1,41 @@ +# Prefill AutoResearch Program + +You are optimizing the two-Mac full-context Prefill system. + +## Ownership + +- You may edit only `candidate.py`. +- Never edit `prepare.py`, benchmark reports, tests, or production metrics. +- The human owns this file. + +## Objective + +Minimize `metric_cold_critic_prefill_s`. Lower is better. + +## Hard constraints + +- Every compute segment must remain at or below 300 seconds. +- Critic must receive the complete Generator response. +- `critic_omitted_tokens` must equal zero. +- Protocol must be `recursive_proof_decomposition_v2`. +- Snapshot mode must remain `final_only`. +- No fallback, local Primary Prefill, failed remote job, sampling, summary, or + semantic simplification is allowed. +- Primary remains decode-only and allens remains Prefill-only. + +## Experiment loop + +1. Read `candidate.py` and `results.tsv`. +2. State one concrete performance hypothesis. +3. Modify only `candidate.py`. +4. Deploy the candidate to allens. +5. Clear Primary and allens caches. +6. Run the fixed full-context acceptance workload. +7. Run `prepare.py` against the resulting report. +8. Keep the candidate only if every hard constraint passes and cold Critic + Prefill time improves. Otherwise restore the previous candidate. +9. Append the result and repeat. + +Do not optimize output wording, scores, prizes, or other proof-irrelevant +content. Optimize only measured Prefill execution while preserving the complete +semantic contract. diff --git a/deploy/install_prefill_worker_launchd.sh b/deploy/install_prefill_worker_launchd.sh index 3a3a6810..92019e46 100755 --- a/deploy/install_prefill_worker_launchd.sh +++ b/deploy/install_prefill_worker_launchd.sh @@ -15,6 +15,7 @@ CACHE_GB="${KAKEYA_WORKER_CACHE_GB:-4}" CACHE_MIN_GB="${KAKEYA_WORKER_CACHE_MIN_GB:-1}" MEMORY_RESERVE_GB="${KAKEYA_WORKER_MEMORY_RESERVE_GB:-0.5}" SNAPSHOT_BYTES_PER_TOKEN="${KAKEYA_SNAPSHOT_BYTES_PER_TOKEN:-400000}" +COMPUTE_CHUNK_TOKENS="${KAKEYA_PREFILL_COMPUTE_CHUNK_TOKENS:-256}" ADAPTIVE_CACHE="${KAKEYA_WORKER_ADAPTIVE_CACHE:-0}" PSK_FILE="${KAKEYA_FLEET_PSK_FILE:-}" CACHE_MODEL_ID="${KAKEYA_CACHE_MODEL_ID:-$KAKEYA_WORKER_MODEL}" @@ -78,6 +79,7 @@ cat > "$PLIST" <--cache-min-gb$CACHE_MIN_GB --memory-reserve-gb$MEMORY_RESERVE_GB --estimated-snapshot-bytes-per-token$SNAPSHOT_BYTES_PER_TOKEN + --prefill-compute-chunk-tokens$COMPUTE_CHUNK_TOKENS $adaptive_xml --sink$SINK --window$WINDOW diff --git a/deploy/launchd/ai.kakeya.prefill-worker-peer.plist b/deploy/launchd/ai.kakeya.prefill-worker-peer.plist index a92dd5f6..f04db179 100644 --- a/deploy/launchd/ai.kakeya.prefill-worker-peer.plist +++ b/deploy/launchd/ai.kakeya.prefill-worker-peer.plist @@ -27,6 +27,7 @@ --adaptive-cache --memory-reserve-gb0.5 --estimated-snapshot-bytes-per-token400000 + --prefill-compute-chunk-tokens256 --prefill-tps1 --max-concurrent-jobs1 --networkthunderbolt diff --git a/docs/ops/distributed-prefill-kv-network.md b/docs/ops/distributed-prefill-kv-network.md index a351d13f..61b4287c 100644 --- a/docs/ops/distributed-prefill-kv-network.md +++ b/docs/ops/distributed-prefill-kv-network.md @@ -365,6 +365,17 @@ The MLX worker exports and compresses only the final chained-prefix snapshot. The final hash commits every preceding token block, so exporting a growing full snapshot at every 64-token boundary is redundant and creates quadratic serialization/compression work for long Critic contexts. +Model compute is segmented independently from the 64-token content-addressed +hash boundary. The default segment is 256 tokens, keeping each allens step below +the five-minute research budget at the measured Prefill rate. Worker job status +updates `tokens_computed` after every segment; Terminal heartbeats display +tokens, percentage, and ETA. + +Karpathy-style optimization lives in `autoresearch/prefill/`. Humans edit +`program.md`, the research agent edits only `candidate.py`, and immutable +`prepare.py` evaluates full-context correctness plus cold Critic Prefill time. +Candidates are retained only when every semantic/topology constraint passes and +the metric improves; `results.tsv` records experiments. Interactive prompt templates are deterministic and contain no per-run nonce, so repeating the same task can reuse allens cold-tier and Primary hot-tier KV. diff --git a/inference_engine/backends/mlx/prefill_worker.py b/inference_engine/backends/mlx/prefill_worker.py index 594237c9..8857609b 100644 --- a/inference_engine/backends/mlx/prefill_worker.py +++ b/inference_engine/backends/mlx/prefill_worker.py @@ -2,7 +2,7 @@ from __future__ import annotations import threading -from typing import Sequence +from typing import Callable, Sequence from inference_engine.backends.mlx.prefill_snapshot import ( export_mlx_prefill_snapshot, @@ -18,11 +18,31 @@ class MLXPrefillComputeEngine: """Serially runs prefill with a loaded MLX verifier and exports one snapshot.""" - def __init__(self, verifier, compatibility: CacheCompatibility) -> None: + def __init__( + self, + verifier, + compatibility: CacheCompatibility, + *, + compute_chunk_tokens: int = 256, + ) -> None: + if compute_chunk_tokens <= 0: + raise ValueError("compute_chunk_tokens must be > 0") self.verifier = verifier self.compatibility = compatibility + self.compute_chunk_tokens = int(compute_chunk_tokens) + self._progress_callback: Callable[[int], None] | None = None self._lock = threading.Lock() + def set_progress_callback( + self, + callback: Callable[[int], None] | None, + ) -> None: + self._progress_callback = callback + + def _report_progress(self, token_count: int) -> None: + if self._progress_callback is not None: + self._progress_callback(int(token_count)) + def compute_prefill( self, token_ids: Sequence[int], @@ -43,18 +63,24 @@ def compute_prefill( with self._lock: if cancelled.is_set(): raise InterruptedError("prefill job cancelled") - first_end = min(size, len(tokens)) + first_end = min(self.compute_chunk_tokens, len(tokens)) self.verifier.prefill(tokens[:first_end]) - for start in range(first_end, len(tokens), size): + self._report_progress(first_end) + for start in range( + first_end, + len(tokens), + self.compute_chunk_tokens, + ): if cancelled.is_set(): raise InterruptedError("prefill job cancelled") - block = tokens[start:start + size] + block = tokens[start:start + self.compute_chunk_tokens] logits = self.verifier.forward_block(block) self.verifier.commit_or_truncate( forwarded=len(block), accepted=len(block), ) self.verifier.next_token_logits = logits[-1].clone() + self._report_progress(min(start + len(block), len(tokens))) # Export exactly once. Intermediate full snapshots make encoding # and compression quadratic in prompt length and are not required # for correctness because the final chained hash commits every diff --git a/inference_engine/distributed/prefill_cache_runtime.py b/inference_engine/distributed/prefill_cache_runtime.py index 94b1c03b..bb492d59 100644 --- a/inference_engine/distributed/prefill_cache_runtime.py +++ b/inference_engine/distributed/prefill_cache_runtime.py @@ -72,6 +72,8 @@ class PrefillReuseStats: bytes_received: int = 0 remote_jobs: int = 0 remote_job_failures: int = 0 + remote_job_tokens_total: int = 0 + remote_job_tokens_computed: int = 0 fallbacks: int = 0 last_fallback_reason: str = "" publish_attempts: int = 0 @@ -356,6 +358,8 @@ def _compute_remote( auth=self.auth, ) self.stats.remote_jobs += 1 + self.stats.remote_job_tokens_total = len(tokens) + self.stats.remote_job_tokens_computed = 0 deadline = time.monotonic() + self.worker_timeout_s while time.monotonic() < deadline: status_request = distributed_pb2.GetPrefillJobStatusRequest( @@ -368,6 +372,10 @@ def _compute_remote( timeout_s=self.lookup_timeout_s, auth=self.auth, ) + self.stats.remote_job_tokens_computed = min( + len(tokens), + int(status.tokens_computed), + ) if status.status == int(PrefillJobState.COMPLETED): return _Hit( source=status.cache_address or target.address, diff --git a/inference_engine/distributed/prefill_worker.py b/inference_engine/distributed/prefill_worker.py index dec13fca..f13ef01f 100644 --- a/inference_engine/distributed/prefill_worker.py +++ b/inference_engine/distributed/prefill_worker.py @@ -258,6 +258,7 @@ def _run(self, job_id: str) -> None: job.state = PrefillJobState.RUNNING started = time.perf_counter() timer = None + engine = None if job.deadline_at: remaining = job.deadline_at - time.time() if remaining <= 0: @@ -274,7 +275,13 @@ def _run(self, job_id: str) -> None: job.job_id, len(job.token_ids) * self.estimated_snapshot_bytes_per_token, ) - blocks = tuple(self._engine_for_current_thread().compute_prefill( + engine = self._engine_for_current_thread() + set_progress = getattr(engine, "set_progress_callback", None) + if callable(set_progress): + set_progress( + lambda count: self._update_job_progress(job.job_id, count), + ) + blocks = tuple(engine.compute_prefill( job.token_ids, job.block_hashes, compression=job.compression, @@ -313,6 +320,10 @@ def _run(self, job_id: str) -> None: job.state = PrefillJobState.FAILED job.failure_reason = f"{type(exc).__name__}: {exc}" finally: + if engine is not None: + set_progress = getattr(engine, "set_progress_callback", None) + if callable(set_progress): + set_progress(None) self.cache_store.release_reservation(job.job_id) if timer is not None: timer.cancel() @@ -320,6 +331,14 @@ def _run(self, job_id: str) -> None: job.compute_ms = (time.perf_counter() - started) * 1000.0 job.finished_at = time.time() + def _update_job_progress(self, job_id: str, token_count: int) -> None: + with self._lock: + job = self._jobs[job_id] + job.tokens_computed = min( + len(job.token_ids), + max(job.tokens_computed, int(token_count)), + ) + def _engine_for_current_thread(self) -> PrefillComputeEngine: if self.engine is not None: return self.engine diff --git a/scripts/agent_gan_repl.py b/scripts/agent_gan_repl.py index 254922d1..cb03f657 100644 --- a/scripts/agent_gan_repl.py +++ b/scripts/agent_gan_repl.py @@ -122,9 +122,15 @@ def finish(self) -> None: class PrefillHeartbeat: - def __init__(self, label: str, interval_s: float = 30.0) -> None: + def __init__( + self, + label: str, + interval_s: float = 30.0, + stats_provider=None, + ) -> None: self.label = label self.interval_s = interval_s + self.stats_provider = stats_provider self.stop = threading.Event() self.started = 0.0 self.thread = None @@ -142,8 +148,23 @@ def __exit__(self, *_args): def _run(self): while not self.stop.wait(self.interval_s): elapsed = time.perf_counter() - self.started + progress = "" + if self.stats_provider is not None: + stats = self.stats_provider() + total = int(stats.get("remote_job_tokens_total", 0)) + computed = int(stats.get("remote_job_tokens_computed", 0)) + if total > 0: + percent = min(100.0, computed / total * 100.0) + eta = ( + elapsed * (total - computed) / computed + if computed > 0 else 0.0 + ) + progress = ( + f" · {computed}/{total} tokens ({percent:.1f}%)" + + (f" · ETA {eta:.0f}s" if computed > 0 else "") + ) print( - f"[allens] {self.label} Prefill still running: {elapsed:.0f}s", + f"[allens] {self.label} Prefill: {elapsed:.0f}s{progress}", flush=True, ) @@ -285,7 +306,7 @@ def get_stats(): f"[allens] Generator Prefill: {len(generator_ids)} tokens...", flush=True, ) - with PrefillHeartbeat("Generator"): + with PrefillHeartbeat("Generator", stats_provider=get_stats): _, generator_warm = _infer( client, eos_ids, generator_ids, 1, get_stats, ) @@ -351,7 +372,7 @@ def get_stats(): f"[allens] Critic Prefill: {len(critic_ids)} tokens...", flush=True, ) - with PrefillHeartbeat("Critic"): + with PrefillHeartbeat("Critic", stats_provider=get_stats): _, critic_warm = _infer( client, eos_ids, critic_ids, 1, get_stats, ) diff --git a/scripts/start_prefill_worker_node.py b/scripts/start_prefill_worker_node.py index 5924c301..49653d0f 100644 --- a/scripts/start_prefill_worker_node.py +++ b/scripts/start_prefill_worker_node.py @@ -126,7 +126,11 @@ def engine_factory() -> MLXPrefillComputeEngine: dtype=torch.bfloat16, device="cpu", )) - return MLXPrefillComputeEngine(verifier, compatibility) + return MLXPrefillComputeEngine( + verifier, + compatibility, + compute_chunk_tokens=args.prefill_compute_chunk_tokens, + ) jobs = PrefillJobStore( None, @@ -298,6 +302,11 @@ def main() -> None: type=int, default=400_000, ) + parser.add_argument( + "--prefill-compute-chunk-tokens", + type=int, + default=256, + ) parser.add_argument("--job-ttl-s", type=float, default=600.0) parser.add_argument("--prefill-tps", type=float, default=20.0) parser.add_argument("--max-concurrent-rpcs", type=int, default=32) diff --git a/tests/inference_engine/bench/test_prefill_autoresearch.py b/tests/inference_engine/bench/test_prefill_autoresearch.py new file mode 100644 index 00000000..3bece3d0 --- /dev/null +++ b/tests/inference_engine/bench/test_prefill_autoresearch.py @@ -0,0 +1,50 @@ +from autoresearch.prefill.prepare import evaluate + + +class Candidate: + PREFILL_COMPUTE_CHUNK_TOKENS = 256 + SNAPSHOT_MODE = "final_only" + MAX_SEGMENT_SECONDS = 300 + + +def _report(**overrides): + stage = { + "name": "agent_critic", + "ok": True, + "complete": True, + "prefix_tokens": 1000, + "warmup_wall_s": 1000, + "review_scope": "full", + "generator_full_tokens": 900, + "critic_context_tokens": 900, + "critic_omitted_tokens": 0, + "critic_protocol": "recursive_proof_decomposition_v2", + "delta": {"fallbacks": 0, "remote_job_failures": 0}, + } + stage.update(overrides) + return {"stages": [stage]} + + +def test_autoresearch_accepts_faster_full_context_candidate(): + result = evaluate(_report(), Candidate) + assert result["accepted"] + assert result["metric_cold_critic_prefill_s"] == 1000 + assert result["estimated_max_segment_s"] == 256 + assert all(result["constraints"].values()) + + +def test_autoresearch_rejects_slow_segment_or_semantic_regression(): + slow = type("Slow", (), { + "PREFILL_COMPUTE_CHUNK_TOKENS": 512, + "SNAPSHOT_MODE": "final_only", + "MAX_SEGMENT_SECONDS": 300, + }) + assert not evaluate(_report(), slow)["accepted"] + assert not evaluate( + _report(critic_omitted_tokens=1), + Candidate, + )["accepted"] + assert not evaluate( + _report(delta={"fallbacks": 1, "remote_job_failures": 1}), + Candidate, + )["accepted"] diff --git a/tests/inference_engine/bridge/test_agent_gan_repl.py b/tests/inference_engine/bridge/test_agent_gan_repl.py index 81f6a535..896429c8 100644 --- a/tests/inference_engine/bridge/test_agent_gan_repl.py +++ b/tests/inference_engine/bridge/test_agent_gan_repl.py @@ -95,10 +95,19 @@ def test_shell_supervisor_restarts_signal_exits_only(): def test_prefill_heartbeat_reports_elapsed_progress(capsys): - with PrefillHeartbeat("Critic", interval_s=0.01): + with PrefillHeartbeat( + "Critic", + interval_s=0.01, + stats_provider=lambda: { + "remote_job_tokens_computed": 256, + "remote_job_tokens_total": 1024, + }, + ): time.sleep(0.025) output = capsys.readouterr().out - assert "Critic Prefill still running" in output + assert "Critic Prefill:" in output + assert "256/1024 tokens (25.0%)" in output + assert "ETA" in output def test_stage_includes_full_context_metrics(): diff --git a/tests/inference_engine/bridge/test_prefill_worker_launchd.py b/tests/inference_engine/bridge/test_prefill_worker_launchd.py index 32fb872f..24450563 100644 --- a/tests/inference_engine/bridge/test_prefill_worker_launchd.py +++ b/tests/inference_engine/bridge/test_prefill_worker_launchd.py @@ -25,6 +25,7 @@ def test_worker_installer_emits_full_cache_compatibility_contract(): "--cache-min-gb", "--memory-reserve-gb", "--estimated-snapshot-bytes-per-token", + "--prefill-compute-chunk-tokens", ): assert f"{flag}" in source assert 'PEER="${KAKEYA_WORKER_PEER:-}"' in source @@ -75,6 +76,11 @@ def test_two_mac_deployment_uses_allens_as_prefill_only(): "400000" in worker ) + assert ( + "--prefill-compute-chunk-tokens" + "256" + in worker + ) assert "--adaptive-cache" in worker assert "--window2048" in worker assert "--prefill-tps1" in worker diff --git a/tests/inference_engine/distributed/test_prefill_worker.py b/tests/inference_engine/distributed/test_prefill_worker.py index 3760be9f..3efa4e08 100644 --- a/tests/inference_engine/distributed/test_prefill_worker.py +++ b/tests/inference_engine/distributed/test_prefill_worker.py @@ -77,8 +77,14 @@ def commit_or_truncate(self, *, forwarded, accepted): assert forwarded == accepted verifier = Verifier() - engine = MLXPrefillComputeEngine(verifier, COMPAT) + engine = MLXPrefillComputeEngine( + verifier, + COMPAT, + compute_chunk_tokens=2, + ) snapshots = [] + progress = [] + engine.set_progress_callback(progress.append) def snapshot(**kwargs): snapshots.append(kwargs) @@ -96,13 +102,17 @@ def snapshot(**kwargs): compression=CompressionCodec.NONE, cancelled=threading.Event(), ) + engine.set_progress_callback(None) assert verifier.forwarded == [1, 2, 3, 4, 5] + assert progress == [2, 4, 5] assert len(blocks) == 1 assert snapshots == [{ "token_count": 5, "block_hash": hashes[-1], "compression": CompressionCodec.NONE, }] + with pytest.raises(ValueError, match="compute_chunk_tokens"): + MLXPrefillComputeEngine(verifier, COMPAT, compute_chunk_tokens=0) @pytest_asyncio.fixture @@ -362,6 +372,46 @@ def compute_prefill(self, token_ids, block_hashes, **_kwargs): jobs.close() +def test_job_store_wires_and_clears_segment_progress_callback(): + class ProgressEngine: + def __init__(self): + self.callback = None + self.callbacks = [] + + def set_progress_callback(self, callback): + self.callback = callback + self.callbacks.append(callback) + + def compute_prefill(self, token_ids, block_hashes, **_kwargs): + self.callback(2) + self.callback(4) + return [ + CacheBlock.create(block_hashes[-1], len(token_ids), b"final"), + ] + + engine = ProgressEngine() + jobs = PrefillJobStore( + engine, + PrefixCacheStore(COMPAT, max_bytes=1024, node_id="progress"), + ) + try: + job = jobs.submit( + request_id="progress", + tenant_id="tenant", + token_ids=[1, 2, 3, 4], + block_hashes=[b"a" * 32, b"b" * 32], + compatibility=COMPAT, + compression=CompressionCodec.NONE, + ) + job.future.result(timeout=1) + assert job.state == PrefillJobState.COMPLETED + assert job.tokens_computed == 4 + assert callable(engine.callbacks[0]) + assert engine.callbacks[-1] is None + finally: + jobs.close() + + def test_factory_engine_is_warmed_and_used_on_same_compute_thread(): cache = PrefixCacheStore(COMPAT, max_bytes=1024, node_id="w") created_on = [] From b102deafe97879752bea0b940216d50562adb850 Mon Sep 17 00:00:00 2001 From: fluffy314 Date: Fri, 17 Jul 2026 16:27:45 +0800 Subject: [PATCH 2/2] test(prefill): cover segmented job progress polling Keep failed and queued status doubles wire-compatible so progress accounting reaches the existing failure and timeout branches. Co-authored-by: Cursor --- .../distributed/test_prefill_cache_runtime_fallback.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/inference_engine/distributed/test_prefill_cache_runtime_fallback.py b/tests/inference_engine/distributed/test_prefill_cache_runtime_fallback.py index 8c07c613..b3d2e4c7 100644 --- a/tests/inference_engine/distributed/test_prefill_cache_runtime_fallback.py +++ b/tests/inference_engine/distributed/test_prefill_cache_runtime_fallback.py @@ -374,6 +374,7 @@ def test_import_budget_and_remote_worker_failures(monkeypatch): lambda *args, **kwargs: type("Status", (), { "status": 4, "failure_reason": "failed", + "tokens_computed": 0, })(), ) assert hook._compute_remote([1, 2], [b"a" * 32]) is None @@ -386,6 +387,7 @@ def test_import_budget_and_remote_worker_failures(monkeypatch): lambda *args, **kwargs: type("Status", (), { "status": 1, "failure_reason": "", + "tokens_computed": 0, })(), ) assert hook._compute_remote([1, 2], [b"a" * 32]) is None