91 lines
3.8 KiB
Python
91 lines
3.8 KiB
Python
import os
|
|
from typing import Any
|
|
|
|
from fastvideo.fastvideo_args import FastVideoArgs
|
|
from fastvideo.logger import init_logger
|
|
from fastvideo.worker.gpu_worker import Worker
|
|
from fastvideo.utils import (run_method, update_environment_variables)
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class WorkerWrapperBase:
|
|
"""
|
|
This class represents one process in an executor/engine. It is responsible
|
|
for lazily initializing the worker and handling the worker's lifecycle.
|
|
We first instantiate the WorkerWrapper, which remembers the worker module
|
|
and class name. Then, when we call `update_environment_variables`, and the
|
|
real initialization happens in `init_worker`.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
fastvideo_args: FastVideoArgs,
|
|
rpc_rank: int = 0,
|
|
) -> None:
|
|
"""
|
|
Initialize the worker wrapper with the given fastvideo_args and rpc_rank.
|
|
Note: rpc_rank is the rank of the worker in the executor. In most cases,
|
|
it is also the rank of the worker in the distributed group. However,
|
|
when multiple executors work together, they can be different.
|
|
e.g. in the case of SPMD-style offline inference with TP=2,
|
|
users can launch 2 engines/executors, each with only 1 worker.
|
|
All workers have rpc_rank=0, but they have different ranks in the TP
|
|
group.
|
|
"""
|
|
self.rpc_rank = rpc_rank
|
|
self.worker: Worker | None = None
|
|
self.fastvideo_args: FastVideoArgs | None = None
|
|
# do not store this `fastvideo_args`, `init_worker` will set the final
|
|
# one.
|
|
|
|
def adjust_rank(self, rank_mapping: dict[int, int]) -> None:
|
|
"""
|
|
Adjust the rpc_rank based on the given mapping.
|
|
It is only used during the initialization of the executor,
|
|
to adjust the rpc_rank of workers after we create all workers.
|
|
"""
|
|
if self.rpc_rank in rank_mapping:
|
|
self.rpc_rank = rank_mapping[self.rpc_rank]
|
|
|
|
def update_environment_variables(self, envs_list: list[dict[str, str]]) -> None:
|
|
envs = envs_list[self.rpc_rank]
|
|
key = 'CUDA_VISIBLE_DEVICES'
|
|
if key in envs and key in os.environ:
|
|
# overwriting CUDA_VISIBLE_DEVICES is desired behavior
|
|
# suppress the warning in `update_environment_variables`
|
|
del os.environ[key]
|
|
update_environment_variables(envs)
|
|
|
|
def init_worker(self, all_kwargs: list[dict[str, Any]]) -> None:
|
|
"""
|
|
Here we inject some common logic before initializing the worker.
|
|
Arguments are passed to the worker class constructor.
|
|
"""
|
|
kwargs = all_kwargs[self.rpc_rank]
|
|
self.fastvideo_args = kwargs.get("fastvideo_args")
|
|
assert self.fastvideo_args is not None, ("fastvideo_args is required to initialize the worker")
|
|
|
|
self.worker = Worker(**kwargs)
|
|
assert self.worker is not None
|
|
|
|
def execute_method(self, method: str | bytes, *args, **kwargs):
|
|
try:
|
|
# method resolution order:
|
|
# if a method is defined in this class, it will be called directly.
|
|
# otherwise, since we define `__getattr__` and redirect attribute
|
|
# query to `self.worker`, the method will be called on the worker.
|
|
return run_method(self, method, args, kwargs)
|
|
except Exception as e:
|
|
# if the driver worker also execute methods,
|
|
# exceptions in the rest worker may cause deadlock in rpc like ray
|
|
# see https://github.com/vllm-project/vllm/issues/3455
|
|
# print the error and inform the user to solve the error
|
|
msg = (f"Error executing method {method!r}. "
|
|
"This might cause deadlock in distributed execution.")
|
|
logger.exception(msg)
|
|
raise e
|
|
|
|
def __getattr__(self, attr):
|
|
return getattr(self.worker, attr)
|