vLLM/vllm/distributed/weight_transfer/factory.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Factory for weight transfer engines with lazy loading."""

import importlib
from collections.abc import Callable
from typing import TYPE_CHECKING

from vllm.distributed.weight_transfer.base import (
    TrainerWeightTransferEngine,
    WeightTransferEngine,
)
from vllm.logger import init_logger

if TYPE_CHECKING:
    import torch

    from vllm.config import VllmConfig
    from vllm.config.weight_transfer import WeightTransferConfig
    from vllm.distributed.weight_transfer.base import (
        TrainerInitInfo,
        VLLMWeightSyncClient,
        WeightSource,
    )

logger = init_logger(__name__)


class WeightTransferEngineFactory:
    """Factory for creating weight transfer engines with lazy loading.

    This factory implements a registry pattern that supports:
    - Lazy loading: Engine modules are only imported when actually needed
    - Extensibility: Custom engines can be registered at runtime
    - Centralized registration: All built-in engines registered in one place
    """

    _registry: dict[str, Callable[[], type[WeightTransferEngine]]] = {}

    @classmethod
    def register_engine(
        cls,
        name: str,
        module_path_or_cls: str | type[WeightTransferEngine],
        class_name: str | None = None,
    ) -> None:
        """Register an engine with lazy-loading or direct class reference.

        Supports two calling conventions:
        1. Lazy loading: register_engine(name, module_path, class_name)
        2. Direct class: register_engine(name, engine_cls)

        Args:
            name: The name to register the engine under (e.g., "nccl")
            module_path_or_cls: Either a module path string for lazy loading,
                or the engine class directly
            class_name: Name of the engine class (required if module_path is string)

        Raises:
            ValueError: If an engine with the same name is already registered
        """
        if name in cls._registry:
            raise ValueError(f"Weight transfer engine '{name}' is already registered.")

        if isinstance(module_path_or_cls, str):
            # Lazy loading path
            module_path = module_path_or_cls
            if class_name is None:
                raise ValueError(
                    "class_name is required when registering with module path"
                )

            def loader() -> type[WeightTransferEngine]:
                module = importlib.import_module(module_path)
                return getattr(module, class_name)

            cls._registry[name] = loader
        else:
            # Direct class registration
            engine_cls = module_path_or_cls
            cls._registry[name] = lambda: engine_cls

    @classmethod
    def create_engine(
        cls,
        config: "WeightTransferConfig",
        vllm_config: "VllmConfig",
        device: "torch.device",
        model: "torch.nn.Module",
    ) -> WeightTransferEngine:
        """Create a weight transfer engine instance.

        Args:
            config: Weight transfer configuration containing the backend name
            vllm_config: The full vLLM config (provides parallel/model config)
            device: The device this worker's model lives on
            model: The local model instance which will receive the weights

        Returns:
            An initialized weight transfer engine instance

        Raises:
            ValueError: If the backend is not registered
        """
        backend = config.backend
        if backend not in cls._registry:
            available = list(cls._registry.keys())
            raise ValueError(
                f"Invalid weight transfer backend: {backend}. "
                f"Available engines: {available}"
            )
        engine_cls = cls._registry[backend]()

        logger.info(
            "Creating weight transfer engine: %s",
            engine_cls.__name__,
        )

        return engine_cls(config, vllm_config, device, model)


class WeightTransferTrainerFactory:
    """Factory for creating trainer-side weight transfer engines.

    Parallel to `WeightTransferEngineFactory`, with its own lazy-import
    registry. The trainer-side and worker-side registries are kept separate:
    they share backend names by convention, but the trainer process never
    instantiates a worker engine and vice versa, so unifying them would only
    couple the import graphs.
    """

    _registry: dict[str, Callable[[], type[TrainerWeightTransferEngine]]] = {}

    @classmethod
    def register_engine(
        cls,
        name: str,
        module_path_or_cls: "str | type[TrainerWeightTransferEngine]",
        class_name: str | None = None,
    ) -> None:
        """Register a trainer engine. Same conventions as
        `WeightTransferEngineFactory.register_engine`."""
        if name in cls._registry:
            raise ValueError(
                f"Weight transfer trainer engine '{name}' is already registered."
            )

        if isinstance(module_path_or_cls, str):
            module_path = module_path_or_cls
            if class_name is None:
                raise ValueError(
                    "class_name is required when registering with module path"
                )

            def loader() -> type[TrainerWeightTransferEngine]:
                module = importlib.import_module(module_path)
                return getattr(module, class_name)

            cls._registry[name] = loader
        else:
            engine_cls = module_path_or_cls
            cls._registry[name] = lambda: engine_cls

    @classmethod
    def trainer_init(
        cls,
        init_info: "TrainerInitInfo",
        *,
        client: "VLLMWeightSyncClient",
        source: "WeightSource | None" = None,
    ) -> TrainerWeightTransferEngine:
        """Build and rendezvous a ready-to-send trainer engine.

        Called on every trainer rank (multi-rank trainers construct on all
        ranks; the sender is resolved inside the engine's ``trainer_init``).

        The trainer side takes no `WeightTransferConfig` and no separate
        `backend` argument: the backend is read from ``init_info.backend`` (a
        `ClassVar` on each `TrainerInitInfo` subclass), and the static wire
        params ride `init_info`.

        Args:
            init_info: Backend-specific trainer init info. Its `backend`
                selects the engine; it also carries the wire params (e.g.
                `packed`).
            client: Inference-side control-plane client.
            source: `WeightSource` of `(name, tensor)` pairs to send each round,
                for full-resync backends (NCCL, IPC). Sparse backend
                omits it and passes its per-round payload to `send_weights`.

        Raises:
            ValueError: If `init_info.backend` is not registered.
        """
        backend = init_info.backend
        if backend not in cls._registry:
            available = list(cls._registry.keys())
            raise ValueError(
                f"Invalid weight transfer backend: {backend}. "
                f"Available trainer engines: {available}"
            )
        engine_cls = cls._registry[backend]()

        logger.info(
            "Creating weight transfer trainer engine: %s",
            engine_cls.__name__,
        )

        return engine_cls.trainer_init(
            init_info=init_info,
            client=client,
            source=source,
        )


# Register built-in weight transfer engines here.
# Registration should be centralized to ensure lazy loading -
# engine modules are only imported when actually used.

WeightTransferEngineFactory.register_engine(
    "nccl",
    "vllm.distributed.weight_transfer.nccl_engine",
    "NCCLWeightTransferEngine",
)

WeightTransferEngineFactory.register_engine(
    "ipc",
    "vllm.distributed.weight_transfer.ipc_engine",
    "IPCWeightTransferEngine",
)

WeightTransferEngineFactory.register_engine(
    "sparse_nccl",
    "vllm.distributed.weight_transfer.sparse_nccl_engine",
    "SparseNCCLWeightTransferEngine",
)

WeightTransferEngineFactory.register_engine(
    "sharded_rdt",
    "vllm.distributed.weight_transfer.sharded_rdt_engine",
    "ShardedRDTWeightTransferEngine",
)


# Trainer-side engines, parallel to the worker registry above.
WeightTransferTrainerFactory.register_engine(
    "nccl",
    "vllm.distributed.weight_transfer.nccl_engine",
    "NCCLTrainerWeightTransferEngine",
)

WeightTransferTrainerFactory.register_engine(
    "ipc",
    "vllm.distributed.weight_transfer.ipc_engine",
    "IPCTrainerWeightTransferEngine",
)

WeightTransferTrainerFactory.register_engine(
    "sparse_nccl",
    "vllm.distributed.weight_transfer.sparse_nccl_engine",
    "SparseNCCLTrainerWeightTransferEngine",
)

WeightTransferTrainerFactory.register_engine(
    "sharded_rdt",
    "vllm.distributed.weight_transfer.sharded_rdt_trainer",
    "ShardedRDTTrainerWeightTransferEngine",
)