| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from dataclasses import dataclass, field |
| 4 | |
| 5 | from vllm.distributed.kv_transfer.kv_connector.v1.base import ( |
| 6 | KVConnectorMetadata, |
| 7 | KVConnectorWorkerMetadata, |
| 8 | ) |
| 9 | from vllm.v1.kv_offload.base import LoadStoreSpec |
| 10 | |
| 11 | ReqId = str |
| 12 | |
| 13 | |
| 14 | @dataclass(slots=True) |
| 15 | class DirectionalTransferStats: |
| 16 | bytes: int = 0 |
| 17 | time: float = 0.0 |
| 18 | sizes: list[int | float] = field(default_factory=list) |
| 19 | |
| 20 | def aggregate( |
| 21 | self, other: "DirectionalTransferStats" |
| 22 | ) -> "DirectionalTransferStats": |
| 23 | return DirectionalTransferStats( |
| 24 | bytes=self.bytes + other.bytes, |
| 25 | time=self.time + other.time, |
| 26 | sizes=[*self.sizes, *other.sizes], |
| 27 | ) |
| 28 | |
| 29 | def record(self, num_bytes: int, time: float) -> None: |
| 30 | self.bytes += num_bytes |
| 31 | self.time += time |
| 32 | self.sizes.append(num_bytes) |
| 33 | |
| 34 | def is_empty(self) -> bool: |
| 35 | return self.bytes == 0 and self.time == 0.0 and not self.sizes |
| 36 | |
| 37 | |
| 38 | @dataclass(slots=True) |
| 39 | class TransferStats: |
| 40 | load: DirectionalTransferStats = field(default_factory=DirectionalTransferStats) |
| 41 | store: DirectionalTransferStats = field(default_factory=DirectionalTransferStats) |
| 42 | |
| 43 | def aggregate(self, other: "TransferStats") -> "TransferStats": |
| 44 | return TransferStats( |
| 45 | load=self.load.aggregate(other.load), |
| 46 | store=self.store.aggregate(other.store), |
| 47 | ) |
| 48 | |
| 49 | def is_empty(self) -> bool: |
| 50 | return self.load.is_empty() and self.store.is_empty() |
| 51 | |
| 52 | |
| 53 | @dataclass |
| 54 | class TransferJob: |
| 55 | """A transfer job bundling request context with transfer spec. |
| 56 | |
| 57 | Used for both loads and stores, keyed by scheduler-assigned job ID. |
| 58 | The worker reports the job ID back when the transfer finishes, |
| 59 | and the scheduler processes the completion. |
| 60 | """ |
| 61 | |
| 62 | req_id: ReqId |
| 63 | src_spec: LoadStoreSpec |
| 64 | dst_spec: LoadStoreSpec |
| 65 | |
| 66 | |
| 67 | @dataclass |
| 68 | class OffloadingConnectorMetadata(KVConnectorMetadata): |
| 69 | # Keyed by scheduler-assigned job IDs. |
| 70 | load_jobs: dict[int, TransferJob] |
| 71 | store_jobs: dict[int, TransferJob] |
| 72 | jobs_to_flush: set[int] | None = None |
| 73 | |
| 74 | |
| 75 | @dataclass |
| 76 | class OffloadingWorkerMetadata(KVConnectorWorkerMetadata): |
| 77 | """Worker -> Scheduler metadata for completed transfer jobs. |
| 78 | |
| 79 | Each worker reports {job_id: 1} for newly completed transfer jobs |
| 80 | (load or store). aggregate() sums counts across workers within a step. |
| 81 | The scheduler accumulates across steps and processes |
| 82 | a transfer completion only when count reaches num_workers. |
| 83 | """ |
| 84 | |
| 85 | completed_jobs: dict[int, int] = field(default_factory=dict) |
| 86 | transfer_stats: TransferStats = field(default_factory=TransferStats) |
| 87 | |
| 88 | def mark_completed(self, job_id: int) -> None: |
| 89 | """Record a transfer job completion from this worker.""" |
| 90 | self.completed_jobs[job_id] = 1 |
| 91 | |
| 92 | def aggregate( |
| 93 | self, other: "KVConnectorWorkerMetadata" |
| 94 | ) -> "KVConnectorWorkerMetadata": |
| 95 | assert isinstance(other, OffloadingWorkerMetadata) |
| 96 | |
| 97 | merged = dict(self.completed_jobs) |
| 98 | for job_id, v in other.completed_jobs.items(): |
| 99 | merged[job_id] = merged.get(job_id, 0) + v |
| 100 | |
| 101 | return OffloadingWorkerMetadata( |
| 102 | completed_jobs=merged, |
| 103 | transfer_stats=self.transfer_stats.aggregate(other.transfer_stats), |
| 104 | ) |