Skip to content

vllm.v1.worker.worker_base

logger module-attribute

logger = init_logger(__name__)

WorkerBase

Bases: WorkerBase

Abstract class for v1 worker, mainly define some methods for v1. For methods shared by v0 and v1, define them in v0 WorkerBase

Source code in vllm/v1/worker/worker_base.py
class WorkerBase(WorkerBaseV0):
    """
    Abstract class for v1 worker, mainly define some methods for v1.
    For methods shared by v0 and v1, define them in v0 WorkerBase
    """

    def __init__(
        self,
        vllm_config: VllmConfig,
        local_rank: int,
        rank: int,
        distributed_init_method: str,
        is_driver_worker: bool = False,
    ):
        """
        Initialize common worker components.

        Args:
            vllm_config: Complete vLLM configuration
            local_rank: Local device index
            rank: Global rank in distributed setup
            distributed_init_method: Distributed initialization method
            is_driver_worker: Whether this worker handles driver 
            responsibilities
        """
        # Configuration storage
        super().__init__(vllm_config=vllm_config)

        self.parallel_config.rank = rank
        self.local_rank = local_rank
        self.rank = rank
        self.distributed_init_method = distributed_init_method
        self.is_driver_worker = is_driver_worker

        # Device and model state
        self.device: Optional[torch.device] = None
        self.model_runner: Optional[nn.Module] = None

    def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
        """Get specifications for KV cache implementation."""
        raise NotImplementedError

    def compile_or_warm_up_model(self) -> None:
        """Prepare model for execution through compilation/warmup."""
        raise NotImplementedError

    def check_health(self) -> None:
        """Basic health check (override for device-specific checks)."""
        return

device instance-attribute

device: Optional[device] = None

distributed_init_method instance-attribute

distributed_init_method = distributed_init_method

is_driver_worker instance-attribute

is_driver_worker = is_driver_worker

local_rank instance-attribute

local_rank = local_rank

model_runner instance-attribute

model_runner: Optional[Module] = None

rank instance-attribute

rank = rank

__init__

__init__(
    vllm_config: VllmConfig,
    local_rank: int,
    rank: int,
    distributed_init_method: str,
    is_driver_worker: bool = False,
)

Initialize common worker components.

Parameters:

Name Type Description Default
vllm_config VllmConfig

Complete vLLM configuration

required
local_rank int

Local device index

required
rank int

Global rank in distributed setup

required
distributed_init_method str

Distributed initialization method

required
is_driver_worker bool

Whether this worker handles driver

False
Source code in vllm/v1/worker/worker_base.py
def __init__(
    self,
    vllm_config: VllmConfig,
    local_rank: int,
    rank: int,
    distributed_init_method: str,
    is_driver_worker: bool = False,
):
    """
    Initialize common worker components.

    Args:
        vllm_config: Complete vLLM configuration
        local_rank: Local device index
        rank: Global rank in distributed setup
        distributed_init_method: Distributed initialization method
        is_driver_worker: Whether this worker handles driver 
        responsibilities
    """
    # Configuration storage
    super().__init__(vllm_config=vllm_config)

    self.parallel_config.rank = rank
    self.local_rank = local_rank
    self.rank = rank
    self.distributed_init_method = distributed_init_method
    self.is_driver_worker = is_driver_worker

    # Device and model state
    self.device: Optional[torch.device] = None
    self.model_runner: Optional[nn.Module] = None

check_health

check_health() -> None

Basic health check (override for device-specific checks).

Source code in vllm/v1/worker/worker_base.py
def check_health(self) -> None:
    """Basic health check (override for device-specific checks)."""
    return

compile_or_warm_up_model

compile_or_warm_up_model() -> None

Prepare model for execution through compilation/warmup.

Source code in vllm/v1/worker/worker_base.py
def compile_or_warm_up_model(self) -> None:
    """Prepare model for execution through compilation/warmup."""
    raise NotImplementedError

get_kv_cache_spec

get_kv_cache_spec() -> dict[str, KVCacheSpec]

Get specifications for KV cache implementation.

Source code in vllm/v1/worker/worker_base.py
def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
    """Get specifications for KV cache implementation."""
    raise NotImplementedError