| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import importlib |
| 5 | from collections.abc import Callable |
| 6 | from typing import TYPE_CHECKING |
| 7 | |
| 8 | from vllm.distributed.ec_transfer.ec_connector.base import ( |
| 9 | ECConnectorBase, |
| 10 | ECConnectorRole, |
| 11 | ) |
| 12 | from vllm.logger import init_logger |
| 13 | |
| 14 | if TYPE_CHECKING: |
| 15 | from vllm.config import ECTransferConfig, VllmConfig |
| 16 | |
| 17 | logger = init_logger(__name__) |
| 18 | |
| 19 | |
| 20 | class ECConnectorFactory: |
| 21 | _registry: dict[str, Callable[[], type[ECConnectorBase]]] = {} |
| 22 | |
| 23 | @classmethod |
| 24 | def register_connector(cls, name: str, module_path: str, class_name: str) -> None: |
| 25 | """Register a connector with a lazy-loading module and class name.""" |
| 26 | if name in cls._registry: |
| 27 | raise ValueError(f"Connector '{name}' is already registered.") |
| 28 | |
| 29 | def loader() -> type[ECConnectorBase]: |
| 30 | module = importlib.import_module(module_path) |
| 31 | return getattr(module, class_name) |
| 32 | |
| 33 | cls._registry[name] = loader |
| 34 | |
| 35 | @classmethod |
| 36 | def create_connector( |
| 37 | cls, |
| 38 | config: "VllmConfig", |
| 39 | role: ECConnectorRole, |
| 40 | ) -> ECConnectorBase: |
| 41 | ec_transfer_config = config.ec_transfer_config |
| 42 | if ec_transfer_config is None: |
| 43 | raise ValueError("ec_transfer_config must be set to create a connector") |
| 44 | connector_cls = cls.get_connector_class(ec_transfer_config) |
| 45 | logger.info( |
| 46 | "Creating connector with name: %s and engine_id: %s", |
| 47 | connector_cls.__name__, |
| 48 | ec_transfer_config.engine_id, |
| 49 | ) |
| 50 | # Connector is explicitly separated into two roles. |
| 51 | # Scheduler connector: |
| 52 | # - Co-locate with scheduler process |
| 53 | # - Should only be used inside the Scheduler class |
| 54 | # Worker connector: |
| 55 | # - Co-locate with worker process |
| 56 | return connector_cls(config, role) |
| 57 | |
| 58 | @classmethod |
| 59 | def get_connector_class( |
| 60 | cls, ec_transfer_config: "ECTransferConfig" |
| 61 | ) -> type[ECConnectorBase]: |
| 62 | """Get the connector class by name.""" |
| 63 | connector_name = ec_transfer_config.ec_connector |
| 64 | if connector_name is None: |
| 65 | raise ValueError("EC connect must not be None") |
| 66 | elif connector_name in cls._registry: |
| 67 | connector_cls = cls._registry[connector_name]() |
| 68 | else: |
| 69 | connector_module_path = ec_transfer_config.ec_connector_module_path |
| 70 | if connector_module_path is None: |
| 71 | raise ValueError(f"Unsupported connector type: {connector_name}") |
| 72 | connector_module = importlib.import_module(connector_module_path) |
| 73 | connector_cls = getattr(connector_module, connector_name) |
| 74 | return connector_cls |
| 75 | |
| 76 | |
| 77 | # Register various connectors here. |
| 78 | # The registration should not be done in each individual file, as we want to |
| 79 | # only load the files corresponding to the current connector. |
| 80 | |
| 81 | ECConnectorFactory.register_connector( |
| 82 | "ECExampleConnector", |
| 83 | "vllm.distributed.ec_transfer.ec_connector.example_connector", |
| 84 | "ECExampleConnector", |
| 85 | ) |
| 86 | |
| 87 | ECConnectorFactory.register_connector( |
| 88 | "ECCPUConnector", |
| 89 | "vllm.distributed.ec_transfer.ec_connector.cpu.connector", |
| 90 | "ECCPUConnector", |
| 91 | ) |
| 92 | |
| 93 | ECConnectorFactory.register_connector( |
| 94 | "ECMooncakeConnector", |
| 95 | "vllm.distributed.ec_transfer.ec_connector.mooncake_ec_connector", |
| 96 | "ECMooncakeConnector", |
| 97 | ) |