diff --git a/deploy/install_prefill_worker_launchd.sh b/deploy/install_prefill_worker_launchd.sh index f027fb0b..7a7b4c38 100755 --- a/deploy/install_prefill_worker_launchd.sh +++ b/deploy/install_prefill_worker_launchd.sh @@ -87,6 +87,7 @@ cat > "$PLIST" < EOF +chmod 644 "$PLIST" launchctl bootout "gui/$(id -u)/$LABEL" 2>/dev/null || true launchctl bootstrap "gui/$(id -u)" "$PLIST" echo "installed $LABEL -> $PLIST" diff --git a/inference_engine/distributed/prefill_worker.py b/inference_engine/distributed/prefill_worker.py index e4b41e36..c30f6032 100644 --- a/inference_engine/distributed/prefill_worker.py +++ b/inference_engine/distributed/prefill_worker.py @@ -9,7 +9,7 @@ from concurrent.futures import Future, ThreadPoolExecutor from dataclasses import dataclass, field from enum import IntEnum -from typing import Protocol, Sequence +from typing import Callable, Protocol, Sequence import grpc @@ -85,9 +85,10 @@ class PrefillJobStore: def __init__( self, - engine: PrefillComputeEngine, + engine: PrefillComputeEngine | None, cache_store: PrefixCacheStore, *, + engine_factory: Callable[[], PrefillComputeEngine] | None = None, max_concurrent_jobs: int = 1, max_jobs: int = 128, completed_ttl_s: float = 600.0, @@ -100,7 +101,10 @@ def __init__( max_prompt_tokens, ) <= 0: raise ValueError("worker limits must be > 0") + if (engine is None) == (engine_factory is None): + raise ValueError("provide exactly one of engine or engine_factory") self.engine = engine + self.engine_factory = engine_factory self.cache_store = cache_store self.max_concurrent_jobs = int(max_concurrent_jobs) self.max_jobs = int(max_jobs) @@ -109,11 +113,16 @@ def __init__( self._jobs: dict[str, PrefillJob] = {} self._requests: dict[tuple[str, str], str] = {} self._lock = threading.RLock() + self._thread_local = threading.local() self._executor = ThreadPoolExecutor( max_workers=self.max_concurrent_jobs, thread_name_prefix="kakeya-prefill-worker", ) + def warmup(self) -> None: + """Construct a factory-backed engine on its eventual compute thread.""" + self._executor.submit(self._engine_for_current_thread).result() + def submit( self, *, @@ -243,7 +252,7 @@ def _run(self, job_id: str) -> None: timer.daemon = True timer.start() try: - blocks = tuple(self.engine.compute_prefill( + blocks = tuple(self._engine_for_current_thread().compute_prefill( job.token_ids, job.block_hashes, compression=job.compression, @@ -289,6 +298,16 @@ def _run(self, job_id: str) -> None: job.compute_ms = (time.perf_counter() - started) * 1000.0 job.finished_at = time.time() + def _engine_for_current_thread(self) -> PrefillComputeEngine: + if self.engine is not None: + return self.engine + engine = getattr(self._thread_local, "engine", None) + if engine is None: + assert self.engine_factory is not None + engine = self.engine_factory() + self._thread_local.engine = engine + return engine + def _gc_locked(self) -> None: cutoff = time.time() - self.completed_ttl_s for job_id, job in list(self._jobs.items()): diff --git a/scripts/start_prefill_worker_node.py b/scripts/start_prefill_worker_node.py index dda9138c..98272f66 100644 --- a/scripts/start_prefill_worker_node.py +++ b/scripts/start_prefill_worker_node.py @@ -83,27 +83,37 @@ async def serve(args) -> None: ) if args.fleet_psk_file else None ) - verifier = MLXSinkWindowVerifier(VerifierConfig( - model_id=args.model_id, - sink_size=args.sink, - window_size=args.window, - dtype=torch.bfloat16, - device="cpu", - )) + if args.max_concurrent_jobs != 1: + raise SystemExit( + "MLX prefill workers require --max-concurrent-jobs 1 so the " + "model and its stream remain on one compute thread", + ) store = PrefixCacheStore( compatibility, max_bytes=int(args.cache_gb * (1 << 30)), node_id=args.node_id, ) - engine = MLXPrefillComputeEngine(verifier, compatibility) + + def engine_factory() -> MLXPrefillComputeEngine: + verifier = MLXSinkWindowVerifier(VerifierConfig( + model_id=args.model_id, + sink_size=args.sink, + window_size=args.window, + dtype=torch.bfloat16, + device="cpu", + )) + return MLXPrefillComputeEngine(verifier, compatibility) + jobs = PrefillJobStore( - engine, + None, store, + engine_factory=engine_factory, max_concurrent_jobs=args.max_concurrent_jobs, max_jobs=args.max_jobs, completed_ttl_s=args.job_ttl_s, max_prompt_tokens=args.max_prompt_tokens, ) + jobs.warmup() def card() -> NodeCapability: inflight, queued, load, queued_tokens = jobs.stats() diff --git a/tests/inference_engine/bridge/test_prefill_worker_launchd.py b/tests/inference_engine/bridge/test_prefill_worker_launchd.py index 4efd689b..e31390ad 100644 --- a/tests/inference_engine/bridge/test_prefill_worker_launchd.py +++ b/tests/inference_engine/bridge/test_prefill_worker_launchd.py @@ -22,6 +22,7 @@ def test_worker_installer_emits_full_cache_compatibility_contract(): assert f"{flag}" in source assert 'PEER="${KAKEYA_WORKER_PEER:-}"' in source assert "--peer" in source + assert 'chmod 644 "$PLIST"' in source def test_head_runtime_discovers_and_uses_worker_cache_port(): diff --git a/tests/inference_engine/distributed/test_prefill_worker.py b/tests/inference_engine/distributed/test_prefill_worker.py index 788697a4..83d880ca 100644 --- a/tests/inference_engine/distributed/test_prefill_worker.py +++ b/tests/inference_engine/distributed/test_prefill_worker.py @@ -162,6 +162,10 @@ def test_job_store_validation_queue_stats_and_gc(): cache = PrefixCacheStore(COMPAT, max_bytes=1024, node_id="w") with pytest.raises(ValueError): PrefillJobStore(_Engine(), cache, max_jobs=0) + with pytest.raises(ValueError, match="exactly one"): + PrefillJobStore(None, cache) + with pytest.raises(ValueError, match="exactly one"): + PrefillJobStore(_Engine(), cache, engine_factory=_Engine) blocking = _Engine() blocking.block.set() jobs = PrefillJobStore(blocking, cache, max_jobs=1, max_prompt_tokens=4) @@ -230,6 +234,42 @@ def test_job_store_validation_queue_stats_and_gc(): completed_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 = [] + computed_on = [] + + class ThreadBoundEngine(_Engine): + def compute_prefill(self, *args, **kwargs): + computed_on.append(threading.get_ident()) + return super().compute_prefill(*args, **kwargs) + + def factory(): + created_on.append(threading.get_ident()) + return ThreadBoundEngine() + + jobs = PrefillJobStore(None, cache, engine_factory=factory) + try: + jobs.warmup() + job = jobs.submit( + request_id="thread-affinity", + tenant_id="tenant", + token_ids=[1, 2], + block_hashes=[b"a" * 32], + compatibility=COMPAT, + compression=CompressionCodec.NONE, + ) + for _ in range(100): + if job.state == PrefillJobState.COMPLETED: + break + time.sleep(0.005) + assert job.state == PrefillJobState.COMPLETED + assert created_on == computed_on + assert created_on[0] != threading.get_ident() + finally: + jobs.close() + + def test_job_store_failure_modes_and_precancelled_run(): cache = PrefixCacheStore(COMPAT, max_bytes=1024, node_id="w")