| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Shared config builders for EC connector unit tests.""" |
| 4 | |
| 5 | import uuid |
| 6 | from unittest.mock import Mock |
| 7 | |
| 8 | import torch |
| 9 | |
| 10 | from vllm.config import ModelConfig, VllmConfig |
| 11 | from vllm.config.ec_transfer import ECRole, ECTransferConfig |
| 12 | from vllm.config.parallel import ParallelConfig |
| 13 | |
| 14 | # The scheduler derives its per-load ack count from the PCP size only under the |
| 15 | # `mp` backend -- which also keeps the world size off the local GPU count, so |
| 16 | # TP > 1 stays expressible on a single-GPU host. |
| 17 | _EXECUTOR_BACKEND = "mp" |
| 18 | _PCP_SIZE = 1 |
| 19 | |
| 20 | |
| 21 | def create_ec_vllm_config( |
| 22 | *, |
| 23 | ec_role: ECRole = "ec_both", |
| 24 | tensor_parallel_size: int = 1, |
| 25 | rank: int = 0, |
| 26 | dtype: torch.dtype = torch.float16, |
| 27 | ) -> Mock: |
| 28 | """Build a `VllmConfig` stand-in for EC connector unit tests. |
| 29 | |
| 30 | `ec_transfer_config` and `parallel_config` are the real config objects, so |
| 31 | role derivation and the rank arithmetic the connector depends on behave as |
| 32 | they do in production. |
| 33 | |
| 34 | `model_config` is a stub because nothing under test needs a real one: the |
| 35 | connector reads only `dtype` here, and the paths that inspect the model |
| 36 | itself are not exercised. Building a real `ModelConfig` would resolve an HF |
| 37 | model and inspect its architecture. |
| 38 | |
| 39 | Args: |
| 40 | ec_role: EC role this instance plays. `is_ec_producer` and |
| 41 | `is_ec_consumer` are derived from it, as in production. |
| 42 | tensor_parallel_size: TP degree. |
| 43 | rank: Global rank of this worker. |
| 44 | dtype: Encoder cache dtype. |
| 45 | |
| 46 | Returns: |
| 47 | A `VllmConfig`-specced mock carrying the configs above. |
| 48 | |
| 49 | """ |
| 50 | parallel_config = ParallelConfig( |
| 51 | tensor_parallel_size=tensor_parallel_size, |
| 52 | prefill_context_parallel_size=_PCP_SIZE, |
| 53 | distributed_executor_backend=_EXECUTOR_BACKEND, |
| 54 | rank=rank, |
| 55 | ) |
| 56 | ec_transfer_config = ECTransferConfig( |
| 57 | # The role properties report False unless a connector is named. |
| 58 | ec_connector="ECCPUConnector", |
| 59 | ec_role=ec_role, |
| 60 | engine_id=str(uuid.uuid4()), |
| 61 | ) |
| 62 | |
| 63 | model_config = Mock(spec=ModelConfig) |
| 64 | model_config.dtype = dtype |
| 65 | |
| 66 | vllm_config = Mock(spec=VllmConfig) |
| 67 | vllm_config.ec_transfer_config = ec_transfer_config |
| 68 | vllm_config.parallel_config = parallel_config |
| 69 | vllm_config.model_config = model_config |
| 70 | vllm_config.instance_id = f"ec-test-{uuid.uuid4()}" |
| 71 | return vllm_config |