| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import os |
| 4 | from dataclasses import dataclass |
| 5 | from typing import TYPE_CHECKING, Any |
| 6 | |
| 7 | import safetensors |
| 8 | |
| 9 | from vllm.config import VllmConfig |
| 10 | from vllm.distributed.ec_transfer.ec_connector.base import ( |
| 11 | ECConnectorBase, |
| 12 | ECConnectorMetadata, |
| 13 | ECConnectorRole, |
| 14 | ) |
| 15 | from vllm.distributed.ec_transfer.ec_connector.utils import ( |
| 16 | PlaceholderMetadataResolver, |
| 17 | collect_ec_item_metadata, |
| 18 | ) |
| 19 | from vllm.logger import init_logger |
| 20 | from vllm.utils.gpu_sync_debug import gpu_sync_allowed |
| 21 | from vllm.v1.core.sched.output import SchedulerOutput |
| 22 | |
| 23 | if TYPE_CHECKING: |
| 24 | from vllm.v1.request import Request |
| 25 | |
| 26 | logger = init_logger(__name__) |
| 27 | |
| 28 | |
| 29 | @dataclass |
| 30 | class MMMeta: |
| 31 | mm_hash: str |
| 32 | num_token: int |
| 33 | |
| 34 | @staticmethod |
| 35 | def make_meta(mm_hash, num_token) -> "MMMeta": |
| 36 | return MMMeta(mm_hash=mm_hash, num_token=num_token) |
| 37 | |
| 38 | |
| 39 | @dataclass |
| 40 | class ECExampleConnectorMetadata(ECConnectorMetadata): |
| 41 | mm_datas: list[MMMeta] |
| 42 | |
| 43 | def __init__(self): |
| 44 | self.mm_datas = [] |
| 45 | |
| 46 | def add_mm_data(self, mm_data: MMMeta): |
| 47 | self.mm_datas.append(mm_data) |
| 48 | |
| 49 | |
| 50 | class ECExampleConnector(ECConnectorBase): |
| 51 | # NOTE: This is Simple debug implementation of the EC connector. |
| 52 | # It save / load the EC cache to / from the disk. |
| 53 | |
| 54 | def __init__(self, vllm_config: "VllmConfig", role: ECConnectorRole): |
| 55 | super().__init__(vllm_config=vllm_config, role=role) |
| 56 | # req_id -> index |
| 57 | self._mm_datas_need_loads: dict[str, int] = {} |
| 58 | self._metadata_resolver = PlaceholderMetadataResolver(vllm_config.model_config) |
| 59 | transfer_config = vllm_config.ec_transfer_config |
| 60 | if transfer_config is not None: |
| 61 | self._storage_path = transfer_config.get_from_extra_config( |
| 62 | "shared_storage_path", "/tmp" |
| 63 | ) |
| 64 | logger.debug(transfer_config) |
| 65 | logger.debug("Shared storage path is %s", self._storage_path) |
| 66 | else: |
| 67 | raise ValueError("ec_transfer_config must be set for ECConnectorBase") |
| 68 | |
| 69 | def start_load_caches(self, encoder_cache, **kwargs) -> None: |
| 70 | """Start loading the cache from the connector into vLLM's encoder cache. |
| 71 | |
| 72 | This method loads the encoder cache based on metadata provided by the scheduler. |
| 73 | It is called before `_gather_mm_embeddings` for the EC Connector. For EC, |
| 74 | the `encoder_cache` and `mm_hash` are stored in `kwargs`. |
| 75 | |
| 76 | Args: |
| 77 | encoder_cache (dict[str, torch.Tensor]): A dictionary mapping multimodal |
| 78 | data hashes (`mm_hash`) to encoder cache tensors. |
| 79 | kwargs (dict): Additional keyword arguments for the connector. |
| 80 | |
| 81 | """ |
| 82 | from vllm.platforms import current_platform |
| 83 | |
| 84 | # Get the metadata |
| 85 | metadata: ECConnectorMetadata = self._get_connector_metadata() |
| 86 | assert isinstance(metadata, ECExampleConnectorMetadata) |
| 87 | assert encoder_cache is not None |
| 88 | if metadata is None: |
| 89 | logger.warning( |
| 90 | "In connector.start_load_caches, but the connector metadata is None" |
| 91 | ) |
| 92 | return |
| 93 | # A bare "cuda" makes safetensors load onto cuda:0. Pin the load to the |
| 94 | # selected device so tensor-parallel workers load the cache locally. |
| 95 | device = current_platform.device_type |
| 96 | if current_platform.is_cuda_alike(): |
| 97 | device = f"{device}:{current_platform.current_device()}" |
| 98 | # Load the EC for each mm data |
| 99 | for mm_data in metadata.mm_datas: |
| 100 | if mm_data.mm_hash in encoder_cache: |
| 101 | continue |
| 102 | filename = self._generate_filename_debug(mm_data.mm_hash) |
| 103 | with gpu_sync_allowed(): |
| 104 | ec_cache = safetensors.torch.load_file(filename, device=device)[ |
| 105 | "ec_cache" |
| 106 | ] |
| 107 | encoder_cache[mm_data.mm_hash] = ec_cache |
| 108 | logger.debug("Success load encoder cache for hash %s", mm_data.mm_hash) |
| 109 | |
| 110 | def save_caches(self, encoder_cache, mm_hash, **kwargs) -> None: |
| 111 | """Save the encoder cache to the connector. |
| 112 | |
| 113 | This method saves the encoder cache from the worker's local storage |
| 114 | to shared storage or another external connector. |
| 115 | |
| 116 | Args: |
| 117 | encoder_cache (dict[str, torch.Tensor]): A dictionary mapping multimodal |
| 118 | data hashes (`mm_hash`) to encoder cache tensors. |
| 119 | mm_hash (str): The hash of the multimodal data whose cache is being saved. |
| 120 | kwargs (dict): Additional keyword arguments for the connector. |
| 121 | |
| 122 | """ |
| 123 | # Return if it is PD Instance |
| 124 | if not self.is_producer: |
| 125 | return |
| 126 | filename = self._generate_filename_debug(mm_hash) |
| 127 | ec_cache = encoder_cache[mm_hash] |
| 128 | with gpu_sync_allowed(): |
| 129 | tensors = {"ec_cache": ec_cache.detach().cpu()} |
| 130 | safetensors.torch.save_file(tensors, filename) |
| 131 | logger.debug("Save cache successful for mm_hash %s", mm_hash) |
| 132 | |
| 133 | def has_cache_item( |
| 134 | self, |
| 135 | identifier: str, |
| 136 | ) -> bool: |
| 137 | """Check if cache exist externally for the media. |
| 138 | |
| 139 | Args: |
| 140 | identifier (str): the identifier of the media. |
| 141 | |
| 142 | Returns: |
| 143 | Bool indicate that media exists in cache or not |
| 144 | |
| 145 | """ |
| 146 | return self._found_match_for_mm_data(identifier) |
| 147 | |
| 148 | def update_state_after_alloc( |
| 149 | self, |
| 150 | request: "Request", |
| 151 | index: int, |
| 152 | ) -> None: |
| 153 | """Update ECConnector state after encoder cache allocation.""" |
| 154 | mm_hash = request.mm_features[index].identifier |
| 155 | # Only load cache if it is consumer and cache exists |
| 156 | if not self.is_consumer or not self.has_cache_item(mm_hash): |
| 157 | return |
| 158 | num_encoder_token = request.get_num_encoder_embeds(index) |
| 159 | self._mm_datas_need_loads[mm_hash] = num_encoder_token |
| 160 | |
| 161 | def build_connector_meta( |
| 162 | self, |
| 163 | scheduler_output: SchedulerOutput, |
| 164 | ) -> ECConnectorMetadata: |
| 165 | """Build the connector metadata for this step. |
| 166 | |
| 167 | This function should NOT modify any fields in the scheduler_output. |
| 168 | Also, calling this function will reset the state of the connector. |
| 169 | This only build for load mm_data only |
| 170 | Args: |
| 171 | scheduler_output (SchedulerOutput): the scheduler output object. |
| 172 | """ |
| 173 | meta = ECExampleConnectorMetadata() |
| 174 | for mm_hash, num_encoder_token in self._mm_datas_need_loads.items(): |
| 175 | meta.add_mm_data(MMMeta.make_meta(mm_hash, num_encoder_token)) |
| 176 | self._mm_datas_need_loads.clear() |
| 177 | return meta |
| 178 | |
| 179 | def request_finished( |
| 180 | self, |
| 181 | request: "Request", |
| 182 | ) -> tuple[bool, dict[str, Any] | None]: |
| 183 | """Report each item's cache key and grid so a consumer can skip the |
| 184 | image transform. |
| 185 | |
| 186 | A consumer only needs the grid to size the prompt's placeholder range; |
| 187 | the embedding itself arrives through this connector. Reporting the grid |
| 188 | the producer actually computed keeps the two sides in agreement without |
| 189 | the caller re-deriving it from the raw media. |
| 190 | """ |
| 191 | if not self.is_producer: |
| 192 | return False, None |
| 193 | |
| 194 | items = collect_ec_item_metadata(request.mm_features, self._metadata_resolver) |
| 195 | if not items: |
| 196 | return False, None |
| 197 | return False, items |
| 198 | |
| 199 | # ============================== |
| 200 | # Helper functions |
| 201 | # ============================== |
| 202 | |
| 203 | def _found_match_for_mm_data(self, mm_hash) -> bool: |
| 204 | """Check if the cache is hit for the request.""" |
| 205 | filename = self._generate_filename_debug(mm_hash) |
| 206 | return os.path.exists(filename) |
| 207 | |
| 208 | def _generate_foldername_debug( |
| 209 | self, |
| 210 | mm_hash: str, |
| 211 | create_folder: bool = True, # <- now defaults to True |
| 212 | ) -> str: |
| 213 | """Return the folder in which the cache for this mm_hash lives. |
| 214 | If `create_folder` is True (default) the directory is created |
| 215 | recursively the first time it is needed. |
| 216 | """ |
| 217 | foldername = os.path.join(self._storage_path, mm_hash) |
| 218 | if create_folder: |
| 219 | os.makedirs(foldername, exist_ok=True) |
| 220 | return foldername |
| 221 | |
| 222 | def _generate_filename_debug(self, mm_hash: str) -> str: |
| 223 | """Return the full path of the safetensors file for this mm_hash. |
| 224 | Ensures the parent directory exists because |
| 225 | `_generate_foldername_debug` is called with its default |
| 226 | (`create_folder=True`). |
| 227 | """ |
| 228 | foldername = self._generate_foldername_debug(mm_hash) # <- folder auto-created |
| 229 | return os.path.join(foldername, "encoder_cache.safetensors") |