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
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
model_runner
instance-attribute
__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
|
|
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
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 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
|