Davids048andClaude Opus 5.5
6c2a250fd0
- ComposedPipelineBase owns a ComponentState: it records each component's path after loading and attaches the state to every stage it registers. from_pretrained(model_path, *, resolved_config, ...) takes a resolved config instead of building FastVideoArgs from keywords or argparse; the training_args alias and TrainingArgs checks are gone, and the training DiT offload is a recorded override. Compile settings come from engine.compile and torch_compile_kwargs(); LoRA settings from pipeline.components.lora_* and training.lora.*. - The worker and executors read engine.num_gpus, engine.parallelism.*, engine.execution_backend, and pipeline.output_type; the Ray executor owns its placement group and runtime env. - VideoGenerator builds from a resolved config only (_from_resolved_config); the prompt_txt fallback is gone (requests carry inputs.prompt_path). The OpenAI server state holds the resolved config (get_resolved_config()). - Tests use resolved configs; the SSIM and performance helpers map their settings to dotted typed paths. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_012M91wnVFmPEvJ9h39r5BH7
91 lines
3.8 KiB
Python
91 lines
3.8 KiB
Python
import os
|
|
from typing import Any
|
|
|
|
from fastvideo.api.resolution import ResolvedGeneratorConfig
|
|
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,
|
|
resolved_config: ResolvedGeneratorConfig,
|
|
rpc_rank: int = 0,
|
|
) -> None:
|
|
"""
|
|
Initialize the worker wrapper with the given resolved_config 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.resolved_config: ResolvedGeneratorConfig | None = None
|
|
# do not store this `resolved_config`, `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.resolved_config = kwargs.get("resolved_config")
|
|
assert self.resolved_config is not None, ("resolved_config 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)
|