| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Metadata for SimpleCPUOffloadConnector.""" |
| 4 | |
| 5 | from dataclasses import dataclass, field |
| 6 | |
| 7 | from vllm.distributed.kv_transfer.kv_connector.v1.base import ( |
| 8 | KVConnectorMetadata, |
| 9 | KVConnectorWorkerMetadata, |
| 10 | ) |
| 11 | |
| 12 | INVALID_JOB_ID = -1 |
| 13 | |
| 14 | |
| 15 | @dataclass |
| 16 | class SimpleCPUOffloadMetadata(KVConnectorMetadata): |
| 17 | """Metadata passed from scheduler to worker for CPU offload operations. |
| 18 | |
| 19 | The worker receives flat block lists keyed by a monotonic event_idx. |
| 20 | Job->req_id translation is handled by the scheduler-side manager |
| 21 | (via inverse maps), so the worker never knows about request identities. |
| 22 | """ |
| 23 | |
| 24 | # Load event per step. INVALID_JOB_ID means no blocks to load this step. |
| 25 | load_event: int = INVALID_JOB_ID |
| 26 | load_gpu_blocks: list[int] = field(default_factory=list) |
| 27 | load_cpu_blocks: list[int] = field(default_factory=list) |
| 28 | # Reverse map: load_event->req_ids, for tracking requests with finished load events |
| 29 | load_event_to_reqs: dict[int, list[str]] = field(default_factory=dict) |
| 30 | |
| 31 | # Store event per step. INVALID_JOB_ID means no blocks to store this step. |
| 32 | store_event: int = INVALID_JOB_ID |
| 33 | store_gpu_blocks: list[int] = field(default_factory=list) |
| 34 | store_cpu_blocks: list[int] = field(default_factory=list) |
| 35 | |
| 36 | # Whether any requests were preempted this step and need flush pending transfers. |
| 37 | need_flush: bool = False |
| 38 | |
| 39 | |
| 40 | @dataclass |
| 41 | class SimpleCPUOffloadWorkerMetadata(KVConnectorWorkerMetadata): |
| 42 | """Worker -> Scheduler metadata for completed store events. |
| 43 | |
| 44 | Each worker reports {event_idx: 1} for newly completed stores. |
| 45 | ``aggregate()`` sums counts across workers within a step. |
| 46 | The scheduler-side manager accumulates across steps and processes |
| 47 | a store completion only when count reaches ``world_size``. |
| 48 | """ |
| 49 | |
| 50 | completed_store_events: dict[int, int] |
| 51 | |
| 52 | def aggregate( |
| 53 | self, other: "KVConnectorWorkerMetadata" |
| 54 | ) -> "KVConnectorWorkerMetadata": |
| 55 | assert isinstance(other, SimpleCPUOffloadWorkerMetadata) |
| 56 | merged = dict(self.completed_store_events) |
| 57 | for k, v in other.completed_store_events.items(): |
| 58 | merged[k] = merged.get(k, 0) + v |
| 59 | return SimpleCPUOffloadWorkerMetadata(completed_store_events=merged) |