[KVOffload][P2P] Generic P2P secondary tier: peer lookup and serving via ParentManager (#48021)
Signed-off-by: Liran Schour <lirans@il.ibm.com> Signed-off-by: liranschour <liranschour@users.noreply.github.com> Co-authored-by: Or Ozeri <or@ozery.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>
This commit is contained in:
liranschour
co-authored by
Or Ozeri
Or Ozeri
Or Ozeri
Or Ozeri
parent
21fd9e85a0
commit
da3a252fd1
@@ -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": "<unique-transfer-id>"}}
|
||||
|
||||
# Decode consumer — pull KV from a specific prefiller (classic P/D)
|
||||
kv_transfer_params = {
|
||||
"remote_prefiller": {
|
||||
"kv_request_id": "<unique-transfer-id>",
|
||||
"remote_host": "<prefiller-node-ip>",
|
||||
"remote_port": 5710,
|
||||
}
|
||||
}
|
||||
|
||||
# P2P consumer — pull whatever the source already has cached
|
||||
kv_transfer_params = {
|
||||
"remote_kv_source": {
|
||||
"kv_request_id": "<unique-transfer-id>",
|
||||
"remote_host": "<source-node-ip>",
|
||||
"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.
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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",
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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):
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user