| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from multiprocessing.synchronize import Lock as LockType |
| 4 | from typing import Literal |
| 5 | |
| 6 | from vllm.config import VllmConfig |
| 7 | |
| 8 | from .base import ( |
| 9 | BaseMultiModalProcessorCache, |
| 10 | BaseMultiModalReceiverCache, |
| 11 | MultiModalProcessorOnlyCache, |
| 12 | ) |
| 13 | from .lru import ( |
| 14 | LruKeyReplicatedReceiverCache, |
| 15 | LruKeyReplicatedSenderCache, |
| 16 | ) |
| 17 | from .shm import ( |
| 18 | ShmObjectStoreReceiverCache, |
| 19 | ShmObjectStoreSenderCache, |
| 20 | ) |
| 21 | |
| 22 | |
| 23 | def _get_cache_type( |
| 24 | vllm_config: VllmConfig, |
| 25 | ) -> Literal[None, "processor_only", "lru", "shm"]: |
| 26 | model_config = vllm_config.model_config |
| 27 | if not model_config.supports_multimodal_inputs: |
| 28 | return None |
| 29 | |
| 30 | # Check if the cache is disabled. |
| 31 | mm_config = model_config.get_multimodal_config() |
| 32 | if mm_config.mm_processor_cache_gb <= 0: |
| 33 | return None |
| 34 | |
| 35 | # Check if IPC caching is supported. |
| 36 | parallel_config = vllm_config.parallel_config |
| 37 | is_ipc_supported = parallel_config._api_process_count == 1 and ( |
| 38 | parallel_config.data_parallel_size == 1 |
| 39 | or parallel_config.data_parallel_external_lb |
| 40 | ) |
| 41 | |
| 42 | if not is_ipc_supported: |
| 43 | return "processor_only" |
| 44 | |
| 45 | mm_config = model_config.get_multimodal_config() |
| 46 | return mm_config.mm_processor_cache_type |
| 47 | |
| 48 | |
| 49 | def processor_cache_from_config( |
| 50 | vllm_config: VllmConfig, |
| 51 | ) -> BaseMultiModalProcessorCache | None: |
| 52 | """Return a `BaseMultiModalProcessorCache`, if enabled.""" |
| 53 | cache_type = _get_cache_type(vllm_config) |
| 54 | if cache_type is None: |
| 55 | return None |
| 56 | elif cache_type == "processor_only": |
| 57 | return MultiModalProcessorOnlyCache(vllm_config.model_config) |
| 58 | elif cache_type == "lru": |
| 59 | return LruKeyReplicatedSenderCache(vllm_config.model_config) |
| 60 | elif cache_type == "shm": |
| 61 | return ShmObjectStoreSenderCache(vllm_config) |
| 62 | else: |
| 63 | raise ValueError(f"Unknown cache type: {cache_type!r}") |
| 64 | |
| 65 | |
| 66 | def processor_only_cache_from_config( |
| 67 | vllm_config: VllmConfig, |
| 68 | ) -> MultiModalProcessorOnlyCache | None: |
| 69 | """Return a `MultiModalProcessorOnlyCache`, if enabled.""" |
| 70 | cache_type = _get_cache_type(vllm_config) |
| 71 | if cache_type is None: |
| 72 | return None |
| 73 | |
| 74 | return MultiModalProcessorOnlyCache(vllm_config.model_config) |
| 75 | |
| 76 | |
| 77 | def engine_receiver_cache_from_config( |
| 78 | vllm_config: VllmConfig, |
| 79 | ) -> BaseMultiModalReceiverCache | None: |
| 80 | """Return a `BaseMultiModalReceiverCache` for the engine process.""" |
| 81 | cache_type = _get_cache_type(vllm_config) |
| 82 | if cache_type in (None, "processor_only", "shm"): |
| 83 | return None |
| 84 | elif cache_type == "lru": |
| 85 | return LruKeyReplicatedReceiverCache(vllm_config.model_config) |
| 86 | else: |
| 87 | raise ValueError(f"Unknown cache type: {cache_type!r}") |
| 88 | |
| 89 | |
| 90 | def worker_receiver_cache_from_config( |
| 91 | vllm_config: VllmConfig, |
| 92 | shared_worker_lock: LockType | None, |
| 93 | ) -> BaseMultiModalReceiverCache | None: |
| 94 | """Return a `BaseMultiModalReceiverCache` for the worker process.""" |
| 95 | cache_type = _get_cache_type(vllm_config) |
| 96 | if cache_type in (None, "processor_only", "lru"): |
| 97 | return None |
| 98 | elif cache_type == "shm": |
| 99 | if shared_worker_lock is None: |
| 100 | raise ValueError( |
| 101 | "Missing `shared_worker_lock` argument from executor. " |
| 102 | "This argument is needed for mm_processor_cache_type='shm'." |
| 103 | ) |
| 104 | |
| 105 | return ShmObjectStoreReceiverCache(vllm_config, shared_worker_lock) |
| 106 | else: |
| 107 | raise ValueError(f"Unknown cache type: {cache_type!r}") |