| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Scheduler-side wire types for the ECCPUConnector. |
| 4 | |
| 5 | Two peer-to-peer message types exchanged between consumer and producer over |
| 6 | ZMQ (XferReq, XferAck) plus a peer-compatibility fingerprint |
| 7 | (compute_ec_compatibility_hash) sent on every XferReq so a producer refuses |
| 8 | to serve a mismatched peer. |
| 9 | """ |
| 10 | |
| 11 | import enum |
| 12 | import hashlib |
| 13 | import json |
| 14 | |
| 15 | import msgspec |
| 16 | |
| 17 | # Bump when the on-wire shape (XferReq/XferAck) changes in a |
| 18 | # backward-incompatible way. Peers reject mismatched versions. |
| 19 | EC_CONNECTOR_VERSION: int = 1 |
| 20 | |
| 21 | |
| 22 | class XferStatus(enum.IntEnum): |
| 23 | """Outcome the producer reports for an XferReq.""" |
| 24 | |
| 25 | OK = 0 |
| 26 | NACK_MISSING = 1 # producer no longer holds the encoding |
| 27 | NACK_INCOMPAT = 2 # peer-compatibility hash mismatch |
| 28 | NACK_VERSION = 3 # connector wire-version mismatch |
| 29 | NACK_INTERNAL = 4 # producer hit an unexpected error |
| 30 | NACK_NOT_READY = 5 # announced, but its GPU->mmap save has not landed yet |
| 31 | |
| 32 | |
| 33 | # The producer expects a retryable condition to clear on its own, so the |
| 34 | # consumer re-requests the read rather than giving up on the encoding. |
| 35 | RETRYABLE_NACKS = frozenset((XferStatus.NACK_NOT_READY,)) |
| 36 | # Ordinary outcomes rather than signs of a broken peer. |
| 37 | EXPECTED_NACKS = RETRYABLE_NACKS | {XferStatus.NACK_MISSING} |
| 38 | |
| 39 | |
| 40 | class XferReq( # type: ignore[call-arg] |
| 41 | msgspec.Struct, |
| 42 | tag="req", |
| 43 | omit_defaults=True, |
| 44 | ): |
| 45 | """Consumer → Producer: I want to READ mm_hash; pin it and tell me where.""" |
| 46 | |
| 47 | mm_hash: str |
| 48 | compatibility_hash: str |
| 49 | # Consumer's session identity. Together with mm_hash it uniquely identifies |
| 50 | # this transfer on both sides: the producer keys its active xfers by |
| 51 | # (session_id, mm_hash), and the NIXL notif_msg carries the same pair. |
| 52 | session_id: str = "" |
| 53 | connector_version: int = EC_CONNECTOR_VERSION |
| 54 | |
| 55 | |
| 56 | class XferAck( # type: ignore[call-arg] |
| 57 | msgspec.Struct, |
| 58 | tag="ack", |
| 59 | omit_defaults=True, |
| 60 | ): |
| 61 | """Producer → Consumer: response to an XferReq. |
| 62 | |
| 63 | On status == OK: producer has pinned source blocks and returns their |
| 64 | indices plus fresh NIXL metadata for the consumer to register and READ. |
| 65 | On any NACK_*: optional fields are empty; consumer falls back to local encode. |
| 66 | """ |
| 67 | |
| 68 | mm_hash: str |
| 69 | status: XferStatus |
| 70 | # Consumer's session_id echoed back from XferReq. Together with mm_hash it |
| 71 | # identifies the exact transfer; the consumer uses both to build notif_msg. |
| 72 | session_id: str = "" |
| 73 | src_block_indices: list[int] = [] |
| 74 | agent_metadata: bytes = b"" |
| 75 | mem_descriptor: bytes = b"" |
| 76 | |
| 77 | |
| 78 | def compute_ec_compatibility_hash( |
| 79 | vllm_version: str, |
| 80 | model: str, |
| 81 | dtype: str, |
| 82 | block_size_bytes: int, |
| 83 | ) -> str: |
| 84 | """Peer-compatibility fingerprint over the four factors that determine |
| 85 | byte-layout compatibility. Producer NACKs any XferReq whose hash differs. |
| 86 | """ |
| 87 | factors = json.dumps( |
| 88 | [vllm_version, model, dtype, block_size_bytes], |
| 89 | separators=(",", ":"), |
| 90 | ) |
| 91 | return hashlib.sha256(factors.encode("utf-8")).hexdigest() |