| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Metadata exchanged by the Mooncake encoder-cache connector.""" |
| 4 | |
| 5 | from __future__ import annotations |
| 6 | |
| 7 | from dataclasses import dataclass, field |
| 8 | |
| 9 | from vllm.distributed.ec_transfer.ec_connector.base import ( |
| 10 | ECConnectorMetadata, |
| 11 | ECConnectorWorkerMetadata, |
| 12 | ) |
| 13 | from vllm.distributed.ec_transfer.ec_connector.mooncake_store_embedding.data import ( |
| 14 | TensorSpec, |
| 15 | ) |
| 16 | |
| 17 | |
| 18 | @dataclass |
| 19 | class ECMooncakeLoadSpec: |
| 20 | """Describe a remote reservation or resident allocation to load.""" |
| 21 | |
| 22 | mm_hash: str |
| 23 | nbytes: int |
| 24 | shape: tuple[int, ...] |
| 25 | dtype: str |
| 26 | transfer_id: str |
| 27 | # The consumer pool still holds this item, so the load is a local handoff: |
| 28 | # no transfer, no producer. |
| 29 | local: bool = False |
| 30 | |
| 31 | |
| 32 | @dataclass |
| 33 | class ECMooncakePushSpec: |
| 34 | """Describe a destination reservation prepared before a tensor is ready.""" |
| 35 | |
| 36 | mm_hash: str |
| 37 | nbytes: int |
| 38 | shape: tuple[int, ...] |
| 39 | dtype: str |
| 40 | consumer_zmq: str |
| 41 | transfer_id: str |
| 42 | request_id: str = "" |
| 43 | |
| 44 | |
| 45 | @dataclass |
| 46 | class ECMooncakeConnectorMetadata(ECConnectorMetadata): |
| 47 | """Worker operations emitted for one Scheduler step.""" |
| 48 | |
| 49 | loads: list[ECMooncakeLoadSpec] = field(default_factory=list) |
| 50 | pushes: list[ECMooncakePushSpec] = field(default_factory=list) |
| 51 | freed: list[str] | None = None |
| 52 | store_candidates: dict[str, TensorSpec] = field(default_factory=dict) |
| 53 | |
| 54 | |
| 55 | @dataclass |
| 56 | class ECMooncakeWorkerMetadata(ECConnectorWorkerMetadata): |
| 57 | """Completion state reported from Workers to the Scheduler.""" |
| 58 | |
| 59 | loaded: set[str] = field(default_factory=set) |
| 60 | failed_loads: set[str] = field(default_factory=set) |
| 61 | # Items the receive pool dropped under pressure. The scheduler assumes an |
| 62 | # evicted item stays resident until told otherwise. |
| 63 | reclaimed: set[str] = field(default_factory=set) |
| 64 | pending_saves: bool = False |
| 65 | failed_saves: set[str] = field(default_factory=set) |
| 66 | |
| 67 | def aggregate(self, other: ECConnectorWorkerMetadata) -> ECMooncakeWorkerMetadata: |
| 68 | assert isinstance(other, ECMooncakeWorkerMetadata) |
| 69 | return ECMooncakeWorkerMetadata( |
| 70 | # Every tensor-parallel rank gathers the embedding from its own |
| 71 | # cache, so an item counts as loaded only where all of them have |
| 72 | # it; one rank falling short must fail the load rather than leave |
| 73 | # the scheduler believing it is ready. |
| 74 | loaded=self.loaded & other.loaded, |
| 75 | failed_loads=self.failed_loads | other.failed_loads, |
| 76 | reclaimed=self.reclaimed | other.reclaimed, |
| 77 | pending_saves=self.pending_saves or other.pending_saves, |
| 78 | failed_saves=self.failed_saves | other.failed_saves, |
| 79 | ) |