diff --git a/docs/features/kv_offloading_usage.md b/docs/features/kv_offloading_usage.md index b5c9d68e26..13dbd5299d 100644 --- a/docs/features/kv_offloading_usage.md +++ b/docs/features/kv_offloading_usage.md @@ -157,7 +157,7 @@ Object keys follow the same run-configuration digest scheme as the filesystem ti The P2P tier (`type: "p2p"`) shares completed KV blocks between vLLM instances over RDMA via NIXL. Each instance binds a control socket on `host:port` and exchanges blocks directly with peers — no shared filesystem required. -PYTHONHASHSEED environment variable must be set to the same fixed value on all nodes. +The `PYTHONHASHSEED` environment variable must be set to the same fixed value (e.g. `"0"`) on all nodes so that block content hashes match across instances (see [Cross-Process Sharing](#cross-process-sharing)). This is enforced: a P2P instance started without `PYTHONHASHSEED` set fails at startup, and each peer's value is verified during the connect handshake — a peer advertising a different `PYTHONHASHSEED` is rejected. | Key | Required | Default | Notes | | --- | --- | --- | --- | @@ -176,6 +176,71 @@ Rather than embedding `host`/`port` in each `secondary_tiers` entry, set them on - `VLLM_P2P_SIDE_CHANNEL_HOST` (default `localhost`): address the P2P control socket binds to. It is used **verbatim** as both the bind address and the identity peers dial back — there is no auto-detection (this mirrors `VLLM_NIXL_SIDE_CHANNEL_HOST`). The default binds the loopback interface only, so peers on another host cannot reach it. **For any cross-host P2P deployment you must set this explicitly to the node's routable IP** (e.g. the pod IP) before launching `vllm serve` — otherwise remote peers will fail to connect. The NIXL agent name is a separate per-process identifier, so peers sharing a `host:port` never collide. - `VLLM_P2P_SIDE_CHANNEL_PORT` (default `5710`): base port for the P2P control socket. The port actually bound is `VLLM_P2P_SIDE_CHANNEL_PORT + data_parallel_index` — one socket per DP replica, matching NIXL (for DP=1 the offset is 0). The peer's port is passed as `remote_port` in `kv_transfer_params`; the router/EPP that selects the DP rank (e.g. via the `X-data-parallel-rank` header) computes `remote_port = base + rank`. The DP-index offset separates replicas *within* one deployment; two co-located *deployments* (a prefiller and a decoder on the same host) still need distinct base ports (e.g. decoder base `5711`) to avoid a bind collision. +#### Orchestration-Layer Protocol + +The P2P tier does not decide *which* peer to pull from — that is the orchestration layer's job (the router/EPP and its scheduler). The orchestrator drives every transfer through a request's `kv_transfer_params` dict: it picks the request's role, allocates a unique transaction ID, and supplies the remote peer's address. All block lookup, hash matching, and NIXL transfer happen at the tier level below; the orchestrator only sets the correct role keys and enforces the allowed combinations. + +Every vLLM instance is a symmetric **peer**. Per request it acts as a **consumer** (pulls KV blocks from a remote peer's CPU cache instead of computing locally) or a **producer** (serves blocks from its own CPU cache to remote consumers) — or both, on the same session, for different requests. Roles are chosen per request by the keys below; there are no fixed prefiller/decoder processes. + +Three role keys are defined, each mapping to a sub-dict. All are optional; a request with none of them uses the tier only as a local CPU cache. + +Each key names the **remote counterpart** this peer transfers with (not this +peer's own role), so the name reads as "the remote ___ I transfer with". + +| Key | Set on | Value fields | Meaning | +| --- | --- | --- | --- | +| `remote_decoder` | prefill producer request | `kv_request_id` | Peer computes KV and keeps it available in CPU cache for the remote decoder to pull. | +| `remote_prefiller` | decode consumer request | `kv_request_id`, `remote_host`, `remote_port` | Peer pulls KV from the remote prefiller at the given address (classic P/D disaggregation). | +| `remote_kv_source` | P2P consumer request | `kv_request_id`, `remote_host`, `remote_port` | Peer looks up and pulls whatever blocks the remote source currently holds in CPU cache. | + +Field semantics: + +- `kv_request_id` (str): unique transaction ID allocated by the orchestrator and pushed to every peer involved in the transfer; used to correlate the lookup, fetch, and transfer-done messages. The producer is implicit — it serves whatever block hashes it currently holds in its CPU cache for that ID. +- `remote_host` (str): IP/hostname of the remote peer's control socket to query. Must be the peer's routable node IP (see [Environment Variables](#environment-variables)). +- `remote_port` (int): the peer's bound control-socket port, i.e. `base + data_parallel_index` for the selected DP rank. + +Allowed and forbidden combinations: + +- **`remote_decoder` + `remote_kv_source`** is the only legal multi-key combination: a prefill producer may *also* act as a P2P consumer for the same request — skipping prefix prefill by pulling cached blocks from a source while still keeping its own computed blocks available for a downstream decoder. +- Forbidden: `remote_prefiller` + `remote_decoder` (contradictory roles), `remote_prefiller` + `remote_kv_source` (two competing fetch sources), and all three together. + +Minimal examples (values that would appear in the request's `kv_transfer_params`): + +```python +# Prefill producer — compute and keep KV for a remote decoder to pull +kv_transfer_params = {"remote_decoder": {"kv_request_id": ""}} + +# Decode consumer — pull KV from a specific prefiller (classic P/D) +kv_transfer_params = { + "remote_prefiller": { + "kv_request_id": "", + "remote_host": "", + "remote_port": 5710, + } +} + +# P2P consumer — pull whatever the source already has cached +kv_transfer_params = { + "remote_kv_source": { + "kv_request_id": "", + "remote_host": "", + "remote_port": 5710, + } +} +``` + +Runtime handshake for a P2P (or P/D) pull, once the orchestrator has set the keys above: + +1. Both peers already have listener threads on their control sockets (see [Environment Variables](#environment-variables)). +2. **Lookup.** The consumer's tiering manager does per-block lookups; in P2P mode the tier returns `None` and registers the key. At `on_schedule_end` the consumer sends one **`LookupMsg`** (`kv_request_id` + block hashes) to the peer, per request, per step. +3. The producer matches those hashes against its local CPU cache and replies with a **`LookupRespMsg`** carrying the hit block hashes. +4. **Resolve.** Retried lookups now return hit / miss / in-flight. The consumer calls `submit_load` for hits only, allocating CPU slots only for hits. +5. The consumer sends a **`FetchMsg`** (`kv_request_id`, block hashes, destination block indexes). +6. The producer performs the **NIXL WRITE** transfer and sends **`TransferDone`** with a success status. +7. On `get_finished`, hits are loaded into GPU as ordinary cache hits; misses are recomputed by the engine. + +In classic **P/D mode** (`remote_prefiller` set, no `remote_kv_source`), the lookup phase (steps 2–4) is skipped: the decode consumer assumes the prefiller holds all of the request's blocks, so every block `lookup()` returns an immediate hit and the consumer jumps straight to the **`FetchMsg`** in step 5. The `LookupMsg`/`LookupRespMsg` round-trip only happens in P2P mode, where the consumer does not know in advance which blocks the peer has cached. + ## Tuning Tips - `cpu_bytes_to_use`: a bigger CPU tier means fewer trips to slower secondary tiers and a higher hit rate. The value is total across all workers, not per-worker. Leave headroom for the rest of the host workload. diff --git a/tests/v1/kv_offload/tiering/p2p/p2p_connector_proxy.py b/tests/v1/kv_offload/tiering/p2p/p2p_connector_proxy.py index 2f0df89902..630fcdef44 100644 --- a/tests/v1/kv_offload/tiering/p2p/p2p_connector_proxy.py +++ b/tests/v1/kv_offload/tiering/p2p/p2p_connector_proxy.py @@ -81,6 +81,8 @@ async def lifespan(app: FastAPI): app.state.decode_dp_iterator = itertools.cycle(range(global_args.decoder_dp_size)) mode = "decoder-first" if global_args.decoder_first else "prefiller-first" + if global_args.p2p: + mode += ",p2p" pd_host = global_args.p2p_connector_host pd_port = global_args.p2p_connector_port n_pref = len(app.state.prefill_clients) @@ -145,6 +147,14 @@ def parse_args(): help="Send decode request before prefill so decoder is already " "waiting when KV blocks arrive (decoder-first mode)", ) + p.add_argument( + "--p2p", + action="store_true", + help="P2P mode: do not inject kv_transfer_params on the prefiller; " + "on the decoder, inject kv_transfer_params with a top-level " + "'remote_kv_source' block ({kv_request_id, remote_host, remote_port}) " + "instead of the default 'remote_prefiller' block.", + ) p.add_argument( "--prefiller-dp-size", type=int, @@ -206,11 +216,14 @@ def _auth_headers(request_id: str) -> dict: async def _prefill(client_info, endpoint, req_data, request_id, dp_rank=None): """Send a prefill-only request (max_tokens=1) to the prefiller.""" data = req_data.copy() - data["kv_transfer_params"] = { - "decode": { - "kv_request_id": request_id, - }, - } + if global_args.p2p: + data.pop("kv_transfer_params", None) + else: + data["kv_transfer_params"] = { + "remote_decoder": { + "kv_request_id": request_id, + }, + } data["stream"] = False data["max_tokens"] = 1 data.pop("max_completion_tokens", None) @@ -259,8 +272,9 @@ async def _handle_completions(api: str, request: Request): # Inject the prefiller's P2PConnector address so the decoder can pull # KV blocks from it. remote_port = base + prefill_rank targets the # replica that produced the KV (base+0 == base when dp=1). + decoder_key = "remote_kv_source" if global_args.p2p else "remote_prefiller" req_data["kv_transfer_params"] = { - "prefill": { + decoder_key: { "kv_request_id": request_id, "remote_host": global_args.p2p_connector_host, "remote_port": global_args.p2p_connector_port + prefill_rank, @@ -311,9 +325,10 @@ async def _handle_completions_decoder_first(api: str, request: Request): prefill_client = _get_next(request.app, "prefill") decode_client = _get_next(request.app, "decode") + decoder_key = "remote_kv_source" if global_args.p2p else "remote_prefiller" decode_data = req_data.copy() decode_data["kv_transfer_params"] = { - "prefill": { + decoder_key: { "kv_request_id": request_id, "remote_host": global_args.p2p_connector_host, "remote_port": global_args.p2p_connector_port + prefill_rank, diff --git a/tests/v1/kv_offload/tiering/p2p/test_manager.py b/tests/v1/kv_offload/tiering/p2p/test_manager.py index 01fd295302..cd6e8f382a 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_manager.py +++ b/tests/v1/kv_offload/tiering/p2p/test_manager.py @@ -21,9 +21,11 @@ from vllm.v1.kv_offload.tiering.p2p import manager as manager_module from vllm.v1.kv_offload.tiering.p2p.manager import ( _UNBOUND_STORE_TIMEOUT_S, P2PSecondaryTierManager, + _annotate_req_context, ) from vllm.v1.kv_offload.tiering.p2p.session import ( LoadResult, + SessionCloseResult, SessionPollResult, StoreResult, ) @@ -33,15 +35,15 @@ from vllm.v1.kv_offload.tiering.p2p.session import ( # --------------------------------------------------------------------------- -def _prefill_kv_params( +def _remote_prefiller_kv_params( remote_host: str = "10.0.0.1", remote_port: int = 8000, kv_request_id: str = "req-1", ) -> dict: - """Decoder-side kv_transfer_params: ``prefill`` sub-dict carries + """Decoder-side kv_transfer_params: ``remote_prefiller`` sub-dict carries kv_request_id + remote_host + remote_port.""" return { - "prefill": { + "remote_prefiller": { "kv_request_id": kv_request_id, "remote_host": remote_host, "remote_port": remote_port, @@ -49,14 +51,34 @@ def _prefill_kv_params( } -def _decode_kv_params(kv_request_id: str = "req-1") -> dict: - """Prefiller-side kv_transfer_params: ``decode`` sub-dict carries +def _remote_kv_source_kv_params( + remote_host: str = "10.0.0.1", + remote_port: int = 8000, + kv_request_id: str = "req-1", +) -> dict: + """Symmetric-P2P consumer kv_transfer_params: ``remote_kv_source`` sub-dict has + the same shape as ``remote_prefiller`` (kv_request_id + remote_host + port).""" + return { + "remote_kv_source": { + "kv_request_id": kv_request_id, + "remote_host": remote_host, + "remote_port": remote_port, + }, + } + + +def _remote_decoder_kv_params(kv_request_id: str = "req-1") -> dict: + """Prefiller-side kv_transfer_params: ``remote_decoder`` sub-dict carries kv_request_id only.""" - return {"decode": {"kv_request_id": kv_request_id}} + return {"remote_decoder": {"kv_request_id": kv_request_id}} def _req_context(kv_params: dict | None = None) -> ReqContext: - return ReqContext(req_id="test", kv_transfer_params=kv_params) + ctx = ReqContext(req_id="test", kv_transfer_params=kv_params) + # Mirror on_new_request: parse the P2P routing state once and cache it, + # so lookup/submit_*/on_request_finished can read it back via get_state. + _annotate_req_context(ctx) + return ctx def _job_metadata( @@ -82,38 +104,78 @@ def _make_manager() -> P2PSecondaryTierManager: """Create a manager with stubbed __init__.""" mgr = P2PSecondaryTierManager.__new__(P2PSecondaryTierManager) mgr._local_id = "127.0.0.1:7777" + mgr._hash_seed = "0" mgr._finished_jobs = [] mgr._failed_req_ids = set() mgr._sessions = {} mgr._kv_to_session = {} mgr._unbound_stores = {} + mgr._failed_serve_ctxs = [] return mgr +def _init_offloading_spec() -> SimpleNamespace: + """Minimal offloading_spec for driving the real __init__.""" + return SimpleNamespace( + config=SimpleNamespace(parallel=SimpleNamespace(data_parallel_index=0)), + blocks_per_chunk=1, + ) + + # --------------------------------------------------------------------------- -# Tests for _remote_id_from_params +# Tests for __init__ PYTHONHASHSEED assertion # --------------------------------------------------------------------------- -class TestRemoteIdFromParams: +class TestInitHashSeedAssertion: + def test_missing_pythonhashseed_raises(self, monkeypatch): + """P2P instance refuses to start when PYTHONHASHSEED is unset.""" + monkeypatch.delenv("PYTHONHASHSEED", raising=False) + with pytest.raises(ValueError, match="PYTHONHASHSEED"): + P2PSecondaryTierManager( + offloading_spec=_init_offloading_spec(), + primary_kv_view=memoryview(bytearray(16)), + ) + + def test_pythonhashseed_set_succeeds(self, monkeypatch): + """With PYTHONHASHSEED set, __init__ records it for the handshake.""" + monkeypatch.setenv("PYTHONHASHSEED", "12345") + monkeypatch.setattr(manager_module, "NixlTransport", lambda *a, **k: object()) + monkeypatch.setattr(manager_module, "ZmqTransport", lambda *a, **k: object()) + monkeypatch.setattr( + manager_module.FileMapper, + "from_offloading_spec", + lambda **k: SimpleNamespace(get_run_config=lambda: {}), + ) + mgr = P2PSecondaryTierManager( + offloading_spec=_init_offloading_spec(), + primary_kv_view=memoryview(bytearray(16)), + ) + assert mgr._hash_seed == "12345" + + +# --------------------------------------------------------------------------- +# Tests for _peer_id_from_params +# --------------------------------------------------------------------------- + + +class TestPeerIdFromParams: def test_valid_params(self): - result = P2PSecondaryTierManager._remote_id_from_params( + result = manager_module._peer_id_from_params( {"remote_host": "10.0.0.1", "remote_port": 8000} ) assert result == "10.0.0.1:8000" def test_missing_host(self): - result = P2PSecondaryTierManager._remote_id_from_params({"remote_port": 8000}) + result = manager_module._peer_id_from_params({"remote_port": 8000}) assert result is None def test_missing_port(self): - result = P2PSecondaryTierManager._remote_id_from_params( - {"remote_host": "10.0.0.1"} - ) + result = manager_module._peer_id_from_params({"remote_host": "10.0.0.1"}) assert result is None def test_empty_dict(self): - result = P2PSecondaryTierManager._remote_id_from_params({}) + result = manager_module._peer_id_from_params({}) assert result is None @@ -130,35 +192,105 @@ class TestLookup: def test_lookup_returns_miss_without_required_fields(self): mgr = _make_manager() - ctx = _req_context(kv_params={"prefill": {"remote_host": "x"}}) + ctx = _req_context(kv_params={"remote_prefiller": {"remote_host": "x"}}) assert mgr.lookup(b"key", ctx) is LookupResult.MISS def test_lookup_returns_hit_for_valid_request(self): mgr = _make_manager() - ctx = _req_context(kv_params=_prefill_kv_params()) + ctx = _req_context(kv_params=_remote_prefiller_kv_params()) assert mgr.lookup(b"key", ctx) is LookupResult.HIT def test_lookup_returns_miss_for_failed_request(self): mgr = _make_manager() mgr._failed_req_ids.add("req-1") - ctx = _req_context(kv_params=_prefill_kv_params(kv_request_id="req-1")) + ctx = _req_context(kv_params=_remote_prefiller_kv_params(kv_request_id="req-1")) assert mgr.lookup(b"key", ctx) is LookupResult.MISS def test_lookup_returns_hit_for_different_request_id(self): mgr = _make_manager() mgr._failed_req_ids.add("req-1") - ctx = _req_context(kv_params=_prefill_kv_params(kv_request_id="req-2")) + ctx = _req_context(kv_params=_remote_prefiller_kv_params(kv_request_id="req-2")) assert mgr.lookup(b"key", ctx) is LookupResult.HIT def test_lookup_returns_miss_without_prefill_key(self): - """No ``prefill`` sub-dict means the request was not routed for + """No ``remote_prefiller`` sub-dict means the request was not routed for remote prefill — local prefill should run instead, so lookup() - returns MISS even when a stale ``decode`` block is present.""" + returns MISS even when a stale ``remote_decoder`` block is present.""" mgr = _make_manager() - ctx = _req_context(kv_params=_decode_kv_params()) + ctx = _req_context(kv_params=_remote_decoder_kv_params()) assert mgr.lookup(b"key", ctx) is LookupResult.MISS +# --------------------------------------------------------------------------- +# Tests for serve_external_requests +# --------------------------------------------------------------------------- + + +class _RecordingParent: + """Minimal ParentManager stub recording on_request_finished calls.""" + + def __init__(self) -> None: + self.finished: list[str] = [] + + def on_new_request(self, ctx): + from vllm.v1.kv_offload.base import RequestOffloadingContext + + return RequestOffloadingContext() + + def lookup(self, key, ctx): + return LookupResult.MISS + + def create_store_job(self, keys, ctx): + raise AssertionError("unreachable") + + def on_request_finished(self, ctx) -> None: + self.finished.append(ctx.req_id) + + +class _RecordingSession: + """Fake P2PSession that records the parent it was served with.""" + + def __init__(self) -> None: + self.served_with: list[object] = [] + + def serve_external_requests(self, parent) -> None: + self.served_with.append(parent) + + +class TestServeExternalRequests: + def test_flushes_failed_serve_ctxs_then_serves_each_session(self): + """serve_external_requests releases the failed serves left by + reaped sessions via parent.on_request_finished (clearing the + queue), then delegates to every live session with the same parent.""" + mgr = _make_manager() + ctx = ReqContext(req_id="p2p:peer:req-1:lu1") + mgr._failed_serve_ctxs = [ctx] + sess_a = _RecordingSession() + sess_b = _RecordingSession() + mgr._sessions = {"a": sess_a, "b": sess_b} # type: ignore[assignment] + + parent = _RecordingParent() + mgr.serve_external_requests(parent) # type: ignore[arg-type] + + # Failed serve released and queue cleared. + assert parent.finished == ["p2p:peer:req-1:lu1"] + assert mgr._failed_serve_ctxs == [] + # Every live session served with the same parent handle. + assert sess_a.served_with == [parent] + assert sess_b.served_with == [parent] + + def test_no_failed_serves_still_serves_sessions(self): + mgr = _make_manager() + sess = _RecordingSession() + mgr._sessions = {"a": sess} # type: ignore[assignment] + + parent = _RecordingParent() + mgr.serve_external_requests(parent) # type: ignore[arg-type] + + assert parent.finished == [] + assert sess.served_with == [parent] + + # --------------------------------------------------------------------------- # Tests for submit_store # --------------------------------------------------------------------------- @@ -166,16 +298,16 @@ class TestLookup: class TestSubmitStore: def test_no_decode_succeeds_immediately(self): - """Without a ``decode`` block, job succeeds immediately.""" + """Without a ``remote_decoder`` block, job succeeds immediately.""" mgr = _make_manager() job = _job_metadata(job_id=1, kv_params={}) mgr.submit_store(job) assert mgr._finished_jobs == [JobResult(job_id=1, success=True)] def test_missing_kv_request_id_fails(self): - """Missing kv_request_id inside ``decode`` fails the job.""" + """Missing kv_request_id inside ``remote_decoder`` fails the job.""" mgr = _make_manager() - params: dict = {"decode": {}} + params: dict = {"remote_decoder": {}} job = _job_metadata(job_id=1, kv_params=params) mgr.submit_store(job) assert mgr._finished_jobs == [JobResult(job_id=1, success=False)] @@ -189,7 +321,7 @@ class TestSubmitStore: job_id=1, keys=[b"k1", b"k2"], block_ids=[3, 4], - kv_params=_decode_kv_params(kv_request_id="req-1"), + kv_params=_remote_decoder_kv_params(kv_request_id="req-1"), ) mgr.submit_store(job) @@ -220,7 +352,7 @@ class TestSubmitStore: job_id=7, keys=[b"k1", b"k2"], block_ids=[3, 4], - kv_params=_decode_kv_params(kv_request_id="req-1"), + kv_params=_remote_decoder_kv_params(kv_request_id="req-1"), ) mgr.submit_store(job) @@ -233,9 +365,9 @@ class TestSubmitStore: def test_extra_top_level_keys_are_ignored(self): """Producer-side kv_transfer_params should not pre-create a session even when a stale caller still passes a top-level - ``remote_host``/``remote_port`` next to ``decode``.""" + ``remote_host``/``remote_port`` next to ``remote_decoder``.""" mgr = _make_manager() - params = _decode_kv_params() + params = _remote_decoder_kv_params() params["remote_host"] = "stale" params["remote_port"] = 12345 job = _job_metadata(job_id=1, kv_params=params) @@ -262,7 +394,7 @@ class TestSubmitLoad: """Empty key list succeeds immediately.""" mgr = _make_manager() job = _job_metadata( - job_id=1, keys=[], block_ids=[], kv_params=_prefill_kv_params() + job_id=1, keys=[], block_ids=[], kv_params=_remote_prefiller_kv_params() ) mgr.submit_load(job) assert mgr._finished_jobs == [JobResult(job_id=1, success=True)] @@ -270,7 +402,7 @@ class TestSubmitLoad: def test_no_session_fails(self): """No session for peer fails and marks request failed.""" mgr = _make_manager() - job = _job_metadata(job_id=1, kv_params=_prefill_kv_params()) + job = _job_metadata(job_id=1, kv_params=_remote_prefiller_kv_params()) mgr.submit_load(job) assert mgr._finished_jobs == [JobResult(job_id=1, success=False)] assert "req-1" in mgr._failed_req_ids @@ -286,7 +418,7 @@ class TestSubmitLoad: job_id=42, keys=[b"k1", b"k2"], block_ids=[5, 6], - kv_params=_prefill_kv_params(kv_request_id="req-42"), + kv_params=_remote_prefiller_kv_params(kv_request_id="req-42"), ) mgr.submit_load(job) @@ -294,6 +426,20 @@ class TestSubmitLoad: assert mgr._finished_jobs == [] assert "req-42" not in mgr._failed_req_ids + def test_missing_consumer_flag_fails(self): + """Peer fields present but neither do_remote_prefill nor + do_p2p_fetch is set — submit_load fails the job rather than + emit a stray FetchMsg.""" + mgr = _make_manager() + params = { + "remote_host": "10.0.0.1", + "remote_port": 8000, + "kv_request_id": "req-1", + } + job = _job_metadata(job_id=1, kv_params=params) + mgr.submit_load(job) + assert mgr._finished_jobs == [JobResult(job_id=1, success=False)] + # --------------------------------------------------------------------------- # Tests for on_request_finished @@ -308,7 +454,7 @@ class TestOnRequestFinished: def test_prunes_failed_req_ids(self): mgr = self._make_with_failed() - ctx = _req_context(kv_params=_prefill_kv_params(kv_request_id="req-1")) + ctx = _req_context(kv_params=_remote_prefiller_kv_params(kv_request_id="req-1")) mgr.on_request_finished(ctx) assert "req-1" not in mgr._failed_req_ids @@ -325,14 +471,26 @@ class TestOnRequestFinished: assert "req-1" in mgr._failed_req_ids def test_decoder_side_calls_session_finish_request(self): - """Decoder-side finish (``prefill`` set) still routes via peer_id + """Decoder-side finish (``remote_prefiller`` set) still routes via peer_id because the consumer addresses the producer it loaded from. The session's finish_request cancels the client-role load.""" mgr = _make_manager() peer_id = "10.0.0.1:8000" session = _FakeSession(peer_id=peer_id) mgr._sessions[peer_id] = session - ctx = _req_context(kv_params=_prefill_kv_params(kv_request_id="req-1")) + ctx = _req_context(kv_params=_remote_prefiller_kv_params(kv_request_id="req-1")) + mgr.on_request_finished(ctx) + assert session.finishes == ["req-1"] + + def test_p2p_consumer_side_calls_session_finish_request(self): + """Symmetric-P2P consumer finish (``remote_kv_source`` set) routes via peer_id + so the session drops any pending lookups (cancel_lookups) and + cancels any inbound load.""" + mgr = _make_manager() + peer_id = "10.0.0.1:8000" + session = _FakeSession(peer_id=peer_id) + mgr._sessions[peer_id] = session + ctx = _req_context(kv_params=_remote_kv_source_kv_params(kv_request_id="req-1")) mgr.on_request_finished(ctx) assert session.finishes == ["req-1"] @@ -342,7 +500,7 @@ class TestOnRequestFinished: mgr = _make_manager() bound = _FakeSession(peer_id="some-peer:1", connected=True) mgr._kv_to_session["req-1"] = bound # type: ignore[assignment] - ctx = _req_context(kv_params=_decode_kv_params(kv_request_id="req-1")) + ctx = _req_context(kv_params=_remote_decoder_kv_params(kv_request_id="req-1")) mgr.on_request_finished(ctx) assert bound.finishes == ["req-1"] assert "req-1" not in mgr._kv_to_session @@ -360,7 +518,7 @@ class TestOnRequestFinished: _UnboundStoreBatch(job_id=10, keys=[b"k"], block_ids=[0]), _UnboundStoreBatch(job_id=11, keys=[b"k2"], block_ids=[1]), ] - ctx = _req_context(kv_params=_decode_kv_params(kv_request_id="req-1")) + ctx = _req_context(kv_params=_remote_decoder_kv_params(kv_request_id="req-1")) mgr.on_request_finished(ctx) assert "req-1" in mgr._unbound_stores assert [b.job_id for b in mgr._unbound_stores["req-1"]] == [10, 11] @@ -378,11 +536,19 @@ class _FakeServerHalf: def __init__(self) -> None: self._inflight: dict[int, object] = {} + @property + def has_inflight_transfers(self) -> bool: + return bool(self._inflight) + class _FakeClientHalf: def __init__(self) -> None: self._inbound: dict[int, object] = {} + @property + def has_active_loads(self) -> bool: + return bool(self._inbound) + class _FakeSession: """Fake bidirectional session that returns canned poll() results.""" @@ -395,8 +561,10 @@ class _FakeSession: loads: list[LoadResult] | None = None, stores: list[StoreResult] | None = None, new_fetch_ids: list[str] | None = None, - close_loads: list[tuple[int, str]] | None = None, + close_jobs: list[int] | None = None, + close_req_ids: list[str] | None = None, close_stores: list[int] | None = None, + close_failed_serves: list[ReqContext] | None = None, ) -> None: self.peer_id = peer_id self.alive = alive @@ -405,18 +573,24 @@ class _FakeSession: self._loads = loads or [] self._stores = stores or [] self._new_fetch_ids = new_fetch_ids or [] - self._close_loads = close_loads or [] + self._close_jobs = close_jobs or [] + self._close_req_ids = close_req_ids or [] self._close_stores = close_stores or [] + self._close_failed_serves = close_failed_serves or [] self.requests: list[tuple[int, str]] = [] self.stores_added: list[tuple[str, list, object, int]] = [] self.attached: list[object] = [] self.finishes: list[str] = [] # Mirror P2PSession._server._inflight (transfer_id → handle) and - # P2PSession._client._inbound for the shutdown-drain and drain_jobs - # paths. Tests populate _server._inflight when needed. + # P2PSession._client.has_active_loads for the shutdown-drain and + # drain_jobs paths. Tests populate _server._inflight when needed. self._server = _FakeServerHalf() self._client = _FakeClientHalf() + @property + def has_pending_work(self) -> bool: + return self._client.has_active_loads or self._server.has_inflight_transfers + def poll(self): result = SessionPollResult( loads=self._loads, @@ -442,7 +616,12 @@ class _FakeSession: self.finishes.append(kv_request_id) def close(self): - return self._close_loads, self._close_stores + return SessionCloseResult( + failed_jobs=self._close_jobs, + failed_req_ids=self._close_req_ids, + failed_stores=self._close_stores, + failed_serves=self._close_failed_serves, + ) class TestGetFinished: @@ -484,7 +663,8 @@ class TestGetFinished: peer_id="dead:1234", alive=False, connected=True, - close_loads=[(20, "req-load")], + close_jobs=[20], + close_req_ids=["req-load"], close_stores=[10, 11], ) mgr._sessions["dead:1234"] = dead # type: ignore[assignment] @@ -498,6 +678,29 @@ class TestGetFinished: assert "dead:1234" not in mgr._sessions assert "req-load" in mgr._failed_req_ids + def test_reap_fails_probes(self): + """A reaped session's in-flight lookups land in _failed_req_ids so + the consumer's lookup() returns MISS instead of RETRY forever.""" + + class FakeData: + def remove_remote_peer(self, pid): + pass + + mgr = self._make() + mgr._data = FakeData() # type: ignore[assignment] + dead = _FakeSession( + peer_id="dead:1234", + alive=False, + connected=True, + close_req_ids=["req-probe-1", "req-probe-2"], + ) + mgr._sessions["dead:1234"] = dead # type: ignore[assignment] + + list(mgr.get_finished_jobs()) + assert "dead:1234" not in mgr._sessions + assert "req-probe-1" in mgr._failed_req_ids + assert "req-probe-2" in mgr._failed_req_ids + def test_unbound_store_kept_within_timeout(self): """Recently-parked unbound stores stay across a poll.""" mgr = self._make() @@ -537,7 +740,7 @@ class TestGetFinished: submitted_at stamp so the unbound-store sweep can age it out.""" mgr = _make_manager() job = _job_metadata( - job_id=1, kv_params=_decode_kv_params(kv_request_id="req-1") + job_id=1, kv_params=_remote_decoder_kv_params(kv_request_id="req-1") ) before = time.monotonic() mgr.submit_store(job) @@ -881,21 +1084,21 @@ class TestBidirectionalManager: b_loads_kv = "req-BtoA-load" # B loads, A serves a_decoder_params = { - "prefill": { + "remote_prefiller": { "kv_request_id": a_loads_kv, "remote_host": "B", "remote_port": 2, }, } b_decoder_params = { - "prefill": { + "remote_prefiller": { "kv_request_id": b_loads_kv, "remote_host": "A", "remote_port": 1, }, } - a_prefiller_params = {"decode": {"kv_request_id": b_loads_kv}} - b_prefiller_params = {"decode": {"kv_request_id": a_loads_kv}} + a_prefiller_params = {"remote_decoder": {"kv_request_id": b_loads_kv}} + b_prefiller_params = {"remote_decoder": {"kv_request_id": a_loads_kv}} # 1. Both sides open client-role sessions toward the peer. mgr_a.on_new_request(_req_context(a_decoder_params)) @@ -1065,7 +1268,8 @@ class TestPollOnce: peer_id=peer_dead, alive=False, connected=True, - close_loads=[(33, "req-33")], + close_jobs=[33], + close_req_ids=["req-33"], close_stores=[44], ) mgr._sessions[peer_dead] = dead # type: ignore[assignment] @@ -1309,13 +1513,13 @@ class TestConnectionDeathMidTransfer: mgr_a, mgr_b = _build_paired_managers() a_decoder_params = { - "prefill": { + "remote_prefiller": { "kv_request_id": "req-load", "remote_host": "B", "remote_port": 2, }, } - a_prefiller_params = {"decode": {"kv_request_id": "req-store"}} + a_prefiller_params = {"remote_decoder": {"kv_request_id": "req-store"}} # Open the outbound session A->B and submit one load + one store. mgr_a.on_new_request(_req_context(a_decoder_params)) @@ -1394,6 +1598,7 @@ class TestBindHostPortDefaults: identity (``host:port``) stays decoupled from the NIXL agent name (a uuid). """ + monkeypatch.setenv("PYTHONHASHSEED", "0") monkeypatch.setattr( manager_module, "FileMapper", diff --git a/tests/v1/kv_offload/tiering/p2p/test_sessions.py b/tests/v1/kv_offload/tiering/p2p/test_sessions.py index d1c6ad0cda..cab4ce6705 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_sessions.py +++ b/tests/v1/kv_offload/tiering/p2p/test_sessions.py @@ -13,9 +13,18 @@ completes its own load. from __future__ import annotations import time +from collections.abc import Sequence +import numpy as np import pytest +from vllm.v1.kv_offload.base import ( + LookupResult, + OffloadKey, + ReqContext, + RequestOffloadingContext, +) +from vllm.v1.kv_offload.tiering.base import JobMetadata from vllm.v1.kv_offload.tiering.p2p.session import ( LoadResult, P2PSession, @@ -33,6 +42,8 @@ from vllm.v1.kv_offload.tiering.p2p.session.protocol import ( ConnectMsg, DisconnectMsg, FetchMsg, + LookupMsg, + LookupRespMsg, TransferDoneMsg, ) from vllm.v1.kv_offload.tiering.p2p.session.server import ( @@ -48,6 +59,11 @@ from vllm.v1.kv_offload.tiering.p2p.session.session import ( # --------------------------------------------------------------------------- +# Shared PYTHONHASHSEED used by the session under test and the fake peer's +# ConnectMsg so the handshake succeeds unless a test overrides one side. +_DEFAULT_HASH_SEED = "0" + + class FakeDataTransport: """Minimal fake DataTransport for testing sessions.""" @@ -146,12 +162,16 @@ class FakeConnection: self._inbox: list[dict] = [] self._sent: list[dict] = [] self._closed = False + # When True, send() raises to simulate a broken/dead connection. + self.fail_send = False @property def alive(self) -> bool: return not self._closed def send(self, msg: dict) -> None: + if self.fail_send: + raise ConnectionError("simulated dead connection") self._sent.append(msg) def recv(self) -> list[dict]: @@ -173,6 +193,7 @@ def _peer_connect_msg( peer_id: str = "peer:8000", block_len: int = 4096, fingerprint: str | None = None, + hash_seed: str = _DEFAULT_HASH_SEED, ) -> dict: """Build a ConnectMsg as if the peer sent it.""" msg = { @@ -182,17 +203,82 @@ def _peer_connect_msg( ConnectMsg.BASE_ADDR: 0x2000, ConnectMsg.NUM_BLOCKS: 16, ConnectMsg.BLOCK_LEN: block_len, + ConnectMsg.HASH_SEED: hash_seed, } if fingerprint is not None: msg[ConnectMsg.CONFIG_FINGERPRINT] = fingerprint return msg +class FakeParent: + """Configurable :class:`ParentManager` for server-role tests. + + ``stored`` is the dict of ready blocks (key → primary block_id). + ``pending`` and ``retry`` script the first lookup() result for those + keys; subsequent lookups behave normally (a key that promised + HIT_PENDING / RETRY can later be promoted to HIT by adding it to + ``stored`` and removing it from ``pending``/``retry``). ``calls`` + captures every parent invocation in order for assertions. + + Injected per-step via ``session.serve_external_requests(parent)`` — + not held by the session, matching how ``TieringOffloadingManager`` + hands the tier a handle valid only for that call. + """ + + def __init__( + self, + stored: dict[OffloadKey, int] | None = None, + pending: set[OffloadKey] | None = None, + retry: set[OffloadKey] | None = None, + ) -> None: + self.stored: dict[OffloadKey, int] = dict(stored or {}) + self.pending: set[OffloadKey] = set(pending or ()) + self.retry: set[OffloadKey] = set(retry or ()) + self._next_job_id: int = 1000 + self.calls: list[tuple] = [] + + def on_new_request(self, ctx: ReqContext) -> RequestOffloadingContext: + self.calls.append(("on_new_request", ctx.req_id)) + return RequestOffloadingContext() + + def lookup(self, key: OffloadKey, ctx: ReqContext) -> LookupResult: + self.calls.append(("lookup", key, ctx.req_id)) + if key in self.pending: + return LookupResult.HIT_PENDING + if key in self.retry: + return LookupResult.RETRY + if key in self.stored: + return LookupResult.HIT + return LookupResult.MISS + + def create_store_job( + self, + keys: Sequence[OffloadKey], + ctx: ReqContext, + ) -> JobMetadata: + keys_list = list(keys) + self.calls.append(("create_store_job", tuple(keys_list), ctx.req_id)) + block_ids = np.array([self.stored[k] for k in keys_list], dtype=np.int32) + job_id = self._next_job_id + self._next_job_id += 1 + return JobMetadata( + job_id=job_id, + keys=keys_list, + block_ids=block_ids, + is_promotion=False, + req_context=ctx, + ) + + def on_request_finished(self, ctx: ReqContext) -> None: + self.calls.append(("on_request_finished", ctx.req_id)) + + def _make_session( conn: FakeConnection | None = None, transport: FakeDataTransport | None = None, peer_id: str = "peer:8000", local_id: str = "local:9000", + local_hash_seed: str = _DEFAULT_HASH_SEED, ) -> tuple[P2PSession, FakeConnection, FakeDataTransport]: if conn is None: conn = FakeConnection(peer_id=peer_id) @@ -203,11 +289,17 @@ def _make_session( local_id=local_id, transport=transport, # type: ignore[arg-type] local_block_len=transport.block_len, + local_hash_seed=local_hash_seed, conn=conn, # type: ignore[arg-type] ) return session, conn, transport +def _serve(session: P2PSession, parent: FakeParent) -> None: + """Resolve enqueued inbound lookups, as the manager does each step.""" + session.serve_external_requests(parent) # type: ignore[arg-type] + + def _activate( session: P2PSession, conn: FakeConnection, peer_id: str = "peer:8000" ) -> None: @@ -217,6 +309,42 @@ def _activate( session.poll() +# --- Accessors for the server role's consolidated per-request state. +# ServerRole keeps one _ServerRequestState per kv_request_id; entries are +# garbage-collected once fully idle, so "no outbound"/"no abort" reads as +# either a missing entry or a None field. These helpers paper over that. + + +def _srv_outbound(session: P2PSession, kv_request_id: str): + """Outbound serve state for a kv_request_id, or None (idle / GC'd).""" + st = session._server._requests.get(kv_request_id) + return st.outbound if st is not None else None + + +def _srv_lookups(session: P2PSession) -> list: + """Every parked inbound _ActiveLookup across all requests.""" + return [ + lu for st in session._server._requests.values() for lu in st.lookups.values() + ] + + +def _srv_abort_started(session: P2PSession, kv_request_id: str) -> float | None: + """Pending-abort start time for a kv_request_id, or None.""" + st = session._server._requests.get(kv_request_id) + return st.abort_started_at if st is not None else None + + +def _srv_inflight_count(session: P2PSession, kv_request_id: str) -> int: + """Inflight-transfer count tracked for a kv_request_id (0 if idle).""" + st = session._server._requests.get(kv_request_id) + return len(st.inflight_tids) if st is not None else 0 + + +def _srv_total_inflight(session: P2PSession) -> int: + """Sum of per-request inflight counts across all requests.""" + return sum(len(st.inflight_tids) for st in session._server._requests.values()) + + # --------------------------------------------------------------------------- # Connect / handshake # --------------------------------------------------------------------------- @@ -300,6 +428,29 @@ class TestConnectHandshake: session.poll() assert "peer:8000" in transport._remote_peers + def test_hash_seed_mismatch_marks_dead(self): + """Mismatched PYTHONHASHSEED rejects peer and marks connection dead.""" + session, conn, transport = _make_session(local_hash_seed="0") + conn.enqueue(_peer_connect_msg(hash_seed="12345")) # mismatch + session.poll() + assert "peer:8000" not in transport._remote_peers + assert not session.alive + assert not any(m[TYPE_KEY] == ConnectAckMsg.TYPE for m in conn._sent) + + def test_hash_seed_match_succeeds(self): + """Matching PYTHONHASHSEED registers the peer and acks.""" + session, conn, transport = _make_session(local_hash_seed="12345") + conn.enqueue(_peer_connect_msg(hash_seed="12345")) + session.poll() + assert "peer:8000" in transport._remote_peers + assert session.alive + assert any(m[TYPE_KEY] == ConnectAckMsg.TYPE for m in conn._sent) + + def test_hash_seed_advertised_in_connect_msg(self): + """Session advertises its own PYTHONHASHSEED in the ConnectMsg.""" + _, conn, _ = _make_session(local_hash_seed="777") + assert conn._sent[0][ConnectMsg.HASH_SEED] == "777" + # --------------------------------------------------------------------------- # Client-role flows @@ -316,7 +467,7 @@ class TestClientFlows: lookup = conn._sent[-1] assert lookup[TYPE_KEY] == FetchMsg.TYPE assert lookup[FetchMsg.KV_REQUEST_ID] == "req-1" - assert lookup[FetchMsg.BLOCK_HASHES] == [b"k1", b"k2"] + assert lookup[FetchMsg.KEYS] == [b"k1", b"k2"] assert lookup[FetchMsg.BLOCK_INDEXES] == [0, 1] def test_transfer_done_success(self): @@ -362,13 +513,46 @@ class TestClientFlows: assert abort[TYPE_KEY] == AbortFetchMsg.TYPE assert abort[AbortFetchMsg.KV_REQUEST_ID] == "req-1" + def test_active_loads_work_list_tracks_in_flight(self): + """collect_results / has_active_loads use the _active_loads work-list, + armed when a fetch is issued and discarded exactly when its load + clears — a probe-only request never enters it, and completion empties + it while the entry may briefly linger for GC.""" + session, conn, _ = _make_session() + _activate(session, conn) + client = session._client + + # A probe-only request has no in-flight load: not in _active_loads. + session.register_lookup("req-probe", b"hp") + assert client._active_loads == set() + assert client.has_active_loads is False + + # Issuing a fetch arms the work-list. + session.request_blocks( + job_id=1, kv_request_id="req-1", keys=[b"k"], block_ids=[0] + ) + assert client._active_loads == {"req-1"} + assert client.has_active_loads is True + + # Completion clears the load and discards it from the work-list. + conn.enqueue( + { + TYPE_KEY: TransferDoneMsg.TYPE, + TransferDoneMsg.KV_REQUEST_ID: "req-1", + TransferDoneMsg.SUCCESS: True, + } + ) + session.poll() + assert client._active_loads == set() + assert client.has_active_loads is False + def test_load_timeout_sends_abort(self): session, conn, _ = _make_session() _activate(session, conn) session.request_blocks( job_id=1, kv_request_id="req-1", keys=[b"k"], block_ids=[0] ) - session._client._inbound["req-1"].submitted_at = time.monotonic() - 60.0 + session._client._requests["req-1"].load.submitted_at = time.monotonic() - 60.0 session.poll() abort = conn._sent[-1] assert abort[TYPE_KEY] == AbortFetchMsg.TYPE @@ -376,7 +560,7 @@ class TestClientFlows: def test_load_abort_ack_timeout_surfaces_failure(self): """After load timeout sends AbortFetch, if no AbortAck arrives within _ABORT_ACK_TIMEOUT_S the request is surfaced as failed and removed - from _inbound — the engine cannot wait forever on a peer that won't + from _requests — the engine cannot wait forever on a peer that won't ack. """ session, conn, _ = _make_session() @@ -385,7 +569,7 @@ class TestClientFlows: job_id=7, kv_request_id="req-7", keys=[b"k"], block_ids=[0] ) # 1) Trip the load timeout to send AbortFetch and stamp aborted_at. - session._client._inbound["req-7"].submitted_at = ( + session._client._requests["req-7"].load.submitted_at = ( time.monotonic() - _LOAD_TIMEOUT_S - 1.0 ) loads = session.poll().loads @@ -395,32 +579,32 @@ class TestClientFlows: and m[AbortFetchMsg.KV_REQUEST_ID] == "req-7" for m in conn._sent ) - assert session._client._inbound["req-7"].aborted_at is not None + assert session._client._requests["req-7"].load.aborted_at is not None # 2) Now backdate aborted_at past the abort-ack timeout. No ack ever # arrived from the peer. - session._client._inbound["req-7"].aborted_at = ( + session._client._requests["req-7"].load.aborted_at = ( time.monotonic() - _ABORT_ACK_TIMEOUT_S - 1.0 ) loads = session.poll().loads assert loads == [LoadResult(job_id=7, kv_request_id="req-7", success=False)] - assert "req-7" not in session._client._inbound + assert "req-7" not in session._client._requests def test_load_abort_ack_clears_request(self): """After load timeout sends AbortFetch, an arriving AbortAckMsg from the peer surfaces the failure cleanly and removes the request from - _inbound — covers the on_abort_ack arrival path.""" + _requests — covers the on_abort_ack arrival path.""" session, conn, _ = _make_session() _activate(session, conn) session.request_blocks( job_id=8, kv_request_id="req-8", keys=[b"k"], block_ids=[0] ) - session._client._inbound["req-8"].submitted_at = ( + session._client._requests["req-8"].load.submitted_at = ( time.monotonic() - _LOAD_TIMEOUT_S - 1.0 ) # First poll: AbortFetch goes out. session.poll() - assert session._client._inbound["req-8"].aborted_at is not None + assert session._client._requests["req-8"].load.aborted_at is not None # Peer acks the abort. conn.enqueue( @@ -431,7 +615,652 @@ class TestClientFlows: ) loads = session.poll().loads assert loads == [LoadResult(job_id=8, kv_request_id="req-8", success=False)] - assert "req-8" not in session._client._inbound + assert "req-8" not in session._client._requests + + +# --------------------------------------------------------------------------- +# Symmetric-P2P lookup flow (do_p2p_fetch) +# --------------------------------------------------------------------------- + + +class TestLookupFlow: + """Consumer-side state machine for do_p2p_fetch lookups.""" + + def test_aggregate_flush_resolve_round_trip(self): + """register_lookup → flush sends one LookupMsg → response + resolves entries → register_lookup returns the cached bool on + every call, and repeat probes never re-issue a LookupMsg.""" + session, conn, _ = _make_session() + _activate(session, conn) + + # Aggregate two keys for the same kv_request_id; both return None. + assert session.register_lookup("req-1", b"hA") is None + assert session.register_lookup("req-1", b"hB") is None + + # Flush sends one LookupMsg with both keys. + sent_before = len(conn._sent) + session.flush_pending_lookups() + new = conn._sent[sent_before:] + assert len(new) == 1 + msg = new[0] + assert msg[TYPE_KEY] == LookupMsg.TYPE + assert msg[LookupMsg.KV_REQUEST_ID] == "req-1" + assert sorted(msg[LookupMsg.KEYS]) == [b"hA", b"hB"] + + # Idempotent re-flush sends nothing — the entries are now in-flight. + sent_before = len(conn._sent) + session.flush_pending_lookups() + assert conn._sent[sent_before:] == [] + + # While in-flight, register_lookup keeps returning None. + assert session.register_lookup("req-1", b"hA") is None + + # Peer answers: hA hit, hB miss. + conn.enqueue( + { + TYPE_KEY: LookupRespMsg.TYPE, + LookupRespMsg.KV_REQUEST_ID: "req-1", + LookupRespMsg.KEYS: [b"hA", b"hB"], + LookupRespMsg.HITS: [True, False], + } + ) + session.poll() + + # register_lookup returns the resolved bool. + assert session.register_lookup("req-1", b"hA") is True + assert session.register_lookup("req-1", b"hB") is False + # The entry is cached, not popped: repeat probes keep returning the + # same result and never re-queue the key, so a flush sends nothing. + assert session.register_lookup("req-1", b"hA") is True + assert session.register_lookup("req-1", b"hB") is False + sent_before = len(conn._sent) + session.flush_pending_lookups() + assert [ + m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupMsg.TYPE + ] == [] + + def test_request_blocks_clears_probe_cache(self): + """A resolved HIT probe is popped when its fetch is issued, so a + re-scheduled request re-probes instead of trusting the stale True + (the served block is unpinned and may have been evicted).""" + session, conn, _ = _make_session() + _activate(session, conn) + + # Probe hA, flush, and let the peer resolve it to a HIT. + assert session.register_lookup("req-1", b"hA") is None + session.flush_pending_lookups() + conn.enqueue( + { + TYPE_KEY: LookupRespMsg.TYPE, + LookupRespMsg.KV_REQUEST_ID: "req-1", + LookupRespMsg.KEYS: [b"hA"], + LookupRespMsg.HITS: [True], + } + ) + session.poll() + assert session.register_lookup("req-1", b"hA") is True + + # Fetch consumes the probe. + session.request_blocks( + job_id=1, kv_request_id="req-1", keys=[b"hA"], block_ids=[0] + ) + assert b"hA" not in session._client._requests["req-1"].probes + + # Re-scheduled probe of the same key is treated as brand-new: it + # returns None and re-queues, so a flush emits a fresh LookupMsg. + assert session.register_lookup("req-1", b"hA") is None + sent_before = len(conn._sent) + session.flush_pending_lookups() + fresh = [m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupMsg.TYPE] + assert len(fresh) == 1 + assert fresh[0][LookupMsg.KEYS] == [b"hA"] + + def test_flush_uses_work_list_not_full_scan(self): + """flush drains a work-list rather than scanning every live request. + + A request with no newly-registered keys is not revisited: after a + flush the work-list is empty, an idle re-flush sends nothing, and a + subsequent register re-arms exactly the one affected id — even while + an unrelated request stays live in ``_requests``. + """ + session, conn, _ = _make_session() + _activate(session, conn) + client = session._client + + # Two requests register keys; both are queued for flush. + session.register_lookup("req-A", b"hA") + session.register_lookup("req-B", b"hB") + assert client._flush_pending == {"req-A", "req-B"} + + # Flush drains the work-list even though both requests stay live. + session.flush_pending_lookups() + assert client._flush_pending == set() + assert set(client._requests) == {"req-A", "req-B"} + + # An idle re-flush visits nothing and sends no LookupMsg. + sent_before = len(conn._sent) + session.flush_pending_lookups() + assert [ + m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupMsg.TYPE + ] == [] + + # A new register re-arms only that id. + session.register_lookup("req-A", b"hA2") + assert client._flush_pending == {"req-A"} + + def test_separate_lookup_msg_per_kv_request_id(self): + """Hashes for different kv_request_ids flush as separate LookupMsgs.""" + session, conn, _ = _make_session() + _activate(session, conn) + + session.register_lookup("req-A", b"h1") + session.register_lookup("req-B", b"h2") + session.register_lookup("req-A", b"h3") + + sent_before = len(conn._sent) + session.flush_pending_lookups() + sent = [m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupMsg.TYPE] + assert len(sent) == 2 + by_req = {m[LookupMsg.KV_REQUEST_ID]: m[LookupMsg.KEYS] for m in sent} + assert sorted(by_req["req-A"]) == [b"h1", b"h3"] + assert by_req["req-B"] == [b"h2"] + + def test_multiple_lookup_msgs_across_steps(self): + """A request's block set may be discovered across scheduler steps: + each step that registers new keys flushes its own LookupMsg for + the same kv_request_id, carrying only the newly-probed keys.""" + session, conn, _ = _make_session() + _activate(session, conn) + + # Step 1: probe hA, hB. + session.register_lookup("req-1", b"hA") + session.register_lookup("req-1", b"hB") + sent_before = len(conn._sent) + session.flush_pending_lookups() + first = [m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupMsg.TYPE] + assert len(first) == 1 + assert first[0][LookupMsg.KV_REQUEST_ID] == "req-1" + assert sorted(first[0][LookupMsg.KEYS]) == [b"hA", b"hB"] + + # Step 2: a new key is discovered for the same request. The + # in-flight keys from step 1 are not re-sent; a second LookupMsg + # goes out carrying only the newly-probed key. + assert session.register_lookup("req-1", b"hA") is None # in-flight no-op + session.register_lookup("req-1", b"hC") + sent_before = len(conn._sent) + session.flush_pending_lookups() + second = [m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupMsg.TYPE] + assert len(second) == 1 + assert second[0][LookupMsg.KV_REQUEST_ID] == "req-1" + assert second[0][LookupMsg.KEYS] == [b"hC"] + + def test_split_response_resolves_across_messages(self): + """Producer may answer one LookupMsg's keys across multiple + LookupRespMsgs — pairs are self-describing so each lands.""" + session, conn, _ = _make_session() + _activate(session, conn) + + session.register_lookup("req-1", b"hA") + session.register_lookup("req-1", b"hB") + session.flush_pending_lookups() + + # Two responses, each carrying one of the two keys. + conn.enqueue( + { + TYPE_KEY: LookupRespMsg.TYPE, + LookupRespMsg.KV_REQUEST_ID: "req-1", + LookupRespMsg.KEYS: [b"hA"], + LookupRespMsg.HITS: [True], + } + ) + conn.enqueue( + { + TYPE_KEY: LookupRespMsg.TYPE, + LookupRespMsg.KV_REQUEST_ID: "req-1", + LookupRespMsg.KEYS: [b"hB"], + LookupRespMsg.HITS: [False], + } + ) + session.poll() + + assert session.register_lookup("req-1", b"hA") is True + assert session.register_lookup("req-1", b"hB") is False + + def test_finish_request_cancels_pending_lookups(self): + """finish_request drops every pending lookup for the kv_request_id.""" + session, conn, _ = _make_session() + _activate(session, conn) + + session.register_lookup("req-1", b"hA") + session.register_lookup("req-1", b"hB") + session.register_lookup("req-2", b"hC") + session.finish_request("req-1") + + # req-1 entries gone, req-2 untouched. + assert "req-1" not in session._client._requests + assert b"hC" in session._client._requests["req-2"].probes + + def test_finish_after_flushed_lookup_sends_empty_fetch(self): + """LookupMsg flushed but no FetchMsg sent (all-miss case) → + finish_request emits an empty FetchMsg so the peer can drop its + lookup state and call parent.on_request_finished.""" + session, conn, _ = _make_session() + _activate(session, conn) + + session.register_lookup("req-1", b"hA") + session.flush_pending_lookups() + sent_before = len(conn._sent) + + session.finish_request("req-1") + + fetches = [m for m in conn._sent[sent_before:] if m[TYPE_KEY] == FetchMsg.TYPE] + assert len(fetches) == 1 + assert fetches[0][FetchMsg.KV_REQUEST_ID] == "req-1" + assert fetches[0][FetchMsg.KEYS] == [] + assert fetches[0][FetchMsg.BLOCK_INDEXES] == [] + + def test_finish_without_flushed_lookup_sends_no_fetch(self): + """No LookupMsg was ever sent → finish_request must not emit an + empty FetchMsg (the peer has no state to release).""" + session, conn, _ = _make_session() + _activate(session, conn) + + # Register but never flush. + session.register_lookup("req-1", b"hA") + sent_before = len(conn._sent) + + session.finish_request("req-1") + + fetches = [m for m in conn._sent[sent_before:] if m[TYPE_KEY] == FetchMsg.TYPE] + assert fetches == [] + + def test_finish_after_real_fetch_sends_no_second_fetch(self): + """A real FetchMsg was already sent for the id → finish_request + must not emit a second (empty) FetchMsg.""" + session, conn, _ = _make_session() + _activate(session, conn) + + session.register_lookup("req-1", b"hA") + session.flush_pending_lookups() + # Resolve the probe to a HIT before fetching, as the manager only + # loads confirmed hits (an unresolved probe yields RETRY). + conn.enqueue( + { + TYPE_KEY: LookupRespMsg.TYPE, + LookupRespMsg.KV_REQUEST_ID: "req-1", + LookupRespMsg.KEYS: [b"hA"], + LookupRespMsg.HITS: [True], + } + ) + session.poll() + session.request_blocks( + job_id=1, kv_request_id="req-1", keys=[b"hA"], block_ids=[7] + ) + sent_before = len(conn._sent) + + session.finish_request("req-1") + + fetches = [m for m in conn._sent[sent_before:] if m[TYPE_KEY] == FetchMsg.TYPE] + assert fetches == [] + + def test_server_lookup_deferred_until_serve_then_all_misses(self): + """``poll()`` only enqueues an inbound LookupMsg — no response is + sent until ``serve_external_requests``. With an all-miss parent + the aggregated LookupRespMsg carries the same keys and + ``hits=[False, ...]``.""" + session, conn, _ = _make_session() + _activate(session, conn) + + sent_before = len(conn._sent) + conn.enqueue( + { + TYPE_KEY: LookupMsg.TYPE, + LookupMsg.KV_REQUEST_ID: "req-1", + LookupMsg.KEYS: [b"hX", b"hY", b"hZ"], + } + ) + session.poll() + + # Dispatch alone must not answer — the parent handle is only valid + # during serve_external_requests. + assert [ + m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupRespMsg.TYPE + ] == [] + + _serve(session, FakeParent()) + + resps = [ + m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupRespMsg.TYPE + ] + assert len(resps) == 1 + resp = resps[0] + assert resp[LookupRespMsg.KV_REQUEST_ID] == "req-1" + assert resp[LookupRespMsg.KEYS] == [b"hX", b"hY", b"hZ"] + assert resp[LookupRespMsg.HITS] == [False, False, False] + + +# --------------------------------------------------------------------------- +# Server-side handling of inbound LookupMsg (ParentManager-driven) +# +# poll() only enqueues the LookupMsg; serve_external_requests(parent) +# resolves it. Tests follow the poll() → _serve() pattern. +# --------------------------------------------------------------------------- + + +def _send_lookup(conn: FakeConnection, kv_request_id: str, keys: list[bytes]): + conn.enqueue( + { + TYPE_KEY: LookupMsg.TYPE, + LookupMsg.KV_REQUEST_ID: kv_request_id, + LookupMsg.KEYS: list(keys), + } + ) + + +def _lookup_resps(conn: FakeConnection, since: int = 0) -> list[dict]: + return [m for m in conn._sent[since:] if m[TYPE_KEY] == LookupRespMsg.TYPE] + + +class TestServerLookupHandling: + def test_immediate_hits_create_one_store_job(self): + """All-HIT batch: one create_store_job call with all keys, one + LookupRespMsg with hits=[True]*N, on_request_finished fires at the + end of serve, and `available` is populated for the eventual fetch.""" + cb = FakeParent(stored={b"hA": 1, b"hB": 2, b"hC": 3}) + session, conn, _ = _make_session() + _activate(session, conn) + + sent_before = len(conn._sent) + _send_lookup(conn, "req-1", [b"hA", b"hB", b"hC"]) + session.poll() + _serve(session, cb) + + resps = _lookup_resps(conn, sent_before) + assert len(resps) == 1 + assert resps[0][LookupRespMsg.KEYS] == [b"hA", b"hB", b"hC"] + assert resps[0][LookupRespMsg.HITS] == [True, True, True] + + kinds = [c[0] for c in cb.calls] + assert kinds.count("create_store_job") == 1 + cs = next(c for c in cb.calls if c[0] == "create_store_job") + assert cs[1] == (b"hA", b"hB", b"hC") + assert cb.calls[-1][0] == "on_request_finished" + + # Hits are pinned in outbound state for the upcoming FetchMsg match. + assert set(_srv_outbound(session, "req-1").available) == { + b"hA", + b"hB", + b"hC", + } + + def test_all_misses_no_store_job_finish_fires(self): + """All-MISS batch: no create_store_job call; one LookupRespMsg + with hits=[False]*N; on_request_finished fires at end of serve.""" + cb = FakeParent() + session, conn, _ = _make_session() + _activate(session, conn) + + sent_before = len(conn._sent) + _send_lookup(conn, "req-1", [b"hA", b"hB"]) + session.poll() + _serve(session, cb) + + resps = _lookup_resps(conn, sent_before) + assert len(resps) == 1 + assert resps[0][LookupRespMsg.HITS] == [False, False] + assert all(c[0] != "create_store_job" for c in cb.calls) + assert cb.calls[-1][0] == "on_request_finished" + + def test_mixed_hit_miss_pending_defers_response_until_aggregate(self): + """HIT/MISS resolutions do not go out on first sight when any + key is still HIT_PENDING / RETRY. The lookup parks until every + key has settled (or the deadline fires), then one + LookupRespMsg carries all keys in wire order.""" + cb = FakeParent( + stored={b"hA": 1}, + pending={b"hB"}, + retry={b"hD"}, + ) + session, conn, _ = _make_session() + _activate(session, conn) + + sent_before = len(conn._sent) + _send_lookup(conn, "req-1", [b"hA", b"hB", b"hC", b"hD"]) + session.poll() + _serve(session, cb) + + # No LookupRespMsg yet — hB and hD are still pending. + assert _lookup_resps(conn, sent_before) == [] + # HIT is still pinned immediately so the eventual FetchMsg matches. + cs_calls = [c for c in cb.calls if c[0] == "create_store_job"] + assert len(cs_calls) == 1 + assert cs_calls[0][1] == (b"hA",) + # Lookup is parked; on_request_finished not yet called. + assert all(c[0] != "on_request_finished" for c in cb.calls) + assert len(_srv_lookups(session)) == 1 + + def test_pending_resolves_then_aggregate_response_fires(self): + """A HIT_PENDING key that becomes HIT on a later poll releases + the deferred aggregate response: one LookupRespMsg carrying + both keys in wire order, and one create_store_job call per + HIT (the second HIT is pinned when it resolves, not when the + response goes out).""" + cb = FakeParent(stored={b"hA": 1}, pending={b"hB"}) + session, conn, _ = _make_session() + _activate(session, conn) + + sent_before = len(conn._sent) + _send_lookup(conn, "req-1", [b"hA", b"hB"]) + session.poll() + _serve(session, cb) + # No response yet — hB still pending. + assert _lookup_resps(conn, sent_before) == [] + + # Promote hB. + cb.pending.discard(b"hB") + cb.stored[b"hB"] = 2 + + # Drive resolver via a second serve_external_requests. + _serve(session, cb) + + resps = _lookup_resps(conn, sent_before) + assert len(resps) == 1 + assert resps[0][LookupRespMsg.KEYS] == [b"hA", b"hB"] + assert resps[0][LookupRespMsg.HITS] == [True, True] + + cs_calls = [c for c in cb.calls if c[0] == "create_store_job"] + assert len(cs_calls) == 2 + assert cs_calls[0][1] == (b"hA",) + assert cs_calls[1][1] == (b"hB",) + # on_request_finished fires once after the aggregate resolve. + assert sum(1 for c in cb.calls if c[0] == "on_request_finished") == 1 + assert b"hA" in _srv_outbound(session, "req-1").available + assert b"hB" in _srv_outbound(session, "req-1").available + + def test_pending_timeout_replies_miss_no_store_job(self): + """A HIT_PENDING key that stays pending past the batch + ``deadline`` is force-MISS and never pinned; the deferred + aggregate response fires with hits=[False].""" + cb = FakeParent(pending={b"hA"}) + session, conn, _ = _make_session() + _activate(session, conn) + + sent_before = len(conn._sent) + _send_lookup(conn, "req-1", [b"hA"]) + session.poll() + _serve(session, cb) + # Initial serve: nothing immediate, lookup parked, no LookupRespMsg. + assert _lookup_resps(conn, sent_before) == [] + + # Forge the deadline into the past to trigger the timeout branch. + lookup = _srv_lookups(session)[0] + lookup.deadline = time.monotonic() - 0.1 + + _serve(session, cb) + + resps = _lookup_resps(conn, sent_before) + assert len(resps) == 1 + assert resps[0][LookupRespMsg.KEYS] == [b"hA"] + assert resps[0][LookupRespMsg.HITS] == [False] + assert all(c[0] != "create_store_job" for c in cb.calls) + assert sum(1 for c in cb.calls if c[0] == "on_request_finished") == 1 + + def test_finish_request_called_per_lookup_msg_not_per_kv_request_id(self): + """Two LookupMsgs for the same kv_request_id get distinct ctxs + and two on_request_finished calls (one per batch).""" + cb = FakeParent(stored={b"hA": 1, b"hB": 2}) + session, conn, _ = _make_session() + _activate(session, conn) + + _send_lookup(conn, "req-1", [b"hA"]) + session.poll() + _serve(session, cb) + _send_lookup(conn, "req-1", [b"hB"]) + session.poll() + _serve(session, cb) + + finish_calls = [c for c in cb.calls if c[0] == "on_request_finished"] + assert len(finish_calls) == 2 + # Distinct synthetic req_ids + assert finish_calls[0][1] != finish_calls[1][1] + # Both namespaced under the same kv_request_id + assert ":req-1:" in finish_calls[0][1] + assert ":req-1:" in finish_calls[1][1] + + def test_close_returns_open_batch_ctxs_as_failed_serves(self): + """Tearing the session down with a parked batch returns the + synthetic ctx as a failed serve (no parent handle at teardown) so + the manager can release the TieringManager's state on its next + serve.""" + cb = FakeParent(pending={b"hA"}) + session, conn, _ = _make_session() + _activate(session, conn) + + _send_lookup(conn, "req-1", [b"hA"]) + session.poll() + _serve(session, cb) + assert len(_srv_lookups(session)) == 1 + assert all(c[0] != "on_request_finished" for c in cb.calls) + + result = session.close() + + assert len(result.failed_serves) == 1 + assert ":req-1:" in result.failed_serves[0].req_id + # close() itself must not call the parent. + assert all(c[0] != "on_request_finished" for c in cb.calls) + + def test_wire_finish_drops_pending_batches_for_kv_request_id(self): + """``ServerRole.finish(kv_request_id)`` drops every parked batch + whose kv_request_id matches and queues its ctx for the next + serve's on_request_finished.""" + cb = FakeParent(pending={b"hA", b"hB"}) + session, conn, _ = _make_session() + _activate(session, conn) + + _send_lookup(conn, "req-1", [b"hA"]) + session.poll() + _serve(session, cb) + _send_lookup(conn, "req-2", [b"hB"]) + session.poll() + _serve(session, cb) + assert len(_srv_lookups(session)) == 2 + + session._server.finish("req-1") + + # req-1 batch dropped from parked lookups; its ctx queued for release. + remaining_kv_request_ids = {b.kv_request_id for b in _srv_lookups(session)} + assert remaining_kv_request_ids == {"req-2"} + queued = session._server._finished_lookup_ctxs + assert len(queued) == 1 + assert ":req-1:" in queued[0].req_id + + # The next serve fires on_request_finished exactly once for req-1. + _serve(session, cb) + finish_calls = [c for c in cb.calls if c[0] == "on_request_finished"] + assert len(finish_calls) == 1 + assert ":req-1:" in finish_calls[0][1] + + def test_incoming_fetch_drops_pending_lookups_for_kv_request_id(self): + """A peer FetchMsg terminates the lookup phase for its id: parked + lookups with matching kv_request_id are dropped and their + ctx queued for on_request_finished; other kv_request_ids untouched.""" + cb = FakeParent(pending={b"hA", b"hB"}) + session, conn, _ = _make_session() + _activate(session, conn) + + _send_lookup(conn, "req-1", [b"hA"]) + session.poll() + _serve(session, cb) + _send_lookup(conn, "req-2", [b"hB"]) + session.poll() + _serve(session, cb) + assert len(_srv_lookups(session)) == 2 + + # Empty FetchMsg: peer signals "lookup phase done" without asking + # for any blocks (the all-miss case). + conn.enqueue( + { + TYPE_KEY: FetchMsg.TYPE, + FetchMsg.KV_REQUEST_ID: "req-1", + FetchMsg.KEYS: [], + FetchMsg.BLOCK_INDEXES: [], + } + ) + session.poll() + + remaining_kv_request_ids = {lu.kv_request_id for lu in _srv_lookups(session)} + assert remaining_kv_request_ids == {"req-2"} + # Dispatch queues the ctx but does not call the parent yet. + queued = session._server._finished_lookup_ctxs + assert len(queued) == 1 + assert ":req-1:" in queued[0].req_id + + # The next serve fires on_request_finished exactly once for req-1. + _serve(session, cb) + finish_calls = [c for c in cb.calls if c[0] == "on_request_finished"] + assert len(finish_calls) == 1 + assert ":req-1:" in finish_calls[0][1] + + def test_lookup_then_fetch_round_trip_emits_store_result(self): + """End-to-end: lookup pins primary slots → fetch matches them → + NIXL transfer completes → StoreResult surfaces with the + create_store_job's job_id (the engine releases the pin).""" + cb = FakeParent(stored={b"hA": 7, b"hB": 8}) + session, conn, transport = _make_session() + _activate(session, conn) + + _send_lookup(conn, "req-1", [b"hA", b"hB"]) + session.poll() + _serve(session, cb) + cs = next(c for c in cb.calls if c[0] == "create_store_job") + # FakeParent issues monotonic job_ids starting at 1000. + expected_job_id = 1000 + + # Consumer issues FetchMsg on the resolved hits. + conn.enqueue( + { + TYPE_KEY: FetchMsg.TYPE, + FetchMsg.KV_REQUEST_ID: "req-1", + FetchMsg.KEYS: [b"hA", b"hB"], + FetchMsg.BLOCK_INDEXES: [20, 21], + } + ) + session.poll() + + # NIXL write_blocks called with our pinned local block_ids. + assert len(transport._transfers) == 1 + _, (_peer, local, remote) = next(iter(transport._transfers.items())) + assert local == [7, 8] + assert remote == [20, 21] + + # Drive the transport completion. + transport._poll_done.append(0) + result = session.poll() + + store_results = [s for s in result.stores if s.success] + assert any(s.job_id == expected_job_id for s in store_results) + # Sanity: kv mention in synthetic ctx. + assert cs[2].startswith("p2p:") # --------------------------------------------------------------------------- @@ -449,7 +1278,7 @@ class TestServerFlows: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1", b"k2"], + FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [10, 11], } ) @@ -467,7 +1296,7 @@ class TestServerFlows: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1"], + FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [5], } ) @@ -485,7 +1314,7 @@ class TestServerFlows: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1"], + FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [5], } ) @@ -508,11 +1337,11 @@ class TestServerFlows: session.poll() ack = next(m for m in conn._sent if m[TYPE_KEY] == AbortAckMsg.TYPE) assert ack[AbortAckMsg.KV_REQUEST_ID] == "req-1" - assert "req-1" not in session._server._pending_aborts + assert _srv_abort_started(session, "req-1") is None def test_abort_fetch_defers_ack_when_cancel_pending(self): """If cancel(mode='wait') reports still-inflight tids, the ack is - deferred and the abort is parked in _pending_aborts.""" + deferred and the abort is parked (abort_started_at set).""" session, conn, transport = _make_session() _activate(session, conn) # Seed an inflight transfer for req-1 that the transport pretends @@ -533,7 +1362,7 @@ class TestServerFlows: session.poll() assert not any(m[TYPE_KEY] == AbortAckMsg.TYPE for m in conn._sent) - assert "req-1" in session._server._pending_aborts + assert _srv_abort_started(session, "req-1") is not None # First attempt happens inside _on_abort_fetch; the per-tick # drain runs again at the end of poll() — both are wait-mode. assert all(mode == "wait" for _, mode in transport._cancel_calls) @@ -558,7 +1387,7 @@ class TestServerFlows: } ) session.poll() - assert "req-1" in session._server._pending_aborts + assert _srv_abort_started(session, "req-1") is not None # Backend finishes draining: transport.poll() will return tid as # DONE, and the next cancel(mode='wait') call sees it's gone. @@ -569,7 +1398,7 @@ class TestServerFlows: ack = next(m for m in conn._sent if m[TYPE_KEY] == AbortAckMsg.TYPE) assert ack[AbortAckMsg.KV_REQUEST_ID] == "req-1" - assert "req-1" not in session._server._pending_aborts + assert _srv_abort_started(session, "req-1") is None assert tid not in session._server._inflight def test_abort_fetch_force_cancels_after_timeout(self): @@ -591,9 +1420,9 @@ class TestServerFlows: } ) session.poll() - assert "req-1" in session._server._pending_aborts + assert _srv_abort_started(session, "req-1") is not None # Backdate past the drain deadline. - session._server._pending_aborts["req-1"] = ( + session._server._requests["req-1"].abort_started_at = ( time.monotonic() - _CANCEL_DRAIN_TIMEOUT_S - 1.0 ) # Even if the transport still claims it can't cancel, the @@ -604,7 +1433,7 @@ class TestServerFlows: ack = next(m for m in conn._sent if m[TYPE_KEY] == AbortAckMsg.TYPE) assert ack[AbortAckMsg.KV_REQUEST_ID] == "req-1" - assert "req-1" not in session._server._pending_aborts + assert _srv_abort_started(session, "req-1") is None assert tid not in session._server._inflight assert ([tid], "immediate") in transport._cancel_calls @@ -627,7 +1456,7 @@ class TestServerFlows: } ) session.poll() - first_started_at = session._server._pending_aborts["req-1"] + first_started_at = _srv_abort_started(session, "req-1") # Second AbortFetchMsg for the same kv_request_id while still # draining must not reset the deadline. @@ -638,7 +1467,7 @@ class TestServerFlows: } ) session.poll() - assert session._server._pending_aborts["req-1"] == first_started_at + assert _srv_abort_started(session, "req-1") == first_started_at # Now let the drain succeed and confirm exactly one ack ever. transport._cancel_still_inflight.discard(tid) @@ -669,7 +1498,7 @@ class TestServerFlows: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1"], + FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [5], } ) @@ -701,7 +1530,7 @@ class TestServerFlows: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1"], + FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [5], } ) @@ -743,12 +1572,12 @@ class TestFinishRequestServerSide: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1"], + FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [5], } ) session.poll() - assert "req-1" in session._server._outbound + assert _srv_outbound(session, "req-1") is not None session.finish_request("req-1") @@ -756,7 +1585,7 @@ class TestFinishRequestServerSide: assert msg is not None assert msg[TransferDoneMsg.KV_REQUEST_ID] == "req-1" assert msg[TransferDoneMsg.SUCCESS] is False - assert "req-1" not in session._server._outbound + assert _srv_outbound(session, "req-1") is None def test_with_inflight_defers_then_fires_on_last_transfer(self): """finish_request with inflight defers; last transfer fires the @@ -768,7 +1597,7 @@ class TestFinishRequestServerSide: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1", b"k2"], + FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [10, 11], } ) @@ -780,8 +1609,8 @@ class TestFinishRequestServerSide: before = len(conn._sent) session.finish_request("req-1") assert len(conn._sent) == before - assert "req-1" in session._server._outbound - assert session._server._outbound["req-1"].finishing + assert _srv_outbound(session, "req-1") is not None + assert _srv_outbound(session, "req-1").finishing # Last inflight settles -> early-fail fires. tid = next(iter(transport._transfers)) @@ -792,7 +1621,7 @@ class TestFinishRequestServerSide: assert msg is not None assert msg[TransferDoneMsg.KV_REQUEST_ID] == "req-1" assert msg[TransferDoneMsg.SUCCESS] is False - assert "req-1" not in session._server._outbound + assert _srv_outbound(session, "req-1") is None def test_full_demand_satisfied_still_sends_success(self): """finish_request must not override a fully-satisfied transfer: @@ -803,7 +1632,7 @@ class TestFinishRequestServerSide: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1"], + FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [10], } ) @@ -831,13 +1660,13 @@ class TestFinishRequestServerSide: session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1) # finish_request first — no demand received yet -> defer. session.finish_request("req-1") - assert "req-1" in session._server._outbound + assert _srv_outbound(session, "req-1") is not None # Fetch arrives now: demand fully satisfied by available. conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1"], + FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [10], } ) @@ -858,7 +1687,7 @@ class TestFinishRequestServerSide: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-2", - FetchMsg.BLOCK_HASHES: [b"k1", b"k2"], + FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [10, 11], } ) @@ -870,7 +1699,7 @@ class TestFinishRequestServerSide: session.poll() msg = next(m for m in conn._sent if m[TYPE_KEY] == TransferDoneMsg.TYPE) assert msg[TransferDoneMsg.SUCCESS] is False - assert "req-2" not in session._server._outbound + assert _srv_outbound(session, "req-2") is None def test_unknown_request_is_noop(self): session, conn, _ = _make_session() @@ -891,7 +1720,7 @@ class TestFinishRequestServerSide: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"demand"], + FetchMsg.KEYS: [b"demand"], FetchMsg.BLOCK_INDEXES: [5], } ) @@ -900,7 +1729,7 @@ class TestFinishRequestServerSide: # matches demand. Without the shortcut, job 42 sits in _store_jobs # for _STORE_TIMEOUT_S. session.add_stored_blocks("req-1", [b"unrelated"], [0], job_id=42) - assert session._server._outbound["req-1"].pending_job_ids == {42} + assert _srv_outbound(session, "req-1").pending_job_ids == {42} session.finish_request("req-1") @@ -908,7 +1737,7 @@ class TestFinishRequestServerSide: msg = self._last_transfer_done(conn) assert msg is not None assert msg[TransferDoneMsg.SUCCESS] is False - assert "req-1" not in session._server._outbound + assert _srv_outbound(session, "req-1") is None # Local store job surfaces on the next poll, success=False. stores = session.poll().stores @@ -925,7 +1754,7 @@ class TestFinishRequestServerSide: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1"], + FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [10], } ) @@ -933,7 +1762,7 @@ class TestFinishRequestServerSide: session.add_stored_blocks("req-1", [b"k1"], [0], job_id=7) # finish_request races with the inflight transfer. session.finish_request("req-1") - assert "req-1" in session._server._outbound # deferred + assert _srv_outbound(session, "req-1") is not None # deferred # Last inflight completes -> _finalize_outbound(success=True) fires. tid = next(iter(transport._transfers)) @@ -944,7 +1773,7 @@ class TestFinishRequestServerSide: assert msg is not None assert msg[TransferDoneMsg.SUCCESS] is True assert StoreResult(job_id=7, success=True) in stores - assert "req-1" not in session._server._outbound + assert _srv_outbound(session, "req-1") is None def test_write_blocks_failure_finalizes_with_failure(self): """write_blocks returning None must not leave the request hanging. @@ -962,7 +1791,7 @@ class TestFinishRequestServerSide: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1"], + FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [10], } ) @@ -973,7 +1802,7 @@ class TestFinishRequestServerSide: session.add_stored_blocks("req-1", [b"k1"], [0], job_id=42) # Outbound was finalized immediately (no other inflight). - assert "req-1" not in session._server._outbound + assert _srv_outbound(session, "req-1") is None # Peer notified with success=False. msg = next(m for m in conn._sent if m[TYPE_KEY] == TransferDoneMsg.TYPE) assert msg[TransferDoneMsg.KV_REQUEST_ID] == "req-1" @@ -996,14 +1825,14 @@ class TestFinishRequestServerSide: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1", b"k2", b"k3"], + FetchMsg.KEYS: [b"k1", b"k2", b"k3"], FetchMsg.BLOCK_INDEXES: [10, 11, 12], } ) session.poll() # Demand registered, no matches yet. assert session._server._inflight == {} - outbound = session._server._outbound["req-1"] + outbound = _srv_outbound(session, "req-1") assert outbound.remaining == 3 assert set(outbound.demanded.keys()) == {b"k1", b"k2", b"k3"} @@ -1024,7 +1853,7 @@ class TestFinishRequestServerSide: assert session._server._inflight == {} assert outbound.remaining == 2 # Not yet finalized — still 2 blocks demanded. - assert "req-1" in session._server._outbound + assert _srv_outbound(session, "req-1") is not None # Round 2: k2 and k3 arrive together. session.add_stored_blocks("req-1", [b"k2", b"k3"], [1, 2], job_id=200) @@ -1038,7 +1867,7 @@ class TestFinishRequestServerSide: stores = session.poll().stores assert StoreResult(job_id=200, success=True) in stores # _finalize_outbound fired — request gone, peer notified with success. - assert "req-1" not in session._server._outbound + assert _srv_outbound(session, "req-1") is None done = next(m for m in conn._sent if m.get(TYPE_KEY) == TransferDoneMsg.TYPE) assert done[TransferDoneMsg.KV_REQUEST_ID] == "req-1" assert done[TransferDoneMsg.SUCCESS] is True @@ -1057,7 +1886,7 @@ class TestFinishRequestServerSide: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1", b"k2"], + FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [10, 11], } ) @@ -1067,7 +1896,7 @@ class TestFinishRequestServerSide: session.add_stored_blocks("req-1", [b"k1"], [0], job_id=100) assert len(session._server._inflight) == 1 tid_1 = next(iter(session._server._inflight)) - outbound = session._server._outbound["req-1"] + outbound = _srv_outbound(session, "req-1") assert outbound.remaining == 2 # decrement happens on completion assert outbound.finishing is False @@ -1078,7 +1907,7 @@ class TestFinishRequestServerSide: assert list(session._server._inflight.keys()) == [tid_1] # Marked finishing, but NOT finalized yet (transfer_1 still inflight). assert outbound.finishing is True - assert "req-1" in session._server._outbound + assert _srv_outbound(session, "req-1") is not None done_msgs = [m for m in conn._sent if m.get(TYPE_KEY) == TransferDoneMsg.TYPE] assert done_msgs == [] @@ -1090,7 +1919,7 @@ class TestFinishRequestServerSide: stores_first = session.poll().stores assert StoreResult(job_id=100, success=True) in stores_first # Outbound state cleaned up; peer notified with success=False. - assert "req-1" not in session._server._outbound + assert _srv_outbound(session, "req-1") is None done = next(m for m in conn._sent if m.get(TYPE_KEY) == TransferDoneMsg.TYPE) assert done[TransferDoneMsg.KV_REQUEST_ID] == "req-1" assert done[TransferDoneMsg.SUCCESS] is False @@ -1123,7 +1952,7 @@ class TestBidirectional: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-srv", - FetchMsg.BLOCK_HASHES: [b"served"], + FetchMsg.KEYS: [b"served"], FetchMsg.BLOCK_INDEXES: [7], } ) @@ -1178,6 +2007,7 @@ class TestPendingSession: local_id="local:9000", transport=transport, # type: ignore[arg-type] local_block_len=4096, + local_hash_seed=_DEFAULT_HASH_SEED, conn=None, ) session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1) @@ -1197,6 +2027,7 @@ class TestPendingSession: local_id="local:9000", transport=transport, # type: ignore[arg-type] local_block_len=4096, + local_hash_seed=_DEFAULT_HASH_SEED, conn=None, ) conn = FakeConnection(peer_id="peer:8000") @@ -1218,13 +2049,16 @@ class TestPendingSession: local_id="local:9000", transport=transport, # type: ignore[arg-type] local_block_len=4096, + local_hash_seed=_DEFAULT_HASH_SEED, conn=None, ) session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1) session.add_stored_blocks("req-2", [b"k2"], [1], job_id=2) - failed_loads, failed_stores = session.close() - assert failed_loads == [] - assert set(failed_stores) == {1, 2} + result = session.close() + assert result.failed_jobs == [] + assert result.failed_req_ids == [] + assert set(result.failed_stores) == {1, 2} + assert result.failed_serves == [] # --------------------------------------------------------------------------- @@ -1246,9 +2080,50 @@ class TestDisconnect: session.request_blocks(1, "req-1", [b"k"], [0]) session.request_blocks(2, "req-2", [b"k"], [0]) session.add_stored_blocks("req-srv", [b"k"], [0], job_id=10) - failed_loads, failed_stores = session.close() - assert set(failed_loads) == {(1, "req-1"), (2, "req-2")} - assert set(failed_stores) == {10} + result = session.close() + assert set(result.failed_jobs) == {1, 2} + assert set(result.failed_req_ids) == {"req-1", "req-2"} + assert set(result.failed_stores) == {10} + assert result.failed_serves == [] + + def test_send_failure_marks_connection_dead(self): + """A raising send must mark the connection dead, not silently drop + the message — otherwise the session lingers alive, is never reaped, + and in-flight lookups/loads toward the dead peer hang forever.""" + session, conn, _ = _make_session() + _activate(session, conn) + assert session.alive + + conn.fail_send = True + # request_blocks flushes a FetchMsg synchronously via _do_send. + session.request_blocks(1, "req-1", [b"k"], [0]) + + assert not session.alive + + def test_close_surfaces_inflight_lookups(self): + """close() reports kv_request_ids whose symmetric-P2P probe is still + unresolved; resolved probes are not reported (their answer is in).""" + session, conn, _ = _make_session() + _activate(session, conn) + + session.register_lookup("req-hit", b"hA") + session.register_lookup("req-inflight", b"hB") + session.flush_pending_lookups() + + # Only req-hit is answered; req-inflight stays in flight. + conn.enqueue( + { + TYPE_KEY: LookupRespMsg.TYPE, + LookupRespMsg.KV_REQUEST_ID: "req-hit", + LookupRespMsg.KEYS: [b"hA"], + LookupRespMsg.HITS: [True], + } + ) + session.poll() + + result = session.close() + assert result.failed_jobs == [] + assert result.failed_req_ids == ["req-inflight"] # --------------------------------------------------------------------------- @@ -1294,7 +2169,7 @@ class TestAdversarial: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-bad", - FetchMsg.BLOCK_HASHES: [b"k1", b"k2"], + FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [1], } ) @@ -1342,7 +2217,7 @@ class TestDispatchErrorHandling: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-bad", - FetchMsg.BLOCK_HASHES: [b"k1", b"k2"], + FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [1], } ) @@ -1372,7 +2247,7 @@ class TestDispatchErrorHandling: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1"], + FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [0], } ) @@ -1395,7 +2270,7 @@ class TestDispatchErrorHandling: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1"], + FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [0], } ) @@ -1422,7 +2297,7 @@ class TestDispatchErrorHandling: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1"], + FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [0], } ) @@ -1438,10 +2313,10 @@ class TestDispatchErrorHandling: class TestInflightPerReqInvariant: - """`_inflight_per_req` is the O(1) replacement for the previous - O(N) scan in `_has_inflight_for`. These tests check that every - mutation site keeps the counter in sync with `_inflight` and that - the lookup is correct under high fan-out. + """Per-request `inflight_tids` is the O(1) replacement for the + previous O(N) scan in `_has_inflight_for`. These tests check that + every mutation site keeps the set in sync with `_inflight` and + that the lookup is correct under high fan-out. """ def test_invariant_holds_through_lifecycle(self): @@ -1452,17 +2327,14 @@ class TestInflightPerReqInvariant: _activate(session, conn) def _invariant_holds() -> bool: - counted = sum(session._server._inflight_per_req.values()) - return counted == len(session._server._inflight) and all( - v > 0 for v in session._server._inflight_per_req.values() - ) + return _srv_total_inflight(session) == len(session._server._inflight) assert _invariant_holds() # Two requests, two blocks each, all dispatched in one batch. session.add_stored_blocks("req-A", [b"a1", b"a2"], [0, 1], job_id=10) session.add_stored_blocks("req-B", [b"b1", b"b2"], [2, 3], job_id=11) - for kv_id, hashes, indexes in ( + for kv_id, keys, indexes in ( ("req-A", [b"a1", b"a2"], [100, 101]), ("req-B", [b"b1", b"b2"], [102, 103]), ): @@ -1470,7 +2342,7 @@ class TestInflightPerReqInvariant: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: kv_id, - FetchMsg.BLOCK_HASHES: hashes, + FetchMsg.KEYS: keys, FetchMsg.BLOCK_INDEXES: indexes, } ) @@ -1493,7 +2365,7 @@ class TestInflightPerReqInvariant: assert _invariant_holds() assert not session._server._has_inflight_for("req-A") - assert "req-A" not in session._server._inflight_per_req # entry was removed + assert _srv_inflight_count(session, "req-A") == 0 # entry drained assert session._server._has_inflight_for("req-B") # Complete req-B; counter must drain to empty. @@ -1508,7 +2380,7 @@ class TestInflightPerReqInvariant: assert _invariant_holds() assert session._server._inflight == {} - assert session._server._inflight_per_req == {} + assert _srv_total_inflight(session) == 0 def test_has_inflight_for_correct_with_many_requests(self): """Populate many inflight xfers across many ids; lookup must @@ -1524,9 +2396,7 @@ class TestInflightPerReqInvariant: tid, _InflightXfer(kv_request_id=kv_id, block_count=1, job_ids={tid}), ) - assert sum(session._server._inflight_per_req.values()) == len( - session._server._inflight - ) + assert _srv_total_inflight(session) == len(session._server._inflight) assert session._server._has_inflight_for("req-0") assert session._server._has_inflight_for("req-99") assert not session._server._has_inflight_for("req-missing") @@ -1539,7 +2409,7 @@ class TestInflightPerReqInvariant: ] for tid in tids_50: session._server._inflight_pop(tid) - assert "req-50" not in session._server._inflight_per_req + assert _srv_inflight_count(session, "req-50") == 0 assert not session._server._has_inflight_for("req-50") # Other ids unaffected. assert session._server._has_inflight_for("req-49") @@ -1559,6 +2429,7 @@ class TestConnectMsgValidation: ConnectMsg.BASE_ADDR: 0x1000, ConnectMsg.NUM_BLOCKS: 8, ConnectMsg.BLOCK_LEN: 4096, + ConnectMsg.HASH_SEED: "0", } def test_valid_message_passes(self): @@ -1600,13 +2471,25 @@ class TestConnectMsgValidation: with pytest.raises(ValueError, match="block_len"): ConnectMsg.validate(msg) + def test_missing_hash_seed(self): + msg = self._valid_msg() + del msg[ConnectMsg.HASH_SEED] + with pytest.raises(ValueError, match="hash_seed"): + ConnectMsg.validate(msg) + + def test_hash_seed_wrong_type(self): + msg = self._valid_msg() + msg[ConnectMsg.HASH_SEED] = 12345 # int, not str + with pytest.raises(ValueError, match="hash_seed"): + ConnectMsg.validate(msg) + class TestFetchMsgValidation: def _valid_msg(self) -> dict: return { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1", b"k2"], + FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [0, 1], } diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index 51f5162a99..cb90599cf8 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -6,9 +6,9 @@ Core abstractions for KV cache offloading in vLLM v1. from abc import ABC, abstractmethod from collections.abc import Collection, Iterable, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, field from enum import Enum, auto -from typing import TYPE_CHECKING, Any, NamedTuple, NewType +from typing import TYPE_CHECKING, Any, NamedTuple, NewType, TypeVar import numpy as np import torch @@ -45,10 +45,23 @@ def get_offload_group_idx(key: OffloadKey) -> int: return int.from_bytes(key[-4:], "big", signed=False) +_T = TypeVar("_T") + + @dataclass class ReqContext: req_id: str kv_transfer_params: dict[str, Any] | None = None + # Per-request scratch space keyed by value type, so a tier can parse + # kv_transfer_params once (in on_new_request) and read the result back + # on later calls for the same request. + _state: dict[type, Any] = field(default_factory=dict, repr=False, init=False) + + def set_state(self, val: Any) -> None: + self._state[type(val)] = val + + def get_state(self, cls: type[_T]) -> _T | None: + return self._state.get(cls) class LookupResult(Enum): diff --git a/vllm/v1/kv_offload/tiering/p2p/manager.py b/vllm/v1/kv_offload/tiering/p2p/manager.py index 95605fd173..f1f557ce38 100644 --- a/vllm/v1/kv_offload/tiering/p2p/manager.py +++ b/vllm/v1/kv_offload/tiering/p2p/manager.py @@ -8,6 +8,7 @@ Owns transports and a single bidirectional P2PSession per remote peer. from __future__ import annotations +import os import time import uuid from collections.abc import Iterable, Sequence @@ -23,12 +24,12 @@ from vllm.v1.kv_offload.base import ( OffloadKey, ReqContext, RequestOffloadingContext, + ScheduleEndContext, ) from vllm.v1.kv_offload.file_mapper import FileMapper from vllm.v1.kv_offload.tiering.base import ( JobMetadata, JobResult, - ScheduleEndContext, SecondaryTierManager, ) from vllm.v1.kv_offload.tiering.p2p.control import ControlTransport, ZmqTransport @@ -37,6 +38,7 @@ from vllm.v1.kv_offload.tiering.p2p.session import P2PSession if TYPE_CHECKING: from vllm.v1.kv_offload.base import OffloadingSpec + from vllm.v1.kv_offload.tiering.base import ParentManager from vllm.v1.kv_offload.tiering.p2p.control.base import ControlConnection logger = init_logger(__name__) @@ -59,24 +61,112 @@ _SHUTDOWN_DRAIN_TIMEOUT_S = 3.0 _DRAIN_SLEEP_S = 0.001 -def _prefill_params(kv_params: dict | None) -> dict | None: - """Return the ``prefill`` sub-dict, or None if absent. +def _remote_prefiller_params(kv_params: dict | None) -> dict | None: + """Return the ``remote_prefiller`` sub-dict, or None if absent. - Set on decoder requests; carries kv_request_id, remote_host, remote_port. + Set on decoder requests to name the remote prefiller they pull from; + carries kv_request_id, remote_host, remote_port. """ if not kv_params: return None - return kv_params.get("prefill") + return kv_params.get("remote_prefiller") -def _decode_params(kv_params: dict | None) -> dict | None: - """Return the ``decode`` sub-dict, or None if absent. +def _remote_decoder_params(kv_params: dict | None) -> dict | None: + """Return the ``remote_decoder`` sub-dict, or None if absent. - Set on prefiller requests; carries kv_request_id. + Set on prefiller requests to name the remote decoder they serve; + carries kv_request_id. """ if not kv_params: return None - return kv_params.get("decode") + return kv_params.get("remote_decoder") + + +def _remote_kv_source_params(kv_params: dict | None) -> dict | None: + """Return the ``remote_kv_source`` sub-dict, or None if absent. + + Set on symmetric-P2P consumer requests to name the remote source they + pull from; carries kv_request_id, remote_host, remote_port. + """ + if not kv_params: + return None + return kv_params.get("remote_kv_source") + + +def _peer_id_from_params(role_params: dict) -> str | None: + """Build ``host:port`` peer_id from a role-scoped sub-dict, or None.""" + host = role_params.get("remote_host") + port = role_params.get("remote_port") + if host and port: + return f"{host}:{port}" + return None + + +@dataclass(slots=True) +class P2PSourceInfo: + """Consumer side: this request fetches from a remote (prefiller or peer).""" + + kv_request_id: str + peer_id: str + do_probe: bool # False for remote_prefiller (PD), True for remote_kv_source + + +@dataclass(slots=True) +class P2PDestInfo: + """Producer side: a remote fetches this request's blocks from us. + + ``kv_request_id`` is None when the ``remote_decoder`` block is present + but malformed (no id); the block's presence still marks the request as + remote-decode, so submit_store must fail rather than store locally. + """ + + kv_request_id: str | None + + +def _parse_source(kv_params: dict | None) -> P2PSourceInfo | None: + """Parse the consumer sub-dict (PD ``remote_prefiller`` or symmetric + ``remote_kv_source``) into a ``P2PSourceInfo``, or None if absent/incomplete.""" + role = _remote_prefiller_params(kv_params) + do_probe = False + if role is None: + role = _remote_kv_source_params(kv_params) + do_probe = True + if not role: + return None + peer_id = _peer_id_from_params(role) + kv_request_id = role.get("kv_request_id") + if peer_id is None or not kv_request_id: + return None + return P2PSourceInfo( + kv_request_id=kv_request_id, + peer_id=peer_id, + do_probe=do_probe, + ) + + +def _parse_dest(kv_params: dict | None) -> P2PDestInfo | None: + """Parse the producer ``remote_decoder`` sub-dict into a ``P2PDestInfo``, + or None if the block is absent (not a remote-decode request).""" + role = _remote_decoder_params(kv_params) + if role is None: + return None + return P2PDestInfo(kv_request_id=role.get("kv_request_id") or None) + + +def _annotate_req_context(req_context: ReqContext) -> None: + """Parse kv_transfer_params once and cache the P2P routing state. + + Called from ``on_new_request``; later calls for the same request read + the cached ``P2PSourceInfo``/``P2PDestInfo`` via ``get_state`` instead + of re-parsing. + """ + source = _parse_source(req_context.kv_transfer_params) + if source is not None: + req_context.set_state(source) + dest = _parse_dest(req_context.kv_transfer_params) + if dest is not None: + req_context.set_state(dest) @dataclass @@ -157,6 +247,20 @@ class P2PSecondaryTierManager(SecondaryTierManager): **kwargs: Reserved for future tier-specific options. """ super().__init__(offloading_spec, primary_kv_view, tier_type) + # Block hashes chain from NONE_HASH, seeded from PYTHONHASHSEED + # (see init_none_hash in v1/core/kv_cache_utils.py). Peers with + # different seeds compute different hashes for identical content, so + # lookups silently miss and no KV crosses the wire. Require it here so + # a misconfigured P2P instance fails at startup rather than degrading + # silently; the value is also verified against each peer on handshake. + hash_seed = os.getenv("PYTHONHASHSEED") + if hash_seed is None: + raise ValueError( + "PYTHONHASHSEED must be set for P2P KV offload so that block " + "hashes match across instances. Set it to a fixed value (e.g. " + "PYTHONHASHSEED=0) on every P2P peer." + ) + self._hash_seed = hash_seed if host is None: host = envs.VLLM_P2P_SIDE_CHANNEL_HOST if port is None: @@ -210,6 +314,12 @@ class P2PSecondaryTierManager(SecondaryTierManager): # kv_request_ids that hit a transport/session failure; On load lookup() # rejects them so the request falls back to local prefill. self._failed_req_ids: set[str] = set() + # Synthetic lookup ctxs from reaped sessions still owing a + # ``parent.on_request_finished`` (the session's failed_serves). The + # dead session had no parent handle at teardown; these are flushed + # at the top of the next ``serve_external_requests`` where the + # handle is valid. + self._failed_serve_ctxs: list[ReqContext] = [] # ------------------------------------------------------------------ # SecondaryTierManager interface @@ -217,66 +327,81 @@ class P2PSecondaryTierManager(SecondaryTierManager): @override def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: - prefill = _prefill_params(req_context.kv_transfer_params) - if ( - not prefill - or not prefill.get("remote_host") - or not prefill.get("remote_port") - or not prefill.get("kv_request_id") - ): + source = req_context.get_state(P2PSourceInfo) + if source is None: + return LookupResult.MISS + if source.kv_request_id in self._failed_req_ids: return LookupResult.MISS - kv_request_id = prefill["kv_request_id"] - if kv_request_id in self._failed_req_ids: - return LookupResult.MISS + # Symmetric-P2P consumer (``remote_kv_source`` sub-dict): probe the + # peer asynchronously. First call registers the (kv_request_id, + # key) entry and returns RETRY; flush_pending_lookups() + # in on_schedule_end batches the LookupMsg; a later step's + # lookup() returns HIT/MISS once LookupRespMsg has arrived. + # PD path (``remote_prefiller`` sub-dict only) keeps the eager HIT. + if source.do_probe: + session = self._sessions.get(source.peer_id) + if session is None: + return LookupResult.MISS + result = session.register_lookup(source.kv_request_id, key) + if result is True: + return LookupResult.HIT + if result is False: + return LookupResult.MISS + return LookupResult.RETRY + + # PD consumer (we are the decoder): all kv blocks should be on the + # prefiller side. Return HIT immediately. return LookupResult.HIT @override def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: - """Open the outbound session toward the producer if needed. + """Parse kv_transfer_params once and open the outbound session. - On the decoder side (``prefill`` set), open a session toward the - producer at remote_host:remote_port so submit_load can issue - FetchMsg as soon as it fires. On the prefiller side, sessions - are created when the consumer's inbound connection arrives in - _accept_new_peers — submit_store no longer pre-creates anything. + Parses the P2P routing state onto ``req_context`` (cached for the + later lookup/submit/finish calls). On the consumer side + (``remote_prefiller`` for PD or ``remote_kv_source`` for symmetric + P2P), open a session toward the producer at remote_host:remote_port + so submit_load can issue FetchMsg as soon as it fires. On the + prefiller side, sessions are created when the consumer's inbound + connection arrives in _accept_new_peers — submit_store no longer + pre-creates anything. """ - prefill = _prefill_params(req_context.kv_transfer_params) - if prefill: - peer_id = self._remote_id_from_params(prefill) - if peer_id: - self._get_or_create_session(peer_id) + _annotate_req_context(req_context) + source = req_context.get_state(P2PSourceInfo) + if source is not None: + self._get_or_create_session(source.peer_id) return RequestOffloadingContext() @override def on_request_finished(self, req_context: ReqContext) -> None: """Cancels pending loads and prunes session-scoped state. - Decoder side (``prefill`` set): looks up the session by peer_id - because the producer's address is what addresses the client-role - load to cancel. Prefiller side (``decode`` set): looks up via - kv_request_id because peer_id is no longer carried on store-time + Consumer side (``remote_prefiller`` for PD or ``remote_kv_source`` + for symmetric-P2P): looks up the session by peer_id because the + producer's address is what addresses the client-role load to + cancel; also drops any pending symmetric-P2P lookup state via + ``session.finish_request``. + Prefiller side (``remote_decoder`` set): looks up via kv_request_id + because peer_id is no longer carried on store-time kv_transfer_params; if a session has bound the id, finish it. If no session has bound the id yet, this is a no-op: parked batches in `_unbound_stores` are left in place and cleaned up only by `_reap_unbound_stores` after `_UNBOUND_STORE_TIMEOUT_S`. """ - kv_params = req_context.kv_transfer_params - if not kv_params: - return - prefill = _prefill_params(kv_params) - decode = _decode_params(kv_params) - kv_request_id = (prefill or decode or {}).get("kv_request_id") + source = req_context.get_state(P2PSourceInfo) + dest = req_context.get_state(P2PDestInfo) + kv_request_id = source.kv_request_id if source is not None else None + if kv_request_id is None and dest is not None: + kv_request_id = dest.kv_request_id if not kv_request_id: return self._failed_req_ids.discard(kv_request_id) - if prefill: - peer_id = self._remote_id_from_params(prefill) - if peer_id: - session = self._sessions.get(peer_id) - if session is not None: - session.finish_request(kv_request_id) + if source is not None: + session = self._sessions.get(source.peer_id) + if session is not None: + session.finish_request(kv_request_id) return # Prefiller-side finish: identify the session via kv_request_id. @@ -293,24 +418,24 @@ class P2PSecondaryTierManager(SecondaryTierManager): assert len(keys) == len(block_ids) - kv_params = job_metadata.req_context.kv_transfer_params - decode = _decode_params(kv_params) + dest = job_metadata.req_context.get_state(P2PDestInfo) logger.debug( - "P2P %s: submit_store ENTRY job_id=%d blocks=%d decode=%s kv_request_id=%s", + "P2P %s: submit_store ENTRY job_id=%d blocks=%d " + "remote_decoder=%s kv_request_id=%s", self._local_id, job_id, len(block_ids), - decode is not None, - (decode or {}).get("kv_request_id"), + dest is not None, + dest.kv_request_id if dest is not None else None, ) - # Absent ``decode`` block => not a remote-decode request: succeed - # locally without parking. An empty/malformed dict is still a - # remote-decode signal and must fail the missing-id check below. - if decode is None: + # Absent ``remote_decoder`` block => not a remote-decode request: + # succeed locally without parking. An empty/malformed dict is still + # a remote-decode signal and must fail the missing-id check below. + if dest is None: self._finished_jobs.append(JobResult(job_id=job_id, success=True)) return - kv_request_id = decode.get("kv_request_id") + kv_request_id = dest.kv_request_id if not kv_request_id: logger.warning( "P2P %s: submit_store missing kv_request_id", @@ -350,32 +475,26 @@ class P2PSecondaryTierManager(SecondaryTierManager): keys = list(job_metadata.keys) block_ids = job_metadata.block_ids - prefill = _prefill_params(job_metadata.req_context.kv_transfer_params) + source = job_metadata.req_context.get_state(P2PSourceInfo) logger.debug( "P2P %s: submit_load ENTRY job_id=%d blocks=%d kv_request_id=%s peer=%s", self._local_id, job_id, len(block_ids), - (prefill or {}).get("kv_request_id"), - self._remote_id_from_params(prefill or {}), + source.kv_request_id if source is not None else None, + source.peer_id if source is not None else None, ) - if ( - not prefill - or not prefill.get("remote_host") - or not prefill.get("remote_port") - or not prefill.get("kv_request_id") - ): + if source is None: logger.debug( - "P2P %s: submit_load job_id=%d FAILED missing prefill params", + "P2P %s: submit_load job_id=%d FAILED missing consumer params", self._local_id, job_id, ) self._finished_jobs.append(JobResult(job_id=job_id, success=False)) return - kv_request_id = prefill["kv_request_id"] - peer_id = self._remote_id_from_params(prefill) - assert peer_id is not None # guaranteed by prefill checks above + kv_request_id = source.kv_request_id + peer_id = source.peer_id if not keys: logger.debug( @@ -442,10 +561,7 @@ class P2PSecondaryTierManager(SecondaryTierManager): warned = False while True: self._poll_once() - pending = any( - s._client._inbound or s._server._inflight - for s in self._sessions.values() - ) + pending = any(s.has_pending_work for s in self._sessions.values()) if not pending: return if not warned and time.monotonic() - start > 5.0: @@ -456,31 +572,45 @@ class P2PSecondaryTierManager(SecondaryTierManager): warned = True time.sleep(_DRAIN_SLEEP_S) + @override + def serve_external_requests(self, parent: ParentManager) -> None: + """Serve inbound peer lookups against the tiering manager. + + Called once per scheduler step (before this tier's + ``on_schedule_end``) with a ``parent`` handle valid only for the + duration of the call — the sole window in which the P2P server + role may query the tiering manager. First release bookkeeping for + the failed serves left by a reaped session, then let every live + session resolve its enqueued inbound LookupMsgs. + """ + if self._failed_serve_ctxs: + for ctx in self._failed_serve_ctxs: + parent.on_request_finished(ctx) + self._failed_serve_ctxs = [] + for session in self._sessions.values(): + session.serve_external_requests(parent) + @override def on_schedule_end(self, context: ScheduleEndContext) -> None: - return + # Flush any p2p lookups aggregated during this step. + # One LookupMsg per (peer, kv_request_id) with unsent entries; + # send-gating happens inside the session if not yet ready. + for session in self._sessions.values(): + session.flush_pending_lookups() # ------------------------------------------------------------------ # Internal # ------------------------------------------------------------------ - @staticmethod - def _remote_id_from_params(role_params: dict) -> str | None: - """Build peer_id from a role-scoped sub-dict (``prefill``/``p2p``).""" - host = role_params.get("remote_host") - port = role_params.get("remote_port") - if host and port: - return f"{host}:{port}" - return None - def _get_or_create_session(self, peer_id: str) -> P2PSession: """Return the existing session for peer_id, or open one outbound. - Decoder-side helper for on_new_request: when ``prefill`` is set, - the consumer must reach the producer at peer_id. If we already - have a session toward that peer (from a prior load or a - peer-initiated inbound), reuse it; otherwise open an outbound - ControlConnection and build a connected session. + Consumer-side helper for on_new_request: when ``remote_prefiller`` + (PD) or ``remote_kv_source`` (symmetric P2P) is set, the consumer must reach the + producer at peer_id. If we already have a session toward that + peer (from a prior load or a peer-initiated inbound), reuse it; + otherwise open an outbound ControlConnection and build a + connected session. """ session = self._sessions.get(peer_id) if session is not None: @@ -491,6 +621,7 @@ class P2PSecondaryTierManager(SecondaryTierManager): local_id=self._local_id, transport=self._data, local_block_len=self._data.block_len, + local_hash_seed=self._hash_seed, conn=conn, ) self._sessions[peer_id] = session @@ -512,6 +643,7 @@ class P2PSecondaryTierManager(SecondaryTierManager): local_id=self._local_id, transport=self._data, local_block_len=self._data.block_len, + local_hash_seed=self._hash_seed, conn=conn, ) logger.info( @@ -546,12 +678,19 @@ class P2PSecondaryTierManager(SecondaryTierManager): ] for kid in stale_kv_ids: del self._kv_to_session[kid] - failed_loads, failed_stores = session.close() - for job_id, kv_request_id in failed_loads: + close_result = session.close() + for job_id in close_result.failed_jobs: self._finished_jobs.append(JobResult(job_id=job_id, success=False)) - self._failed_req_ids.add(kv_request_id) - for job_id in failed_stores: + for job_id in close_result.failed_stores: self._finished_jobs.append(JobResult(job_id=job_id, success=False)) + # Fail every client-side request (in-flight loads plus unresolved + # symmetric-P2P probes) toward the dead peer so lookup() returns + # MISS (local prefill) instead of RETRY forever — even if a fresh + # session to the same peer is later opened by another request. + self._failed_req_ids.update(close_result.failed_req_ids) + # Release the TieringManager's per-request bookkeeping for the + # dead session's synthetic lookups on the next serve_external_requests. + self._failed_serve_ctxs.extend(close_result.failed_serves) self._data.remove_remote_peer(pid) logger.warning("P2P %s: peer %s down", self._local_id, pid) @@ -648,6 +787,9 @@ class P2PSecondaryTierManager(SecondaryTierManager): def shutdown(self) -> None: self._drain_inflight_for_shutdown() for session in self._sessions.values(): + # Orphan ctxs from close() are intentionally dropped: the manager + # is being torn down, so there is no next serve_external_requests + # to flush them and no TieringManager left to release. session.close() self._sessions.clear() self._kv_to_session.clear() diff --git a/vllm/v1/kv_offload/tiering/p2p/session/__init__.py b/vllm/v1/kv_offload/tiering/p2p/session/__init__.py index 82148e8340..2efaf09e34 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/__init__.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/__init__.py @@ -3,6 +3,7 @@ from vllm.v1.kv_offload.tiering.p2p.session.session import ( LoadResult, P2PSession, + SessionCloseResult, SessionPollResult, StoreResult, ) @@ -10,6 +11,7 @@ from vllm.v1.kv_offload.tiering.p2p.session.session import ( __all__ = [ "LoadResult", "P2PSession", + "SessionCloseResult", "SessionPollResult", "StoreResult", ] diff --git a/vllm/v1/kv_offload/tiering/p2p/session/client.py b/vllm/v1/kv_offload/tiering/p2p/session/client.py index f080cb79a3..ed91f7f9b0 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/client.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/client.py @@ -11,16 +11,19 @@ callback injected by the coordinator (which gates on ConnectAck). from __future__ import annotations +import enum import time from collections.abc import Callable, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import TYPE_CHECKING, NamedTuple from vllm.logger import init_logger +from vllm.v1.kv_offload.base import OffloadKey from vllm.v1.kv_offload.tiering.p2p.session.protocol import ( TYPE_KEY, AbortFetchMsg, FetchMsg, + LookupMsg, ) if TYPE_CHECKING: @@ -32,16 +35,60 @@ _LOAD_TIMEOUT_S = 30.0 _ABORT_ACK_TIMEOUT_S = 10.0 +class ClientPhase(enum.Enum): + """Lifecycle of a request's client-side lookup/fetch signalling. + + Advances monotonically. Only ``finish`` reads it, to decide + whether a terminal empty FetchMsg is owed to release the peer's + lookup state (owed only from ``PROBING``: a LookupMsg went out but no + FetchMsg has since closed the peer's lookup phase). + """ + + REGISTERED = enum.auto() # keys in probes/unsent, nothing sent yet + PROBING = enum.auto() # LookupMsg flushed, awaiting responses + FETCH_SENT = enum.auto() # FetchMsg sent (real or terminal empty) + + @dataclass -class _InboundRequestState: - """Client-role state for a single load request.""" +class _InboundLoadState: + """Client-role state for a single in-flight load request. + + Lives on ``_ClientRequestState.load`` for the duration of a fetch; + the owning kv_request_id is the dict key, so it isn't stored here. + """ job_id: int # opaque ID assigned by the manager to this load request - kv_request_id: str submitted_at: float aborted_at: float | None = None +@dataclass +class _ClientRequestState: + """Per-kv_request_id client-side state. + + One entry per kv_request_id we're driving. Lookup-phase fields are + used only by symmetric P2P (``do_p2p_fetch``); PD-only loads leave + ``probes``/``unsent`` empty and drive just ``phase`` and ``load``. An + entry is dropped once every field is idle — see ``ClientRole._maybe_prune``. + """ + + # -- Lookup phase (symmetric P2P only; untouched for PD) -- + # Probe outcome per OffloadKey: None while in-flight (registered/sent + # but unresolved), True/False once a LookupRespMsg lands. There is no + # timeout — finish (via finish_request) is guaranteed after + # the request's lookup() calls and clears every probe, so an + # unanswered probe simply stays None until then. + probes: dict[OffloadKey, bool | None] = field(default_factory=dict) + # OffloadKeys registered but not yet flushed onto the wire. Drained and + # cleared by the next flush_pending_lookups. + unsent: list[OffloadKey] = field(default_factory=list) + + # Monotonic lookup/fetch signalling phase; see ``ClientPhase``. + phase: ClientPhase = ClientPhase.REGISTERED + # Set while a fetch is in flight; cleared on completion/abort/timeout. + load: _InboundLoadState | None = None + + class LoadResult(NamedTuple): """Result from a session poll, client side.""" @@ -50,6 +97,21 @@ class LoadResult(NamedTuple): success: bool +class ClientCloseResult(NamedTuple): + """What the manager must fail when a peer session is torn down. + + ``failed_jobs`` is the ``job_id`` of every load still in flight; each + becomes a ``JobResult(success=False)``. ``failed_req_ids`` is every + kv_request_id whose lookup() would otherwise defer forever on the dead + peer — the in-flight loads plus any request holding an unresolved + symmetric-P2P probe. ``failed_jobs`` is the subset of ``failed_req_ids`` + that had a load job. + """ + + failed_jobs: list[int] + failed_req_ids: list[str] + + class ClientRole: """Client-side load state machine for one peer session. @@ -61,9 +123,55 @@ class ClientRole: def __init__(self, peer_id: str, send: Callable[[dict], None]) -> None: self._peer_id = peer_id self._send = send - self._inbound: dict[str, _InboundRequestState] = {} + # All per-kv_request_id state lives here. Entries are created + # lazily by request_blocks / register_lookup and dropped by + # _maybe_prune once every field is idle. + self._requests: dict[str, _ClientRequestState] = {} + # kv_request_ids with unsent lookup keys for the next flush to + # visit — the work-list that keeps flush_pending_lookups from + # scanning every request each scheduler step. Mirrors the server's + # _serve_pending. Populated by register_lookup, drained by + # flush_pending_lookups, and discarded on finish/close. + self._flush_pending: set[str] = set() + # kv_request_ids with a fetch in flight (``st.load is not None``) — + # the work-list collect_results walks for timeouts, and the + # has_active_loads predicate, instead of scanning every request. + # Kept in exact sync with ``st.load``: armed in request_blocks, + # discarded wherever load is cleared, and cleared on close. + self._active_loads: set[str] = set() self._completed_loads: list[LoadResult] = [] + # ------------------------------------------------------------------ + # State helpers + # ------------------------------------------------------------------ + + def _get_or_create_request(self, kv_request_id: str) -> _ClientRequestState: + """Get or create the state entry for a kv_request_id.""" + st = self._requests.get(kv_request_id) + if st is None: + st = _ClientRequestState() + self._requests[kv_request_id] = st + return st + + def _maybe_prune(self, kv_request_id: str) -> None: + """Drop the entry once it holds no live load or lookup state. + + The sticky ``phase`` is only read by ``finish``. A probe + clears when its fetch is issued (``request_blocks``) or when the + request finishes (``finish``/``close``); in the former case + ``load`` is set and keeps the entry alive, in the latter the phase + is no longer needed — so dropping on emptiness never loses a phase + still in use. + """ + st = self._requests.get(kv_request_id) + if st is not None and st.load is None and not st.probes and not st.unsent: + del self._requests[kv_request_id] + + @property + def has_active_loads(self) -> bool: + """True if any kv_request_id has a fetch in flight.""" + return bool(self._active_loads) + # ------------------------------------------------------------------ # Public API # ------------------------------------------------------------------ @@ -72,7 +180,7 @@ class ClientRole: self, job_id: JobId, kv_request_id: str, - keys: Sequence[bytes], + keys: Sequence[OffloadKey], block_ids: Sequence[int], send_ready: bool, ) -> None: @@ -86,44 +194,101 @@ class ClientRole: len(block_ids), send_ready, ) - self._inbound[kv_request_id] = _InboundRequestState( + st = self._get_or_create_request(kv_request_id) + st.load = _InboundLoadState( job_id=job_id, - kv_request_id=kv_request_id, submitted_at=time.monotonic(), ) + self._active_loads.add(kv_request_id) + st.phase = ClientPhase.FETCH_SENT self._send( { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: kv_request_id, - FetchMsg.BLOCK_HASHES: list(keys), + FetchMsg.KEYS: list(keys), FetchMsg.BLOCK_INDEXES: [int(idx) for idx in block_ids], } ) + # Issuing the fetch ends this request's lookup phase, so drop all + # probe state. Once the peer serves this fetch both sides unpin, so + # the producer may evict the block; a stale cached True would + # otherwise let a re-scheduled lookup() return HIT without + # re-probing, pointing at a block the producer no longer holds. + # Clearing forces a fresh LookupMsg on re-schedule so the producer + # answers from current state. For a symmetric-P2P request (probes + # populated) every fetched block was a confirmed HIT; a PD-only load + # never probes, so probes is empty and the clear is a no-op. + if st.probes: + assert all(st.probes.get(key) is True for key in keys) + st.probes.clear() - def cancel(self, kv_request_id: str) -> None: - """Cancel a pending load. Sends AbortFetchMsg if still active.""" - req = self._inbound.pop(kv_request_id, None) - if req is not None and req.aborted_at is None: + def finish(self, kv_request_id: str) -> None: + """Finish a request: abort any in-flight load and release lookup state. + + Called from the session's ``finish_request``. The two branches are + mutually exclusive: ``load`` is set only by ``request_blocks``, which + also advances ``phase`` to ``FETCH_SENT``, and nothing moves it back + to ``PROBING`` — so a fetch in flight never coexists with the + ``PROBING`` phase. + + - Fetch in flight (``load`` set, phase ``FETCH_SENT``): send an + AbortFetchMsg unless the load is already aborting, then drop it. + - Outstanding lookups (``PROBING``): a LookupMsg was flushed but no + FetchMsg has closed the peer's lookup phase. Every FetchMsg the + server receives in p2p mode is its "request finished" signal (it + releases lookup state and fires ``cb.finish_request``); when the + client's lookups all missed no FetchMsg is otherwise sent, so emit + a terminal empty one purely to trigger those semantics. In + ``REGISTERED`` the peer never received a LookupMsg and in + ``FETCH_SENT`` a FetchMsg already closed the phase, so neither owes + a terminal FetchMsg. + + Then drop all probe/lookup state and prune the entry. + """ + st = self._requests.get(kv_request_id) + if st is None: + return + if st.load is not None: + if st.load.aborted_at is None: + self._send( + { + TYPE_KEY: AbortFetchMsg.TYPE, + AbortFetchMsg.KV_REQUEST_ID: kv_request_id, + } + ) + st.load = None + self._active_loads.discard(kv_request_id) + elif st.phase is ClientPhase.PROBING: + st.phase = ClientPhase.FETCH_SENT self._send( { - TYPE_KEY: AbortFetchMsg.TYPE, - AbortFetchMsg.KV_REQUEST_ID: kv_request_id, + TYPE_KEY: FetchMsg.TYPE, + FetchMsg.KV_REQUEST_ID: kv_request_id, + FetchMsg.KEYS: [], + FetchMsg.BLOCK_INDEXES: [], } ) + st.probes.clear() + st.unsent.clear() + self._flush_pending.discard(kv_request_id) + self._maybe_prune(kv_request_id) def on_transfer_done(self, kv_request_id: str, success: bool) -> None: """Handle a TransferDoneMsg from the peer.""" - req = self._inbound.pop(kv_request_id, None) - if req is not None: + st = self._requests.get(kv_request_id) + if st is not None and st.load is not None: self._completed_loads.append( LoadResult( - job_id=req.job_id, + job_id=st.load.job_id, kv_request_id=kv_request_id, success=success, ) ) + st.load = None + self._active_loads.discard(kv_request_id) + self._maybe_prune(kv_request_id) else: - # No matching _inbound entry: either a duplicate + # No matching in-flight load: either a duplicate # transfer_done from the peer (protocol violation) or a # benign race with a local cancel/abort/timeout that # already popped the entry. We don't track terminated ids, @@ -137,23 +302,26 @@ class ClientRole: def on_abort_ack(self, kv_request_id: str) -> None: """Handle an AbortAckMsg from the peer.""" - req = self._inbound.pop(kv_request_id, None) - if req is not None: + st = self._requests.get(kv_request_id) + if st is not None and st.load is not None: logger.warning( "P2PSession %s: load request %s (job_id=%d) timed out; " "load job completed with failure. If this recurs, ensure " "PYTHONHASHSEED is set to the same value on all nodes.", self._peer_id, kv_request_id, - req.job_id, + st.load.job_id, ) self._completed_loads.append( LoadResult( - job_id=req.job_id, + job_id=st.load.job_id, kv_request_id=kv_request_id, success=False, ) ) + st.load = None + self._active_loads.discard(kv_request_id) + self._maybe_prune(kv_request_id) else: # See on_transfer_done: same ambiguity (duplicate ack # vs. raced with local cancel/timeout that already popped). @@ -164,19 +332,141 @@ class ClientRole: kv_request_id, ) + # ------------------------------------------------------------------ + # Symmetric-P2P lookup (do_p2p_fetch=true) + # ------------------------------------------------------------------ + + def register_lookup(self, kv_request_id: str, key: bytes) -> bool | None: + """Register or resolve one (kv_request_id, key) probe. + + Idempotent across scheduler steps: + - First call: creates a pending entry, returns None. + - Subsequent calls while in-flight: returns None. + - Once a LookupRespMsg has resolved the entry: returns the cached + bool result on every call without popping it. + + A resolved entry is retained until its fetch is issued + (``request_blocks`` pops it) or the request finishes + (``finish`` clears all entries for the id). A request's + block set can be re-probed across steps, so popping on read would + make a repeat probe of an already-resolved key look brand-new and + re-queue it, emitting a redundant LookupMsg for an answer we + already hold. Keeping the entry until fetch makes repeat probes + free; clearing it at fetch forces a fresh probe if the request is + re-scheduled, since the block is unpinned once served. + """ + st = self._get_or_create_request(kv_request_id) + okey = OffloadKey(key) + if okey in st.probes: + return st.probes[okey] + st.probes[okey] = None + st.unsent.append(okey) + self._flush_pending.add(kv_request_id) + logger.debug( + "P2P LOOKUP client %s: REGISTER kv_request_id=%s key=%s (unsent=%d)", + self._peer_id, + kv_request_id, + key.hex()[:16], + len(st.unsent), + ) + return None + + def flush_pending_lookups(self) -> None: + """Send a LookupMsg for each kv_request_id with unsent entries. + + Called once per scheduler step from the manager's + ``on_schedule_end()``. A request's block set may be discovered + across several scheduler steps, so more than one LookupMsg can + go out per kv_request_id — one per step that registered new + keys. register_lookup() de-dups in-flight and already-resolved + (req_id, key) pairs, so each LookupMsg carries only the keys + first probed in that step. The peer's lookup phase for the id is + still closed by exactly one FetchMsg, which the client contract + guarantees is sent after every lookup for the id has resolved + (see request_blocks / finish). Send-gating is handled by + the injected ``_send`` callback (queues until ConnectAckMsg if + needed). + + Only requests that registered new keys since the last flush are + visited — the ``_flush_pending`` work-list avoids scanning every + live request each scheduler step. + """ + for req_id in self._flush_pending: + st = self._requests.get(req_id) + if st is None or not st.unsent: + continue + # Record that the peer now holds lookup state for this id so + # finish knows a terminal empty FetchMsg may be owed. + # Only promote from REGISTERED: once a fetch has gone out + # (FETCH_SENT) a later LookupMsg must not regress the phase, as + # no terminal FetchMsg is owed for an already-fetched request. + if st.phase is ClientPhase.REGISTERED: + st.phase = ClientPhase.PROBING + logger.debug( + "P2P LOOKUP client %s: SEND LookupMsg kv_request_id=%s keys=%d", + self._peer_id, + req_id, + len(st.unsent), + ) + self._send( + { + TYPE_KEY: LookupMsg.TYPE, + LookupMsg.KV_REQUEST_ID: req_id, + LookupMsg.KEYS: list(st.unsent), + } + ) + st.unsent = [] + self._flush_pending.clear() + + def on_lookup_resp( + self, + kv_request_id: str, + keys: Sequence[bytes], + hits: Sequence[bool], + ) -> None: + """Apply per-pair hit/miss results from a peer. + + Pairs that don't match a known entry (already cancelled or + never asked) are silently dropped — the producer is free to + split or coalesce responses. + """ + n_hit = sum(1 for hit in hits if hit) + logger.debug( + "P2P LOOKUP client %s: RECV LookupRespMsg kv_request_id=%s " + "keys=%d hits=%d misses=%d", + self._peer_id, + kv_request_id, + len(keys), + n_hit, + len(hits) - n_hit, + ) + st = self._requests.get(kv_request_id) + if st is None: + return + for h, hit in zip(keys, hits): + key = OffloadKey(h) + if key in st.probes: + st.probes[key] = hit + def collect_results(self) -> list[LoadResult]: - """Walk timeouts and drain completed loads. + """Walk load timeouts and drain completed loads. Active requests past ``_LOAD_TIMEOUT_S`` get an AbortFetchMsg sent and enter the aborting phase. Aborting requests past ``_ABORT_ACK_TIMEOUT_S`` are surfaced as failed loads. + + Lookups have no timeout: an unanswered probe stays None (RETRY) + until finish_request clears it — see ``_ClientRequestState.probes``. """ now = time.monotonic() to_remove: list[str] = [] - for req_id, req in self._inbound.items(): - if req.aborted_at is None: - if now - req.submitted_at >= _LOAD_TIMEOUT_S: - req.aborted_at = now + for req_id in self._active_loads: + st = self._requests[req_id] + assert st.load is not None + load = st.load + if load.aborted_at is None: + if now - load.submitted_at >= _LOAD_TIMEOUT_S: + load.aborted_at = now logger.warning( "P2PSession %s: %s timed out, sending abort", self._peer_id, @@ -189,11 +479,11 @@ class ClientRole: } ) else: - if now - req.aborted_at >= _ABORT_ACK_TIMEOUT_S: + if now - load.aborted_at >= _ABORT_ACK_TIMEOUT_S: to_remove.append(req_id) self._completed_loads.append( LoadResult( - job_id=req.job_id, + job_id=load.job_id, kv_request_id=req_id, success=False, ) @@ -204,15 +494,31 @@ class ClientRole: req_id, ) for req_id in to_remove: - self._inbound.pop(req_id) + self._requests[req_id].load = None + self._active_loads.discard(req_id) + self._maybe_prune(req_id) results = self._completed_loads self._completed_loads = [] return results - def close(self) -> list[tuple[int, str]]: - """Tear down. Returns ``(job_id, kv_request_id)`` for pending loads.""" - failed = [(req.job_id, req.kv_request_id) for req in self._inbound.values()] - self._inbound.clear() + def close(self) -> ClientCloseResult: + """Tear down, reporting work the dead peer can no longer complete. + + A request is failed if it has a load in flight (its job fails) or + holds an unresolved symmetric-P2P probe (its lookup() would defer + forever on an answer that can never arrive). See ``ClientCloseResult``. + """ + failed_jobs = [ + st.load.job_id for st in self._requests.values() if st.load is not None + ] + failed_req_ids = [ + req_id + for req_id, st in self._requests.items() + if st.load is not None or any(hit is None for hit in st.probes.values()) + ] + self._requests.clear() + self._flush_pending.clear() + self._active_loads.clear() self._completed_loads.clear() - return failed + return ClientCloseResult(failed_jobs=failed_jobs, failed_req_ids=failed_req_ids) diff --git a/vllm/v1/kv_offload/tiering/p2p/session/protocol.py b/vllm/v1/kv_offload/tiering/p2p/session/protocol.py index 3a2ab5fd88..8988c7e79f 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/protocol.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/protocol.py @@ -26,6 +26,14 @@ Block Transfer Flow (happy path) 1. Client sends FetchMsg with a kv_request_id and lists of block keys + remote indexes where it wants the data written. + In p2p mode FetchMsg is also the server-side "request finished" + signal for the id: no further ``cb.create_store_job`` will fire + (parked LookupMsg batches are popped, so pending-key resolution + cannot promote a HIT after this point), all server-side lookup + state for the id is released, and ``cb.finish_request`` fires on + each dropped batch. The client emits exactly one FetchMsg per + lookup-touched request, including an empty one when no blocks + end up being fetched. 2. Server matches requested blocks against locally stored blocks: - Blocks already available are transferred immediately via RDMA. - Blocks not yet available are recorded as "demanded" and @@ -119,6 +127,9 @@ class ConnectMsg: BLOCK_LEN: Size in bytes of each block (must match between peers). CONFIG_FINGERPRINT: SHA-256 prefix of the model configuration. Peers with different fingerprints are incompatible. + HASH_SEED: The peer's PYTHONHASHSEED. Block hashes chain from a seed + derived from it, so peers with different values compute different + hashes for identical content and must not exchange blocks. """ TYPE = "connect" @@ -128,6 +139,7 @@ class ConnectMsg: NUM_BLOCKS = "num_blocks" BLOCK_LEN = "block_len" CONFIG_FINGERPRINT = "config_fingerprint" + HASH_SEED = "hash_seed" @staticmethod def validate(msg: dict) -> None: @@ -137,6 +149,7 @@ class ConnectMsg: _require_non_neg_int(msg, ConnectMsg.BASE_ADDR) _require_pos_int(msg, ConnectMsg.NUM_BLOCKS) _require_pos_int(msg, ConnectMsg.BLOCK_LEN) + _require(msg, ConnectMsg.HASH_SEED, str) class ConnectAckMsg: @@ -165,37 +178,106 @@ class DisconnectMsg: class FetchMsg: - """Client → Server: request blocks by key. + """Client → Server: request blocks by key and close the lookup phase. + + In p2p mode FetchMsg is also the server-side "request finished" + signal for ``kv_request_id``: on receipt the server (a) fires no + further ``cb.create_store_job`` for this id — parked LookupMsg + batches are popped, so ``_resolve_pending_lookups`` cannot promote + a HIT_PENDING / RETRY key into a fresh pin after this point — and + (b) calls ``cb.finish_request(batch.ctx)`` on each dropped batch + so the TieringManager can release per-batch bookkeeping. In the + all-miss case the client emits an empty FetchMsg (``KEYS`` + and ``BLOCK_INDEXES`` both empty) purely to fire this signal. Fields: KV_REQUEST_ID: Identifies this block transfer request. - BLOCK_HASHES: List of block keys (OffloadKey bytes). - BLOCK_INDEXES: List of remote block indexes (same length as BLOCK_HASHES). + KEYS: List of block keys (OffloadKey bytes). May be empty. + BLOCK_INDEXES: List of remote block indexes (same length as KEYS). """ TYPE = "fetch" KV_REQUEST_ID = "kv_request_id" - BLOCK_HASHES = "block_hashes" + KEYS = "keys" BLOCK_INDEXES = "block_indexes" @staticmethod def validate(msg: dict) -> None: """Raise ValueError if any field has an invalid type or value.""" _require(msg, FetchMsg.KV_REQUEST_ID, str) - _require_list(msg, FetchMsg.BLOCK_HASHES) + _require_list(msg, FetchMsg.KEYS) _require_list(msg, FetchMsg.BLOCK_INDEXES) - hashes = msg[FetchMsg.BLOCK_HASHES] + keys = msg[FetchMsg.KEYS] indexes = msg[FetchMsg.BLOCK_INDEXES] - if len(hashes) != len(indexes): + if len(keys) != len(indexes): raise ValueError( - f"block_hashes/block_indexes length mismatch: " - f"{len(hashes)} vs {len(indexes)}" + f"keys/block_indexes length mismatch: {len(keys)} vs {len(indexes)}" ) for idx in indexes: if not isinstance(idx, int) or idx < 0: raise ValueError(f"block_indexes: invalid index {idx!r}") +class LookupMsg: + """Client → Server: probe which block keys the peer holds. + + Sent on the consumer side under symmetric P2P (do_p2p_fetch=true) + after the consumer has aggregated per-block lookups across a + scheduler step. The producer replies with one or more LookupRespMsg + covering the requested keys. + + Fields: + KV_REQUEST_ID: Identifies this lookup transaction. + KEYS: List of block keys (OffloadKey bytes) to probe. + """ + + TYPE = "lookup" + KV_REQUEST_ID = "kv_request_id" + KEYS = "keys" + + @staticmethod + def validate(msg: dict) -> None: + """Raise ValueError if any field has an invalid type or value.""" + _require(msg, LookupMsg.KV_REQUEST_ID, str) + _require_list(msg, LookupMsg.KEYS) + + +class LookupRespMsg: + """Server → Client: per-key hit/miss answer for a prior LookupMsg. + + Carries two parallel arrays of equal length so each (key, + hit) pair is self-describing. The producer is free to split or + coalesce responses across multiple LookupMsgs for the same + KV_REQUEST_ID — the consumer matches each pair back to its + pending entry by (KV_REQUEST_ID, key). + + Fields: + KV_REQUEST_ID: The lookup transaction this responds to. + KEYS: List of block keys answered by this message. + HITS: Parallel list of bools — True if the producer holds the + corresponding block, False otherwise. + """ + + TYPE = "lookup_resp" + KV_REQUEST_ID = "kv_request_id" + KEYS = "keys" + HITS = "hits" + + @staticmethod + def validate(msg: dict) -> None: + """Raise ValueError if any field has an invalid type or value.""" + _require(msg, LookupRespMsg.KV_REQUEST_ID, str) + _require_list(msg, LookupRespMsg.KEYS) + _require_list(msg, LookupRespMsg.HITS) + keys = msg[LookupRespMsg.KEYS] + hits = msg[LookupRespMsg.HITS] + if len(keys) != len(hits): + raise ValueError(f"keys/hits length mismatch: {len(keys)} vs {len(hits)}") + for hit in hits: + if not isinstance(hit, bool): + raise ValueError(f"hits: invalid value {hit!r}") + + class TransferDoneMsg: """Server → Client: all blocks transferred for a request. diff --git a/vllm/v1/kv_offload/tiering/p2p/session/server.py b/vllm/v1/kv_offload/tiering/p2p/session/server.py index 2a6c8c6924..3709ecd9c5 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/server.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/server.py @@ -18,26 +18,32 @@ coordinator's ``_dispatch_message`` can reuse its existing from __future__ import annotations import time -from collections.abc import Callable, Sequence +from collections.abc import Callable, Iterable, Sequence from dataclasses import dataclass, field from typing import TYPE_CHECKING, NamedTuple from vllm.logger import init_logger -from vllm.v1.kv_offload.base import OffloadKey +from vllm.v1.kv_offload.base import LookupResult, OffloadKey, ReqContext from vllm.v1.kv_offload.tiering.p2p.session.protocol import ( TYPE_KEY, AbortAckMsg, + LookupRespMsg, TransferDoneMsg, ) if TYPE_CHECKING: - from vllm.v1.kv_offload.tiering.base import JobId + from vllm.v1.kv_offload.tiering.base import JobId, ParentManager from vllm.v1.kv_offload.tiering.p2p.data import DataTransport logger = init_logger(__name__) _STORE_TIMEOUT_S = 30.0 _CANCEL_DRAIN_TIMEOUT_S = 10.0 +# Cap on time the server holds a HIT_PENDING / RETRY key from an inbound +# LookupMsg before falling back to MISS. Long enough that an in-flight +# primary write or a just-started promotion typically completes; short +# enough that the consumer doesn't sit idle on a stuck producer. +_LOOKUP_PENDING_TIMEOUT_S = 5.0 class StoreResult(NamedTuple): @@ -69,8 +75,8 @@ class _MatchResult(NamedTuple): class _OutboundRequestState: """Server-role state for a single peer fetch request. - The owning ``kv_request_id`` is the dict key in - ``ServerRole._outbound`` and is not duplicated on the value. + The owning ``kv_request_id`` is the ``ServerRole._requests`` dict key + and is not duplicated on the value. """ demand_received: bool = False @@ -90,7 +96,7 @@ class _OutboundRequestState: def add_stored_blocks( self, - block_hashes: Sequence[OffloadKey], + keys: Sequence[OffloadKey], block_ids: Sequence[int], job_id: int, ) -> _MatchResult: @@ -98,13 +104,13 @@ class _OutboundRequestState: self.pending_job_ids.add(job_id) local_idxs: list[int] = [] remote_idxs: list[int] = [] - for block_hash, local_idx in zip(block_hashes, block_ids): - remote_idx = self.demanded.pop(block_hash, None) + for key, local_idx in zip(keys, block_ids): + remote_idx = self.demanded.pop(key, None) if remote_idx is not None: local_idxs.append(local_idx) remote_idxs.append(remote_idx) else: - self.available[block_hash] = (job_id, local_idx) + self.available[key] = (job_id, local_idx) return _MatchResult( local_idxs=local_idxs, remote_idxs=remote_idxs, @@ -113,25 +119,25 @@ class _OutboundRequestState: def add_fetch_demand( self, - block_hashes: Sequence[OffloadKey], + keys: Sequence[OffloadKey], block_indexes: Sequence[int], ) -> _MatchResult: """Register the peer's fetch demand. Returns matched pairs.""" self.demand_received = True - self.remaining = len(block_hashes) + self.remaining = len(keys) local_idxs: list[int] = [] remote_idxs: list[int] = [] job_ids: set[int] = set() - for block_hash, remote_idx in zip(block_hashes, block_indexes): - stored_entry = self.available.pop(block_hash, None) + for key, remote_idx in zip(keys, block_indexes): + stored_entry = self.available.pop(key, None) if stored_entry is not None: stored_job_id, local_idx = stored_entry local_idxs.append(local_idx) remote_idxs.append(remote_idx) job_ids.add(stored_job_id) else: - self.demanded[block_hash] = remote_idx + self.demanded[key] = remote_idx return _MatchResult( local_idxs=local_idxs, remote_idxs=remote_idxs, @@ -139,6 +145,78 @@ class _OutboundRequestState: ) +@dataclass +class _ActiveLookup: + """In-flight state for one inbound LookupMsg. + + Aggregates per-key HIT/MISS resolutions and defers the single + outbound LookupRespMsg until every key has been resolved or the + ``deadline`` fires (remaining ``pending`` keys then force-MISS). + Exactly one LookupRespMsg is emitted per LookupMsg, carrying every + key in the original wire order. + """ + + lookup_id: int + kv_request_id: str + ctx: ReqContext + # Keys from the inbound LookupMsg, preserved in wire order so + # the aggregated response goes back in the same order. + keys: list[OffloadKey] = field(default_factory=list) + # Per-key resolution: True = HIT, False = MISS. A key is present + # here once definitively resolved; still-pending keys are only + # in ``pending``. + resolved: dict[OffloadKey, bool] = field(default_factory=dict) + # Keys still awaiting resolution (HIT_PENDING / RETRY from + # ``parent.lookup``); re-polled by ``_resolve_pending_lookups``. + pending: set[OffloadKey] = field(default_factory=set) + # Absolute deadline (``time.monotonic``). Once reached, remaining + # ``pending`` keys are force-resolved to MISS so the consumer + # can fall back instead of waiting on a stuck producer. + deadline: float = 0.0 + + +class _PendingLookup(NamedTuple): + """A raw inbound LookupMsg awaiting resolution. + + Enqueued by ``on_lookup`` during dispatch and drained by the next + ``serve_external_requests``. The deadline for any resulting + HIT_PENDING / RETRY key is measured from ``enqueued_at``. + """ + + keys: list[OffloadKey] + enqueued_at: float + + +@dataclass +class _ServerRequestState: + """Per-kv_request_id server-side state. + + Consolidates the outbound serve/transfer state, the symmetric-P2P + inbound-lookup state, abort bookkeeping, and the inflight-transfer + count for one kv_request_id. The owning id is the ``_requests`` dict + key and is not duplicated here. An entry is dropped once every field + is idle — see ``ServerRole._maybe_prune``. + """ + + # Outbound serve state (PD + symmetric producer side). None until the + # first add_stored_blocks / on_fetch; reset to None on finalize/abort. + outbound: _OutboundRequestState | None = None + # Raw inbound LookupMsgs not yet processed against the ParentManager. + pending_lookups: list[_PendingLookup] = field(default_factory=list) + # Per-LookupMsg state parked with HIT_PENDING / RETRY keys, keyed by + # the (globally unique) lookup_id and re-polled each serve. + lookups: dict[int, _ActiveLookup] = field(default_factory=dict) + # Start time of a pending abort drain (``time.monotonic``); None when + # no abort is in progress. + abort_started_at: float | None = None + # Transfer ids in ``ServerRole._inflight`` for this id. Kept in sync + # via _inflight_add / _inflight_pop so a non-empty set is an exact + # "has any inflight transfer" predicate and the abort drain can + # enumerate this request's transfers without scanning all of + # ``_inflight``. + inflight_tids: set[int] = field(default_factory=set) + + class ServerRole: """Server-side store/serve state machine for one peer session. @@ -157,22 +235,55 @@ class ServerRole: self._transport = transport self._send = send - self._outbound: dict[str, _OutboundRequestState] = {} + # All per-kv_request_id state lives here. Entries are created + # lazily and dropped by _maybe_prune once every field is idle. + self._requests: dict[str, _ServerRequestState] = {} + # kv_request_ids with lookup work (unprocessed pending_lookups or + # parked lookups) for the next serve to visit — the work-list that + # keeps serve_external_requests from scanning every request. + self._serve_pending: set[str] = set() # transfer_id → xfer. Mutate ONLY via _inflight_add / _inflight_pop - # so the per-request count below stays in sync. + # so the per-request inflight_tids stays in sync. self._inflight: dict[int, _InflightXfer] = {} - # kv_request_id → number of entries in _inflight for that id. - # Kept in sync with _inflight; entries that hit zero are removed - # so `kv_request_id in self._inflight_per_req` is an exact - # "has any inflight transfer" predicate (O(1) replacement for - # the previous O(N) scan). - self._inflight_per_req: dict[str, int] = {} self._store_jobs: dict[int, float] = {} # job_id → submitted_at - self._pending_aborts: dict[str, float] = {} # kv_request_id → start # StoreResults queued by _finalize_outbound for the next poll # tick to surface. Mirrors the deferred-result pattern used for # load timeouts. self._pending_store_results: list[StoreResult] = [] + # Synthetic lookup ctxs whose ``on_request_finished`` still needs to + # fire but which were closed outside a serve window (FetchMsg / local + # finish popped their parked lookup). Drained via + # ``parent.on_request_finished`` in ``serve_external_requests``. + self._finished_lookup_ctxs: list[ReqContext] = [] + self._lookup_id_counter: int = 0 + # kv_request_ids with a parked abort awaiting drain — work-list so + # drain_pending_aborts doesn't scan every request each poll tick. + self._parked_aborts: set[str] = set() + + # ------------------------------------------------------------------ + # State helpers + # ------------------------------------------------------------------ + + def _get_or_create_request(self, kv_request_id: str) -> _ServerRequestState: + """Get or create the state entry for a kv_request_id.""" + st = self._requests.get(kv_request_id) + if st is None: + st = _ServerRequestState() + self._requests[kv_request_id] = st + return st + + def _maybe_prune(self, kv_request_id: str) -> None: + """Drop the entry once it holds no live state.""" + st = self._requests.get(kv_request_id) + if ( + st is not None + and st.outbound is None + and not st.inflight_tids + and not st.lookups + and not st.pending_lookups + and st.abort_started_at is None + ): + del self._requests[kv_request_id] # ------------------------------------------------------------------ # Public API @@ -187,19 +298,41 @@ class ServerRole: ) -> None: """New blocks stored locally — match against pending fetch demand.""" self._store_jobs[job_id] = time.monotonic() - req = self._outbound.setdefault(kv_request_id, _OutboundRequestState()) - result = req.add_stored_blocks(keys, block_ids, job_id) - if result.local_idxs and req.demand_received: + st = self._get_or_create_request(kv_request_id) + if st.outbound is None: + st.outbound = _OutboundRequestState() + result = st.outbound.add_stored_blocks(keys, block_ids, job_id) + if result.local_idxs and st.outbound.demand_received: self._submit_transfer(kv_request_id, result) def on_fetch( self, kv_request_id: str, - block_hashes: Sequence[OffloadKey], + keys: Sequence[OffloadKey], block_indexes: Sequence[int], ) -> None: """Handle a FetchMsg from the peer. + In p2p mode FetchMsg is the server-side "request finished" + signal for ``kv_request_id``. Three consequences flow from that: + + - No further ``parent.create_store_job`` will fire for this id: + the request's parked ``lookups`` are popped here (via + ``_finish_inbound_lookups``) before the next + ``serve_external_requests`` runs ``_resolve_pending_lookups``, + so any HIT_PENDING / RETRY key that would otherwise later + promote to HIT and pin a slot is dropped instead. Any raw + not-yet-processed LookupMsg for this id is dropped too. + - All server-side lookup state for the id is cleaned up (the + ``lookups`` entries themselves). + - The synthetic ctx is queued for ``parent.on_request_finished`` + (fired by the next ``serve_external_requests``) so the + TieringManager can release per-lookup bookkeeping. + + The client contract guarantees exactly one FetchMsg per + lookup-touched request — including an empty one when no + blocks end up being fetched. + Raises ``ValueError`` on a duplicate fetch for the same ``kv_request_id``; the coordinator's dispatch loop turns that into a protocol-error disconnect. @@ -208,18 +341,28 @@ class ServerRole: "P2PSession %s: fetch RECEIVED kv_request_id=%s blocks=%d", self._peer_id, kv_request_id, - len(block_hashes), + len(keys), ) - existing = self._outbound.get(kv_request_id) + st = self._requests.get(kv_request_id) + existing = st.outbound if st is not None else None if existing is not None and existing.demand_received: # A second fetch for the same kv_request_id would overwrite # `remaining` and leak inflight bookkeeping. Treat as a # protocol violation. raise ValueError(f"duplicate fetch for kv_request_id={kv_request_id}") - req = self._outbound.setdefault(kv_request_id, _OutboundRequestState()) - result = req.add_fetch_demand(block_hashes, block_indexes) + st = self._get_or_create_request(kv_request_id) + if st.outbound is None: + st.outbound = _OutboundRequestState() + req = st.outbound + result = req.add_fetch_demand(keys, block_indexes) if result.local_idxs: self._submit_transfer(kv_request_id, result) + # Close the peer's request as far as the server's lookup phase + # is concerned: pop parked lookups so no further + # ``parent.create_store_job`` fires for this id, and queue their + # ctxs for ``parent.on_request_finished``. Done before the + # finalize path below so bookkeeping releases in-order. + self._finish_inbound_lookups(kv_request_id) # Prefiller-first mode: finish_request may have run before # fetch arrived. If so, finalize once we know what was # demanded — fully satisfied → success, else early-fail. @@ -230,9 +373,9 @@ class ServerRole: """Handle an AbortFetchMsg from the peer.""" # Abort for an unknown id may be a benign race/duplicate or a # real protocol violation; we don't track completed ids, so warn. - if kv_request_id not in self._outbound and not self._has_inflight_for( - kv_request_id - ): + st = self._requests.get(kv_request_id) + has_outbound = st is not None and st.outbound is not None + if not has_outbound and not self._has_inflight_for(kv_request_id): logger.warning( "P2PSession %s: abort_fetch for unknown kv_request_id=%s " "(no outbound or inflight state); benign race or stale", @@ -242,9 +385,285 @@ class ServerRole: # Idempotent: receiving AbortFetchMsg again before we've sent the # ack just triggers another drain attempt without resetting the # deadline. - self._pending_aborts.setdefault(kv_request_id, time.monotonic()) + st = self._get_or_create_request(kv_request_id) + if st.abort_started_at is None: + st.abort_started_at = time.monotonic() + self._parked_aborts.add(kv_request_id) self._drain_abort(kv_request_id) + def on_lookup( + self, + kv_request_id: str, + keys: Sequence[OffloadKey], + ) -> None: + """Enqueue a LookupMsg from a symmetric-P2P consumer. + + Dispatch runs during ``session.poll()`` where the + :class:`ParentManager` handle is not available, so this only + records the raw request. It is resolved — querying the tiering + manager and emitting the aggregated ``LookupRespMsg`` — by the + next :meth:`serve_external_requests`, the sole window in which + parent calls are valid. + """ + logger.debug( + "P2P LOOKUP server %s: RECV LookupMsg kv_request_id=%s keys=%d", + self._peer_id, + kv_request_id, + len(keys), + ) + self._get_or_create_request(kv_request_id).pending_lookups.append( + _PendingLookup(keys=list(keys), enqueued_at=time.monotonic()) + ) + self._serve_pending.add(kv_request_id) + + def serve_external_requests(self, parent: ParentManager) -> None: + """Resolve inbound peer lookups against the tiering manager. + + Called once per scheduler step with a ``parent`` handle valid + only for this call. Drains newly-enqueued LookupMsgs, re-polls + any parked HIT_PENDING / RETRY keys, and releases the + bookkeeping for lookups closed since the last serve. + """ + for kv_request_id in list(self._serve_pending): + st = self._requests.get(kv_request_id) + if st is None: + self._serve_pending.discard(kv_request_id) + continue + if st.pending_lookups: + pending = st.pending_lookups + st.pending_lookups = [] + for pl in pending: + self._process_inbound_lookup( + kv_request_id, pl.keys, pl.enqueued_at, parent + ) + self._resolve_pending_lookups(kv_request_id, parent) + st = self._requests.get(kv_request_id) + if st is None or (not st.pending_lookups and not st.lookups): + self._serve_pending.discard(kv_request_id) + self._maybe_prune(kv_request_id) + + if self._finished_lookup_ctxs: + for ctx in self._finished_lookup_ctxs: + parent.on_request_finished(ctx) + self._finished_lookup_ctxs = [] + + def _poll_lookup_keys( + self, + lookup: _ActiveLookup, + keys: Iterable[OffloadKey], + parent: ParentManager, + ) -> list[OffloadKey]: + """Poll ``keys`` against the tiering manager and pin any HITs. + + For each key not already definitively resolved, query + ``parent.lookup`` and record the outcome on ``lookup``: HIT / MISS + land in ``resolved`` and clear ``pending``; HIT_PENDING / RETRY park + in ``pending`` for a later serve. Newly-HIT keys are pinned in one + batch via :meth:`_pin_and_register_hits` and also returned (for + caller logging). + + Shared by the first-sighting pass (:meth:`_process_inbound_lookup`) + and the re-poll pass (:meth:`_resolve_pending_lookups`); callers pass + a de-duplicated ``keys`` collection. + """ + new_hits: list[OffloadKey] = [] + for h in keys: + if h in lookup.resolved: + continue + result = parent.lookup(h, lookup.ctx) + if result is LookupResult.HIT: + new_hits.append(h) + lookup.resolved[h] = True + lookup.pending.discard(h) + elif result is LookupResult.MISS: + lookup.resolved[h] = False + lookup.pending.discard(h) + else: + lookup.pending.add(h) + if new_hits: + self._pin_and_register_hits( + lookup.kv_request_id, new_hits, lookup.ctx, parent + ) + return new_hits + + def _process_inbound_lookup( + self, + kv_request_id: str, + keys: list[OffloadKey], + enqueued_at: float, + parent: ParentManager, + ) -> None: + """Resolve one enqueued LookupMsg against ``parent``. + + For each key, query the tiering manager via ``parent.lookup`` + and pin any HITs immediately via ``parent.create_store_job`` + (plumbed into the existing ``add_stored_blocks`` matching path so + the eventual FetchMsg finds them). HIT_PENDING / RETRY keys park + in :class:`_ActiveLookup` for re-polling by + :meth:`_resolve_pending_lookups`. + + The outbound ``LookupRespMsg`` is deferred until every key has + settled to HIT or MISS (or the batch-level ``deadline`` fires, + forcing any stragglers to MISS). One LookupRespMsg goes out per + LookupMsg — carrying every key in wire order — after which + ``parent.on_request_finished`` fires and the entry is dropped. + """ + self._lookup_id_counter += 1 + lookup_id = self._lookup_id_counter + ctx = ReqContext(req_id=f"p2p:{self._peer_id}:{kv_request_id}:lu{lookup_id}") + lookup = _ActiveLookup( + lookup_id=lookup_id, + kv_request_id=kv_request_id, + ctx=ctx, + keys=list(keys), + deadline=enqueued_at + _LOOKUP_PENDING_TIMEOUT_S, + ) + + # Open per-request bookkeeping for this synthetic ctx before the + # first lookup; released by ``on_request_finished`` once every + # key has settled. + parent.on_new_request(ctx) + + # dict.fromkeys de-duplicates keys within the LookupMsg while + # preserving wire order — each unique key is polled once. + hit_keys = self._poll_lookup_keys(lookup, dict.fromkeys(lookup.keys), parent) + + logger.debug( + "P2P LOOKUP server %s: RESOLVED kv_request_id=%s hits=%d misses=%d " + "pending=%d", + self._peer_id, + kv_request_id, + len(hit_keys), + sum(1 for v in lookup.resolved.values() if not v), + len(lookup.pending), + ) + + if lookup.pending: + self._get_or_create_request(kv_request_id).lookups[lookup_id] = lookup + else: + # Every key resolved on first sight — emit the aggregated + # response now and close the synthetic request. + self._finalize_lookup(lookup, parent) + + def _pin_and_register_hits( + self, + kv_request_id: str, + keys: list[OffloadKey], + ctx: ReqContext, + parent: ParentManager, + ) -> None: + """Pin primary slots for HIT keys and feed them into the + existing ``add_stored_blocks`` matching path. + + Caller has already confirmed every key is HIT (single-threaded + scheduler ⇒ no eviction race), so the JobMetadata returned by + ``parent.create_store_job`` carries parallel ``keys``/``block_ids`` + of length ``len(keys)``. + """ + meta = parent.create_store_job(keys, ctx) + self.add_stored_blocks( + kv_request_id, + list(meta.keys), + list(meta.block_ids), + meta.job_id, + ) + + def _resolve_pending_lookups( + self, kv_request_id: str, parent: ParentManager + ) -> None: + """Re-poll a request's deferred LookupMsg keys; finalize when ready. + + Walks every parked :class:`_ActiveLookup` for ``kv_request_id`` and + re-calls ``parent.lookup`` per still-pending key, moving HIT/MISS + results into ``resolved``. If ``deadline`` has passed, remaining + ``pending`` keys are force-resolved to MISS so the consumer + can fall back instead of waiting on a stuck producer. Newly-HIT + keys are pinned via ``parent.create_store_job`` in one call per + affected lookup. Lookups whose ``pending`` empties out get their + aggregated LookupRespMsg sent by :meth:`_finalize_lookup`. + """ + st = self._requests.get(kv_request_id) + if st is None or not st.lookups: + return + now = time.monotonic() + finished_lookups: list[int] = [] + for lookup_id, lookup in st.lookups.items(): + self._poll_lookup_keys(lookup, list(lookup.pending), parent) + + if lookup.pending and now >= lookup.deadline: + for h in lookup.pending: + lookup.resolved[h] = False + lookup.pending.clear() + + if not lookup.pending: + finished_lookups.append(lookup_id) + + for lookup_id in finished_lookups: + lookup = st.lookups.pop(lookup_id) + self._finalize_lookup(lookup, parent) + + def _finalize_lookup(self, lookup: _ActiveLookup, parent: ParentManager) -> None: + """Emit the aggregated LookupRespMsg and close the synthetic request. + + Called once per lookup when ``pending`` is empty — either every + key resolved to HIT or MISS, or the deadline forced remaining + stragglers to MISS. Preserves the wire order of the inbound + LookupMsg so the client can zip keys and hits positionally. + """ + if lookup.keys: + hits = [lookup.resolved[h] for h in lookup.keys] + n_hit = sum(1 for v in hits if v) + logger.debug( + "P2P LOOKUP server %s: SEND LookupRespMsg kv_request_id=%s " + "keys=%d hits=%d misses=%d", + self._peer_id, + lookup.kv_request_id, + len(lookup.keys), + n_hit, + len(hits) - n_hit, + ) + self._send( + { + TYPE_KEY: LookupRespMsg.TYPE, + LookupRespMsg.KV_REQUEST_ID: lookup.kv_request_id, + LookupRespMsg.KEYS: list(lookup.keys), + LookupRespMsg.HITS: hits, + } + ) + parent.on_request_finished(lookup.ctx) + + def _finish_inbound_lookups(self, kv_request_id: str) -> None: + """Close the server-side lookup phase for ``kv_request_id``. + + Pops every parked ``_ActiveLookup`` for this id (so + ``_resolve_pending_lookups`` cannot promote a HIT_PENDING / + RETRY key into a fresh ``parent.create_store_job`` after this + point) and queues each ``lookup.ctx`` for + ``parent.on_request_finished`` (fired by the next + ``serve_external_requests``, since no parent handle is available + during dispatch) so the TieringManager can release per-lookup + bookkeeping. Any still-unprocessed raw LookupMsg for this id is + dropped — it never got ``on_new_request``, so nothing is owed. + The aggregated LookupRespMsg is skipped — the client already + knows the request is over (it just sent a terminal FetchMsg, or + is finishing locally). + + Called on the two events that mean "no more lookup traffic for + ``kv_request_id`` is expected on this session": the terminal + FetchMsg from the peer (client contract: exactly one FetchMsg + per lookup-touched request, even if empty), and a local + ``finish``. Whichever fires second is a no-op. + """ + st = self._requests.get(kv_request_id) + if st is None: + return + st.pending_lookups.clear() + for lookup in st.lookups.values(): + self._finished_lookup_ctxs.append(lookup.ctx) + st.lookups.clear() + self._serve_pending.discard(kv_request_id) + self._maybe_prune(kv_request_id) + def finish(self, kv_request_id: str) -> None: """Mark an outbound request finishing. @@ -253,13 +672,23 @@ class ServerRole: peer to stop waiting (TransferDoneMsg success=False) instead of letting it hit _LOAD_TIMEOUT_S. + Also drops any in-flight lookups for this kv_request_id and + queues their ctxs for ``parent.on_request_finished`` so the + TieringManager can release per-request bookkeeping. (For + symmetric P2P this path is rarely hit since the producer has no + local request lifecycle for the consumer's id; this cleanup is + mostly active on the PD side.) + If the decoder hasn't sent fetch yet (no demand received), defer — on_fetch will finalize once demand arrives. If inflight transfers exist for this id, defer — the last completing transfer in collect_results will fire the message. """ - req = self._outbound.get(kv_request_id) + self._finish_inbound_lookups(kv_request_id) + + st = self._requests.get(kv_request_id) + req = st.outbound if st is not None else None if req is None: return req.finishing = True @@ -268,14 +697,19 @@ class ServerRole: if self._has_inflight_for(kv_request_id): return # Remaining > 0 here: if it had hit 0, the poll-done success - # branch would have already popped _outbound and we'd have + # branch would have already cleared outbound and we'd have # returned at `req is None` above. Helper derives success from # remaining and emits StoreResult(success=False) for any # leftover pending jobs. self._finalize_outbound(kv_request_id) def collect_results(self) -> list[StoreResult]: - """Drain timeouts, deferred results, and transport completions.""" + """Drain timeouts, deferred results, and transport completions. + + Inbound LookupMsg resolution (including re-polling HIT_PENDING / + RETRY keys) is NOT done here — it runs in + ``serve_external_requests`` where the ParentManager is available. + """ results: list[StoreResult] = self._timeout_pending_store_jobs() if self._pending_store_results: @@ -306,15 +740,9 @@ class ServerRole: tid, ) continue - req = self._outbound.get(xfer.kv_request_id) - for job_id in xfer.job_ids: - if self._store_jobs.pop(job_id, None) is None: - # Already reported (timeout, cancellation, etc.) — - # don't double-emit a contradictory success result. - continue - results.append(StoreResult(job_id=job_id, success=True)) - if req is not None: - req.pending_job_ids.discard(job_id) + results.extend(self._settle_xfer_jobs(xfer, success=True)) + st = self._requests.get(xfer.kv_request_id) + req = st.outbound if st is not None else None if req is not None and req.demand_received: req.remaining -= xfer.block_count assert req.remaining >= 0, ( @@ -324,6 +752,7 @@ class ServerRole: self._finalize_outbound(xfer.kv_request_id, success=True) elif req.finishing and not self._has_inflight_for(xfer.kv_request_id): self._finalize_outbound(xfer.kv_request_id, success=False) + self._maybe_prune(xfer.kv_request_id) failed_kv_request_ids: set[str] | None = None for tid in poll_result.failed: @@ -341,16 +770,11 @@ class ServerRole: if failed_kv_request_ids is None: failed_kv_request_ids = set() failed_kv_request_ids.add(xfer.kv_request_id) - req_for_xfer = self._outbound.get(xfer.kv_request_id) - for job_id in xfer.job_ids: - if self._store_jobs.pop(job_id, None) is None: - # Already reported (timeout, cancellation, etc.) — - # don't double-emit. - continue - results.append(StoreResult(job_id=job_id, success=False)) - if req_for_xfer is not None: - req_for_xfer.pending_job_ids.discard(job_id) - req = self._outbound.pop(xfer.kv_request_id, None) + results.extend(self._settle_xfer_jobs(xfer, success=False)) + st = self._requests.get(xfer.kv_request_id) + req = st.outbound if st is not None else None + if st is not None: + st.outbound = None if req is not None and req.demand_received: self._send( { @@ -359,6 +783,7 @@ class ServerRole: TransferDoneMsg.SUCCESS: False, } ) + self._maybe_prune(xfer.kv_request_id) # Cancel other inflight for the same failed kv_request_ids if failed_kv_request_ids: @@ -370,6 +795,8 @@ class ServerRole: for tid in ids_to_cancel: self._inflight_pop(tid) self._transport.cancel(ids_to_cancel) + for kv_request_id in failed_kv_request_ids: + self._maybe_prune(kv_request_id) return results @@ -384,54 +811,90 @@ class ServerRole: def drain_pending_aborts(self) -> None: """Re-attempt every parked abort once per poll tick.""" - if not self._pending_aborts: - return - for kv_request_id in list(self._pending_aborts): + for kv_request_id in list(self._parked_aborts): self._drain_abort(kv_request_id) - def close(self) -> list[int]: - """Tear down. Cancels inflight, returns failed store job ids.""" + def close(self) -> tuple[list[int], list[ReqContext]]: + """Tear down. Cancels inflight. + + Returns ``(failed_store_job_ids, failed_serves)`` where + ``failed_serves`` are synthetic lookup ctxs still owing a + ``parent.on_request_finished``. The session is going away with no + parent handle in hand, so the manager flushes these in its next + ``serve_external_requests``. + """ failed_stores = list(self._store_jobs.keys()) self._store_jobs.clear() if self._inflight: self._transport.cancel(list(self._inflight.keys())) self._inflight.clear() - self._inflight_per_req.clear() - self._outbound.clear() - self._pending_aborts.clear() self._pending_store_results.clear() - return failed_stores + # Surface every synthetic ctx still owing on_request_finished so + # the manager can release the TieringManager's per-request + # bookkeeping: parked lookups plus any already queued from a + # FetchMsg / finish that closed them before this teardown. + failed_serves = [ + lu.ctx for st in self._requests.values() for lu in st.lookups.values() + ] + failed_serves.extend(self._finished_lookup_ctxs) + self._requests.clear() + self._serve_pending.clear() + self._parked_aborts.clear() + self._finished_lookup_ctxs.clear() + return failed_stores, failed_serves + + @property + def has_inflight_transfers(self) -> bool: + """True if any outbound store transfer is still in flight.""" + return bool(self._inflight) # ------------------------------------------------------------------ # Internal — inflight bookkeeping # ------------------------------------------------------------------ def _has_inflight_for(self, kv_request_id: str) -> bool: - return kv_request_id in self._inflight_per_req + st = self._requests.get(kv_request_id) + return st is not None and bool(st.inflight_tids) def _inflight_add(self, tid: int, xfer: _InflightXfer) -> None: - """Insert an inflight transfer and bump the per-request count.""" + """Insert an inflight transfer and record it on the request.""" self._inflight[tid] = xfer - self._inflight_per_req[xfer.kv_request_id] = ( - self._inflight_per_req.get(xfer.kv_request_id, 0) + 1 - ) + self._get_or_create_request(xfer.kv_request_id).inflight_tids.add(tid) def _inflight_pop(self, tid: int) -> _InflightXfer | None: - """Pop an inflight transfer and decrement the per-request count. + """Pop an inflight transfer and drop it from the request's set. - Removes the per-request entry once the count hits zero so the - dict stays bounded and `_has_inflight_for` remains exact. + Callers are responsible for the ``_maybe_prune`` that may follow once + the request's other state has also cleared. """ xfer = self._inflight.pop(tid, None) if xfer is None: return None - new_count = self._inflight_per_req.get(xfer.kv_request_id, 0) - 1 - if new_count > 0: - self._inflight_per_req[xfer.kv_request_id] = new_count - else: - self._inflight_per_req.pop(xfer.kv_request_id, None) + st = self._requests.get(xfer.kv_request_id) + if st is not None: + st.inflight_tids.discard(tid) return xfer + def _settle_xfer_jobs( + self, xfer: _InflightXfer, success: bool + ) -> list[StoreResult]: + """Emit StoreResults for a completed transfer's store jobs. + + Pops each attached job from ``_store_jobs`` and clears it from the + request's pending set. A job already popped (via timeout, cancel, + etc.) is skipped so we never double-emit a contradictory result. + """ + results: list[StoreResult] = [] + st = self._requests.get(xfer.kv_request_id) + req = st.outbound if st is not None else None + for job_id in xfer.job_ids: + if self._store_jobs.pop(job_id, None) is None: + continue + results.append(StoreResult(job_id=job_id, success=success)) + if req is not None: + req.pending_job_ids.discard(job_id) + return results + # ------------------------------------------------------------------ # Internal — finalize / abort drain # ------------------------------------------------------------------ @@ -452,9 +915,10 @@ class ServerRole: The same flag is used for both the peer's TransferDoneMsg and the StoreResult(s) emitted for any leftover pending job_ids. """ - req = self._outbound.pop(kv_request_id, None) - if req is None: - return + st = self._requests[kv_request_id] + assert st.outbound is not None + req = st.outbound + st.outbound = None if success is None: success = req.remaining == 0 for job_id in req.pending_job_ids: @@ -469,6 +933,7 @@ class ServerRole: TransferDoneMsg.SUCCESS: success, } ) + self._maybe_prune(kv_request_id) def _drain_abort(self, kv_request_id: str) -> None: """One drain attempt for a pending abort. @@ -479,21 +944,15 @@ class ServerRole: inflight, or after ``_CANCEL_DRAIN_TIMEOUT_S`` falls back to ``mode="immediate"`` and acks anyway. """ - self._outbound.pop(kv_request_id, None) - ids = [ - tid - for tid, xfer in self._inflight.items() - if xfer.kv_request_id == kv_request_id - ] + st = self._requests[kv_request_id] + st.outbound = None + ids = list(st.inflight_tids) if not ids: self._finalize_abort(kv_request_id) return - started_at = self._pending_aborts.get(kv_request_id) - expired = ( - started_at is not None - and time.monotonic() - started_at >= _CANCEL_DRAIN_TIMEOUT_S - ) + assert st.abort_started_at is not None + expired = time.monotonic() - st.abort_started_at >= _CANCEL_DRAIN_TIMEOUT_S if expired: for tid in ids: self._inflight_pop(tid) @@ -521,13 +980,16 @@ class ServerRole: self._finalize_abort(kv_request_id) def _finalize_abort(self, kv_request_id: str) -> None: - self._pending_aborts.pop(kv_request_id, None) + st = self._requests[kv_request_id] + st.abort_started_at = None + self._parked_aborts.discard(kv_request_id) self._send( { TYPE_KEY: AbortAckMsg.TYPE, AbortAckMsg.KV_REQUEST_ID: kv_request_id, } ) + self._maybe_prune(kv_request_id) # ------------------------------------------------------------------ # Internal — transfers and store-job timeouts @@ -579,7 +1041,8 @@ class ServerRole: # nothing else is in flight, finalize now so the peer # and the local store jobs don't wait for finish_request # or for _STORE_TIMEOUT_S / _LOAD_TIMEOUT_S. - req = self._outbound.get(kv_request_id) + st = self._requests.get(kv_request_id) + req = st.outbound if st is not None else None if req is not None: req.finishing = True if not self._has_inflight_for(kv_request_id): diff --git a/vllm/v1/kv_offload/tiering/p2p/session/session.py b/vllm/v1/kv_offload/tiering/p2p/session/session.py index fc5f1feab5..7d19913b0a 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/session.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/session.py @@ -34,6 +34,8 @@ from vllm.v1.kv_offload.tiering.p2p.session.protocol import ( ConnectMsg, DisconnectMsg, FetchMsg, + LookupMsg, + LookupRespMsg, TransferDoneMsg, ) from vllm.v1.kv_offload.tiering.p2p.session.server import ( @@ -42,7 +44,8 @@ from vllm.v1.kv_offload.tiering.p2p.session.server import ( ) if TYPE_CHECKING: - from vllm.v1.kv_offload.tiering.base import JobId + from vllm.v1.kv_offload.base import ReqContext + from vllm.v1.kv_offload.tiering.base import JobId, ParentManager from vllm.v1.kv_offload.tiering.p2p.data import DataTransport logger = init_logger(__name__) @@ -71,6 +74,22 @@ class SessionPollResult(NamedTuple): new_fetch_ids: list[str] +class SessionCloseResult(NamedTuple): + """Result of tearing down a P2PSession. + + `failed_jobs`/`failed_stores` are the in-flight jobs (client loads / + server stores) the manager must mark failed. `failed_req_ids` is every + client-side kv_request_id whose lookup() must fail (in-flight loads plus + unresolved probes); `failed_serves` is the server-side lookup state the + dead peer can no longer resolve. + """ + + failed_jobs: list[int] # client load job_ids + failed_req_ids: list[str] # client kv_request_ids (loads + probes) + failed_stores: list[int] # server store job_ids + failed_serves: list[ReqContext] # server-side lookup ctxs needing release + + class P2PSession: """Bidirectional session — coordinator over ClientRole + ServerRole. @@ -93,12 +112,14 @@ class P2PSession: local_id: str, transport: DataTransport, local_block_len: int, + local_hash_seed: str, conn: ControlConnection | None = None, ) -> None: self.peer_id = peer_id self._local_id = local_id self._transport = transport self._local_block_len = local_block_len + self._local_hash_seed = local_hash_seed self._conn: ControlConnection | None = None self._send_ready = False # True after the peer acked our ConnectMsg @@ -115,7 +136,11 @@ class P2PSession: self._new_fetch_ids: list[str] = [] self._client = ClientRole(peer_id=peer_id, send=self._send) - self._server = ServerRole(peer_id=peer_id, transport=transport, send=self._send) + self._server = ServerRole( + peer_id=peer_id, + transport=transport, + send=self._send, + ) if conn is not None: self.attach_connection(conn) @@ -139,6 +164,11 @@ class P2PSession: """True after the peer acked our ConnectMsg (we may send freely).""" return self._send_ready + @property + def has_pending_work(self) -> bool: + """True while inbound loads or outbound transfers are outstanding.""" + return self._client.has_active_loads or self._server.has_inflight_transfers + # ------------------------------------------------------------------ # Connection lifecycle # ------------------------------------------------------------------ @@ -162,7 +192,7 @@ class P2PSession: self, job_id: JobId, kv_request_id: str, - keys: Sequence[bytes], + keys: Sequence[OffloadKey], block_ids: Sequence[int], ) -> None: """Send fetch to the peer.""" @@ -183,13 +213,40 @@ class P2PSession: def finish_request(self, kv_request_id: str) -> None: """Called when the request is finishing locally. - Cancels any inbound load (client role) and finalizes any - outbound serving (server role) for this id. Roles that aren't - active for this id are silent no-ops. + Finishes the client role (aborts any inbound load and drops any + pending symmetric-P2P lookup state) and finalizes any outbound + serving (server role) for this id. Roles that aren't active for + this id are silent no-ops. """ - self._client.cancel(kv_request_id) + self._client.finish(kv_request_id) self._server.finish(kv_request_id) + def register_lookup(self, kv_request_id: str, key: bytes) -> bool | None: + """Register or resolve one (kv_request_id, key) probe. + + Called from the manager's lookup() for symmetric-P2P consumers + (``remote_kv_source`` sub-dict in kv_transfer_params). See + ``ClientRole.register_lookup`` for the state-machine contract. + """ + return self._client.register_lookup(kv_request_id, key) + + def flush_pending_lookups(self) -> None: + """Flush any aggregated symmetric-P2P lookups for this peer. + + Called once per scheduler step from the manager's + ``on_schedule_end()``. Send-gating is handled inside the + client's ``_send`` callback (queues until ConnectAckMsg). + """ + self._client.flush_pending_lookups() + + def serve_external_requests(self, parent: ParentManager) -> None: + """Resolve inbound peer lookups against the tiering manager. + + Delegates to the server role; the ``parent`` handle is valid + only for the duration of this call. + """ + self._server.serve_external_requests(parent) + def poll(self) -> SessionPollResult: """Process incoming messages, drive transfers, apply timeouts.""" if self._conn is None: @@ -214,14 +271,22 @@ class P2PSession: loads=loads, stores=stores, new_fetch_ids=new_fetch_ids ) - def close(self) -> tuple[list[tuple[int, str]], list[int]]: - """Shut down. Returns (failed_loads, failed_stores). + def close(self) -> SessionCloseResult: + """Shut down. - failed_loads: list of (job_id, kv_request_id) pairs. - failed_stores: list of job_ids. + failed_jobs: client load job_ids to fail. + failed_req_ids: client kv_request_ids to fail — in-flight loads plus + requests with an unresolved symmetric-P2P probe toward the + now-dead peer. The manager fails these so the consumer's lookup() + falls back to local prefill instead of deferring forever on an + answer that can never arrive. + failed_stores: server store job_ids to fail. + failed_serves: synthetic lookup ctxs still owing + ``parent.on_request_finished`` (the manager flushes these on + its next ``serve_external_requests``). """ - failed_loads = self._client.close() - failed_stores = self._server.close() + client_result = self._client.close() + failed_stores, failed_serves = self._server.close() if self._conn is not None: with contextlib.suppress(Exception): @@ -229,7 +294,12 @@ class P2PSession: self._conn.close() self._conn = None - return failed_loads, failed_stores + return SessionCloseResult( + failed_jobs=client_result.failed_jobs, + failed_req_ids=client_result.failed_req_ids, + failed_stores=failed_stores, + failed_serves=failed_serves, + ) # ------------------------------------------------------------------ # Message dispatch @@ -299,9 +369,9 @@ class P2PSession: elif msg_type == FetchMsg.TYPE: FetchMsg.validate(msg) kv_request_id = msg[FetchMsg.KV_REQUEST_ID] - block_hashes = [ + keys = [ OffloadKey(bh if isinstance(bh, bytes) else bytes(bh)) - for bh in msg[FetchMsg.BLOCK_HASHES] + for bh in msg[FetchMsg.KEYS] ] block_indexes = msg[FetchMsg.BLOCK_INDEXES] # Run the server-role state machine inline as today — @@ -310,7 +380,7 @@ class P2PSession: # the manager (after poll() returns) can replay any parked # submit_store batches; their add_stored_blocks calls hit # the demand recorded here and submit transfers immediately. - self._server.on_fetch(kv_request_id, block_hashes, block_indexes) + self._server.on_fetch(kv_request_id, keys, block_indexes) self._new_fetch_ids.append(kv_request_id) elif msg_type == AbortFetchMsg.TYPE: AbortFetchMsg.validate(msg) @@ -324,6 +394,23 @@ class P2PSession: elif msg_type == AbortAckMsg.TYPE: AbortAckMsg.validate(msg) self._client.on_abort_ack(msg[AbortAckMsg.KV_REQUEST_ID]) + elif msg_type == LookupMsg.TYPE: + LookupMsg.validate(msg) + kv_request_id = msg[LookupMsg.KV_REQUEST_ID] + keys = [ + OffloadKey(bh if isinstance(bh, bytes) else bytes(bh)) + for bh in msg[LookupMsg.KEYS] + ] + self._server.on_lookup(kv_request_id, keys) + elif msg_type == LookupRespMsg.TYPE: + LookupRespMsg.validate(msg) + kv_request_id = msg[LookupRespMsg.KV_REQUEST_ID] + keys = [ + OffloadKey(bh if isinstance(bh, bytes) else bytes(bh)) + for bh in msg[LookupRespMsg.KEYS] + ] + hits = msg[LookupRespMsg.HITS] + self._client.on_lookup_resp(kv_request_id, keys, hits) elif msg_type == DisconnectMsg.TYPE: if self._conn is not None: self._conn.mark_dead() @@ -361,6 +448,12 @@ class P2PSession: f"config fingerprint mismatch from {self.peer_id}: " f"remote={remote_fp!r}, local={local_fp!r}" ) + if msg[ConnectMsg.HASH_SEED] != self._local_hash_seed: + raise ValueError( + f"PYTHONHASHSEED mismatch from {self.peer_id}: " + f"remote={msg[ConnectMsg.HASH_SEED]!r}, " + f"local={self._local_hash_seed!r}" + ) self._transport.add_remote_peer( self.peer_id, agent_metadata=msg[ConnectMsg.AGENT_METADATA], @@ -410,6 +503,7 @@ class P2PSession: ConnectMsg.NUM_BLOCKS: self._transport.num_blocks, ConnectMsg.BLOCK_LEN: self._transport.block_len, ConnectMsg.CONFIG_FINGERPRINT: self._transport.config_fingerprint, + ConnectMsg.HASH_SEED: self._local_hash_seed, } ) @@ -433,8 +527,15 @@ class P2PSession: self._conn.send(msg) logger.debug("P2PSession %s: sent %s", self.peer_id, msg.get(TYPE_KEY)) except Exception: + # A send failure means the connection is broken. Swallowing it + # silently strands every in-flight lookup/load toward this peer: + # the session stays alive, is never reaped, and the consumer's + # lookup() keeps returning RETRY until the HTTP client times out. + # Mark the connection dead so the manager reaps the session on + # its next poll and surfaces the stranded work as failures. logger.warning( - "P2PSession %s: failed to send %s", + "P2PSession %s: send of %s failed — marking connection dead", self.peer_id, msg.get(TYPE_KEY), ) + self._conn.mark_dead()