822 lines
34 KiB
Python
822 lines
34 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import atexit
|
|
import contextlib
|
|
from dataclasses import dataclass
|
|
from enum import Enum
|
|
import faulthandler
|
|
import logging
|
|
import logging.handlers
|
|
import multiprocessing as mp
|
|
from multiprocessing.connection import Connection
|
|
from multiprocessing.queues import Queue
|
|
import os
|
|
import queue
|
|
import signal
|
|
import time
|
|
from collections.abc import Callable
|
|
from multiprocessing.process import BaseProcess
|
|
from typing import Any, cast
|
|
|
|
import psutil
|
|
import torch
|
|
|
|
from fastvideo.distributed.parallel_state import get_dp_group, get_tp_group
|
|
import fastvideo.envs as envs
|
|
from fastvideo.fastvideo_args import FastVideoArgs
|
|
from fastvideo.logger import init_logger
|
|
from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
|
|
from fastvideo.utils import (decorate_logs, get_distributed_init_method, get_exception_traceback, get_loopback_ip,
|
|
get_mp_context, get_open_port, kill_itself_when_parent_died, force_spawn)
|
|
from fastvideo.worker.executor import Executor
|
|
from fastvideo.worker.worker_base import WorkerWrapperBase
|
|
|
|
logger = init_logger(__name__)
|
|
_RPC_ERROR_KEY = "__fastvideo_rpc_error__"
|
|
|
|
|
|
def _raise_for_rpc_errors(method: str | Callable, responses: list[Any]) -> None:
|
|
errors = []
|
|
for rank, response in enumerate(responses):
|
|
if isinstance(response, dict) and response.get(_RPC_ERROR_KEY):
|
|
errors.append(f"worker {rank}: {response.get('error', 'unknown error')}")
|
|
if errors:
|
|
raise RuntimeError(f"RPC {method!r} failed: " + "; ".join(errors))
|
|
|
|
|
|
def _make_queue_log_handler(log_queue: Queue) -> logging.Handler:
|
|
"""Create a QueueHandler that forwards fastvideo logs to a multiprocessing queue."""
|
|
return logging.handlers.QueueHandler(log_queue)
|
|
|
|
|
|
class StreamingTaskType(str, Enum):
|
|
"""
|
|
Enumeration for different streaming task types.
|
|
|
|
Inherits from str to allow string comparison for backward compatibility.
|
|
"""
|
|
RESET = "reset"
|
|
STEP = "step"
|
|
CLEAR = "clear"
|
|
EXIT = "exit"
|
|
|
|
|
|
@dataclass
|
|
class StreamingTask:
|
|
"""Task submitted to worker via input queue."""
|
|
task_type: StreamingTaskType
|
|
# For STEP tasks:
|
|
keyboard_action: torch.Tensor | None = None
|
|
mouse_action: torch.Tensor | None = None
|
|
# For RESET tasks:
|
|
batch: ForwardBatch | None = None
|
|
fastvideo_args: FastVideoArgs | None = None
|
|
|
|
|
|
@dataclass
|
|
class StreamingResult:
|
|
"""Result returned from worker via output queue."""
|
|
task_type: StreamingTaskType
|
|
output_batch: ForwardBatch | None = None
|
|
error: Exception | None = None
|
|
|
|
|
|
class MultiprocExecutor(Executor):
|
|
|
|
def _init_executor(self) -> None:
|
|
self.world_size = self.fastvideo_args.num_gpus
|
|
self.shutting_down = False
|
|
|
|
set_multiproc_executor_envs()
|
|
|
|
# Check if master_port is provided in fastvideo_args
|
|
master_port = get_open_port(self.fastvideo_args.master_port)
|
|
distributed_init_method = get_distributed_init_method(get_loopback_ip(), master_port)
|
|
logger.info("Use master port: %s", master_port)
|
|
|
|
# Create streaming queues BEFORE spawning workers
|
|
ctx = get_mp_context()
|
|
self._streaming_input_queue: Queue | None = ctx.Queue()
|
|
self._streaming_output_queue: Queue | None = ctx.Queue()
|
|
self._streaming_enabled = False
|
|
|
|
unready_workers: list[UnreadyWorkerProcHandle] = []
|
|
success = False
|
|
try:
|
|
for rank in range(self.world_size):
|
|
unready_workers.append(
|
|
WorkerMultiprocProc.make_worker_process(
|
|
fastvideo_args=self.fastvideo_args,
|
|
local_rank=rank,
|
|
rank=rank,
|
|
distributed_init_method=distributed_init_method,
|
|
streaming_input_queue=self._streaming_input_queue,
|
|
streaming_output_queue=self._streaming_output_queue,
|
|
log_queue=self._log_queue,
|
|
))
|
|
|
|
# Workers must be created before wait_for_ready to avoid
|
|
# deadlock, since worker.init_device() does a device sync.
|
|
self.workers = WorkerMultiprocProc.wait_for_ready(unready_workers)
|
|
|
|
success = True
|
|
finally:
|
|
if not success:
|
|
# Clean up the worker procs if there was a failure.
|
|
# Close death_writers first to signal workers to exit
|
|
self._ensure_worker_termination([uw.proc for uw in unready_workers])
|
|
|
|
# Register shutdown on exit
|
|
atexit.register(self.shutdown)
|
|
|
|
def execute_forward(self, forward_batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
|
|
responses = self.collective_rpc("execute_forward",
|
|
kwargs={
|
|
"forward_batch": forward_batch,
|
|
"fastvideo_args": fastvideo_args
|
|
})
|
|
output = responses[0]["output_batch"]
|
|
|
|
logging_info = None
|
|
logging_info = responses[0]["logging_info"] if envs.FASTVIDEO_STAGE_LOGGING else None
|
|
|
|
# Get extra dict (contains audio, peak_memory_mb, etc.)
|
|
extra = responses[0].get("extra", {})
|
|
|
|
result_batch = ForwardBatch(data_type=forward_batch.data_type,
|
|
output=output,
|
|
logging_info=logging_info,
|
|
extra=extra,
|
|
trajectory_latents=responses[0].get("trajectory_latents"),
|
|
trajectory_timesteps=responses[0].get("trajectory_timesteps"))
|
|
|
|
return result_batch
|
|
|
|
def execute_streaming_reset(self, forward_batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> dict[str, Any]:
|
|
responses = self.collective_rpc("execute_streaming_reset",
|
|
kwargs={
|
|
"forward_batch": forward_batch,
|
|
"fastvideo_args": fastvideo_args,
|
|
})
|
|
return responses[0]
|
|
|
|
def execute_streaming_step(self, keyboard_action: Any, mouse_action: Any) -> ForwardBatch:
|
|
responses = self.collective_rpc("execute_streaming_step",
|
|
kwargs={
|
|
"keyboard_action": keyboard_action,
|
|
"mouse_action": mouse_action,
|
|
})
|
|
return responses[0]
|
|
|
|
async def execute_streaming_step_async(self, keyboard_action: Any, mouse_action: Any) -> ForwardBatch:
|
|
responses = await self.collective_rpc_async("execute_streaming_step",
|
|
kwargs={
|
|
"keyboard_action": keyboard_action,
|
|
"mouse_action": mouse_action,
|
|
})
|
|
return responses[0]
|
|
|
|
def execute_streaming_clear(self) -> dict[str, Any]:
|
|
responses = self.collective_rpc("execute_streaming_clear")
|
|
return responses[0]
|
|
|
|
def enable_streaming(self) -> None:
|
|
if self._streaming_enabled:
|
|
return
|
|
|
|
self.collective_rpc("start_streaming_queue_loop")
|
|
self._streaming_enabled = True
|
|
|
|
def disable_streaming(self) -> None:
|
|
if not self._streaming_enabled:
|
|
return
|
|
|
|
if self._streaming_input_queue is not None:
|
|
self._streaming_input_queue.put(StreamingTask(task_type=StreamingTaskType.EXIT))
|
|
self._streaming_enabled = False
|
|
|
|
self._streaming_input_queue = None
|
|
self._streaming_output_queue = None
|
|
|
|
def submit_reset(self, forward_batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> None:
|
|
if not self._streaming_enabled:
|
|
self.enable_streaming()
|
|
|
|
self._streaming_input_queue.put(
|
|
StreamingTask(
|
|
task_type=StreamingTaskType.RESET,
|
|
batch=forward_batch,
|
|
fastvideo_args=fastvideo_args,
|
|
))
|
|
|
|
def submit_step(self, keyboard_action: torch.Tensor | None, mouse_action: torch.Tensor | None) -> None:
|
|
if not self._streaming_enabled:
|
|
raise RuntimeError("Streaming mode not enabled. Call enable_streaming() first.")
|
|
|
|
self._streaming_input_queue.put(
|
|
StreamingTask(
|
|
task_type=StreamingTaskType.STEP,
|
|
keyboard_action=keyboard_action,
|
|
mouse_action=mouse_action,
|
|
))
|
|
|
|
def submit_clear(self) -> None:
|
|
if self._streaming_enabled and self._streaming_input_queue is not None:
|
|
self._streaming_input_queue.put(StreamingTask(task_type=StreamingTaskType.CLEAR))
|
|
|
|
def get_result(self, timeout: float | None = None) -> StreamingResult | None:
|
|
if not self._streaming_enabled or self._streaming_output_queue is None:
|
|
return None
|
|
|
|
try:
|
|
if timeout == 0:
|
|
return self._streaming_output_queue.get_nowait()
|
|
else:
|
|
return self._streaming_output_queue.get(timeout=timeout)
|
|
except queue.Empty:
|
|
return None
|
|
|
|
def wait_result(self) -> StreamingResult:
|
|
if not self._streaming_enabled or self._streaming_output_queue is None:
|
|
raise RuntimeError("Streaming mode not enabled.")
|
|
|
|
return self._streaming_output_queue.get()
|
|
|
|
def set_lora_adapter(self,
|
|
lora_nickname: str,
|
|
lora_path: str | None = None,
|
|
strength: float = 1.0,
|
|
accumulate: bool = False) -> None:
|
|
responses = self.collective_rpc("set_lora_adapter",
|
|
kwargs={
|
|
"lora_nickname": lora_nickname,
|
|
"lora_path": lora_path,
|
|
"strength": strength,
|
|
"accumulate": accumulate
|
|
})
|
|
for i, response in enumerate(responses):
|
|
if response["status"] != "lora_adapter_set":
|
|
raise RuntimeError(f"Worker {i} failed to set LoRA adapter to {lora_path}")
|
|
|
|
def unmerge_lora_weights(self) -> None:
|
|
responses = self.collective_rpc("unmerge_lora_weights", kwargs={})
|
|
for i, response in enumerate(responses):
|
|
if response["status"] != "lora_adapter_unmerged":
|
|
raise RuntimeError(f"Worker {i} failed to unmerge LoRA weights")
|
|
|
|
def merge_lora_weights(self) -> None:
|
|
responses = self.collective_rpc("merge_lora_weights", kwargs={})
|
|
for i, response in enumerate(responses):
|
|
if response["status"] != "lora_adapter_merged":
|
|
raise RuntimeError(f"Worker {i} failed to merge LoRA weights")
|
|
|
|
def set_log_queue(self, log_queue: Queue | None) -> None:
|
|
"""Forward worker logs to the given queue. Call before generate_video."""
|
|
self.collective_rpc("set_log_queue", kwargs={"log_queue": log_queue})
|
|
|
|
def clear_log_queue(self) -> None:
|
|
"""Stop forwarding worker logs to the queue. Call after generate_video."""
|
|
self.collective_rpc("clear_log_queue")
|
|
|
|
def collective_rpc(self,
|
|
method: str | Callable,
|
|
timeout: float | None = None,
|
|
args: tuple = (),
|
|
kwargs: dict | None = None) -> list[Any]:
|
|
kwargs = kwargs or {}
|
|
|
|
try:
|
|
for worker in self.workers:
|
|
worker.pipe.send({"method": method, "args": args, "kwargs": kwargs})
|
|
|
|
responses = []
|
|
for worker in self.workers:
|
|
response = worker.pipe.recv()
|
|
responses.append(response)
|
|
_raise_for_rpc_errors(method, responses)
|
|
return responses
|
|
except TimeoutError as e:
|
|
raise TimeoutError(f"RPC call to {method} timed out.") from e
|
|
except KeyboardInterrupt as e:
|
|
# if we catch a KeyboardInterrupt, user wants to stop the execution.
|
|
# we need to send a signal to all workers to stop.
|
|
logger.info("Received KeyboardInterrupt, sending SIGINT to all workers")
|
|
for worker in self.workers:
|
|
if worker.proc.pid is not None:
|
|
os.kill(worker.proc.pid, signal.SIGINT)
|
|
raise e
|
|
except Exception as e:
|
|
raise e
|
|
|
|
async def collective_rpc_async(self,
|
|
method: str | Callable,
|
|
timeout: float | None = None,
|
|
args: tuple = (),
|
|
kwargs: dict | None = None) -> list[Any]:
|
|
kwargs = kwargs or {}
|
|
loop = asyncio.get_running_loop()
|
|
|
|
for worker in self.workers:
|
|
worker.pipe.send({"method": method, "args": args, "kwargs": kwargs})
|
|
|
|
async def recv_from_worker(worker: WorkerProcHandle) -> Any:
|
|
return await loop.run_in_executor(None, worker.pipe.recv)
|
|
|
|
responses = await asyncio.gather(*[recv_from_worker(worker) for worker in self.workers])
|
|
result = list(responses)
|
|
_raise_for_rpc_errors(method, result)
|
|
return result
|
|
|
|
def shutdown(self) -> None:
|
|
"""Properly shut down the executor and its workers"""
|
|
if hasattr(self, 'shutting_down') and self.shutting_down:
|
|
return # Prevent multiple shutdown calls
|
|
|
|
logger.info("Shutting down MultiprocExecutor...")
|
|
|
|
# Check if workers were initialized (they might not be if initialization failed)
|
|
if not hasattr(self, 'workers') or not self.workers:
|
|
logger.info("No workers to shut down.")
|
|
return
|
|
|
|
self.shutting_down = True
|
|
|
|
# First try gentle termination
|
|
try:
|
|
# Send termination message to all workers
|
|
for worker in self.workers:
|
|
with contextlib.suppress(Exception):
|
|
worker.pipe.send({"method": "shutdown", "args": (), "kwargs": {}})
|
|
|
|
# Give workers some time to exit gracefully
|
|
start_time = time.perf_counter()
|
|
while time.perf_counter() - start_time < 5.0: # 5 seconds timeout
|
|
if all(not worker.proc.is_alive() for worker in self.workers):
|
|
break
|
|
time.sleep(0.1)
|
|
|
|
# Force terminate any remaining workers
|
|
for worker in self.workers:
|
|
if worker.proc.is_alive():
|
|
worker.proc.terminate()
|
|
|
|
# Final timeout for terminate
|
|
start_time = time.perf_counter()
|
|
while time.perf_counter() - start_time < 2.0: # 2 seconds timeout
|
|
if all(not worker.proc.is_alive() for worker in self.workers):
|
|
break
|
|
time.sleep(0.1)
|
|
|
|
# Kill if still alive
|
|
for worker in self.workers:
|
|
if worker.proc.is_alive():
|
|
worker.proc.kill()
|
|
worker.proc.join(timeout=1.0)
|
|
|
|
except Exception as e:
|
|
logger.error("Error during shutdown: %s", e)
|
|
# Last resort, try to kill all workers
|
|
for worker in self.workers:
|
|
with contextlib.suppress(Exception):
|
|
if worker.proc.is_alive():
|
|
worker.proc.kill()
|
|
|
|
# Clean up pipes
|
|
for worker in self.workers:
|
|
with contextlib.suppress(Exception):
|
|
worker.pipe.close()
|
|
|
|
self.workers = []
|
|
logger.info("MultiprocExecutor shutdown complete")
|
|
|
|
@staticmethod
|
|
def _ensure_worker_termination(worker_procs: list[BaseProcess]):
|
|
"""Ensure that all worker processes are terminated. Assumes workers have
|
|
received termination requests. Waits for processing, then sends
|
|
termination and kill signals if needed."""
|
|
|
|
def wait_for_termination(procs: list[BaseProcess], timeout: float) -> bool:
|
|
if not time:
|
|
# If we are in late stage shutdown, the interpreter may replace
|
|
# `time` with `None`.
|
|
return all(not proc.is_alive() for proc in procs)
|
|
start_time = time.time()
|
|
while time.time() - start_time < timeout:
|
|
if all(not proc.is_alive() for proc in procs):
|
|
return True
|
|
time.sleep(0.1)
|
|
return False
|
|
|
|
# Send SIGTERM if still running
|
|
active_procs = [proc for proc in worker_procs if proc.is_alive()]
|
|
for p in active_procs:
|
|
p.terminate()
|
|
if not wait_for_termination(active_procs, 4):
|
|
# Send SIGKILL if still running
|
|
active_procs = [p for p in active_procs if p.is_alive()]
|
|
for p in active_procs:
|
|
p.kill()
|
|
|
|
def __del__(self):
|
|
"""Ensure cleanup on garbage collection"""
|
|
self.shutdown()
|
|
|
|
def __enter__(self):
|
|
"""Support for context manager protocol"""
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
"""Ensure cleanup when exiting context"""
|
|
self.shutdown()
|
|
|
|
|
|
@dataclass
|
|
class UnreadyWorkerProcHandle:
|
|
"""WorkerProcess handle before READY."""
|
|
proc: BaseProcess
|
|
rank: int
|
|
pipe: Connection
|
|
ready_pipe: Connection
|
|
|
|
|
|
@dataclass
|
|
class WorkerProcHandle:
|
|
proc: BaseProcess
|
|
rank: int
|
|
pipe: Connection
|
|
|
|
@classmethod
|
|
def from_unready_handle(cls, unready_handle: UnreadyWorkerProcHandle) -> WorkerProcHandle:
|
|
return cls(
|
|
proc=unready_handle.proc,
|
|
rank=unready_handle.rank,
|
|
pipe=unready_handle.pipe,
|
|
)
|
|
|
|
|
|
class WorkerMultiprocProc:
|
|
"""Adapter that runs one Worker in busy loop."""
|
|
|
|
READY_STR = "READY"
|
|
|
|
def __init__(
|
|
self,
|
|
fastvideo_args: FastVideoArgs,
|
|
local_rank: int,
|
|
rank: int,
|
|
distributed_init_method: str,
|
|
pipe: Connection,
|
|
streaming_input_queue: Queue | None = None,
|
|
streaming_output_queue: Queue | None = None,
|
|
_initial_log_handler: logging.Handler | None = None,
|
|
**kwargs: Any,
|
|
):
|
|
self.rank = rank
|
|
self.pipe = pipe
|
|
self.streaming_input_queue = streaming_input_queue
|
|
self.streaming_output_queue = streaming_output_queue
|
|
self._initial_log_handler = _initial_log_handler
|
|
wrapper = WorkerWrapperBase(fastvideo_args=fastvideo_args, rpc_rank=rank)
|
|
|
|
all_kwargs: list[dict] = [{} for _ in range(fastvideo_args.num_gpus)]
|
|
all_kwargs[rank] = {
|
|
"fastvideo_args": fastvideo_args,
|
|
"local_rank": local_rank,
|
|
"rank": rank,
|
|
"distributed_init_method": distributed_init_method,
|
|
}
|
|
wrapper.init_worker(all_kwargs)
|
|
self.worker = wrapper
|
|
|
|
# Initialize device
|
|
self.worker.init_device()
|
|
|
|
# Set process title and log prefix
|
|
self.setup_proc_title_and_log_prefix()
|
|
|
|
@staticmethod
|
|
def make_worker_process(
|
|
fastvideo_args: FastVideoArgs,
|
|
local_rank: int,
|
|
rank: int,
|
|
distributed_init_method: str,
|
|
streaming_input_queue: Queue | None = None,
|
|
streaming_output_queue: Queue | None = None,
|
|
log_queue: Queue | None = None,
|
|
) -> UnreadyWorkerProcHandle:
|
|
context = get_mp_context()
|
|
executor_pipe, worker_pipe = context.Pipe(duplex=True)
|
|
reader, writer = context.Pipe(duplex=False)
|
|
|
|
process_kwargs = {
|
|
"fastvideo_args": fastvideo_args,
|
|
"local_rank": local_rank,
|
|
"rank": rank,
|
|
"distributed_init_method": distributed_init_method,
|
|
"pipe": worker_pipe,
|
|
"ready_pipe": writer,
|
|
"streaming_input_queue": streaming_input_queue,
|
|
"streaming_output_queue": streaming_output_queue,
|
|
"log_queue": log_queue,
|
|
}
|
|
# Run EngineCore busy loop in background process.
|
|
proc = context.Process(target=WorkerMultiprocProc.worker_main,
|
|
kwargs=process_kwargs,
|
|
name=f"FVWorkerProc-{rank}",
|
|
daemon=True)
|
|
|
|
proc.start()
|
|
worker_pipe.close()
|
|
return UnreadyWorkerProcHandle(proc, rank, executor_pipe, reader)
|
|
|
|
@staticmethod
|
|
def worker_main(*args, **kwargs):
|
|
""" Worker initialization and execution loops.
|
|
This runs a background process """
|
|
|
|
log_queue = kwargs.pop("log_queue", None)
|
|
# Add log handler before model loading so we capture fsdp_load, cuda, etc.
|
|
if log_queue is not None:
|
|
_handler = _make_queue_log_handler(log_queue)
|
|
logging.getLogger("fastvideo").addHandler(_handler)
|
|
kwargs["_initial_log_handler"] = _handler
|
|
|
|
# Signal handler used for graceful termination.
|
|
# SystemExit exception is only raised once to allow this and worker
|
|
# processes to terminate without error
|
|
shutdown_requested = False
|
|
|
|
def signal_handler(signum, frame):
|
|
nonlocal shutdown_requested
|
|
if not shutdown_requested:
|
|
shutdown_requested = True
|
|
raise SystemExit()
|
|
|
|
# Either SIGTERM or SIGINT will terminate the worker
|
|
signal.signal(signal.SIGTERM, signal_handler)
|
|
signal.signal(signal.SIGINT, signal_handler)
|
|
kill_itself_when_parent_died()
|
|
faulthandler.enable()
|
|
parent_process = psutil.Process().parent()
|
|
|
|
worker = None
|
|
ready_pipe = kwargs.pop("ready_pipe")
|
|
rank = kwargs.get("rank")
|
|
|
|
try:
|
|
worker = WorkerMultiprocProc(*args, **kwargs)
|
|
|
|
# Send READY once we know everything is loaded
|
|
ready_pipe.send({
|
|
"status": WorkerMultiprocProc.READY_STR,
|
|
})
|
|
|
|
ready_pipe.close()
|
|
ready_pipe = None
|
|
|
|
worker.worker_busy_loop()
|
|
|
|
except Exception as exc:
|
|
if ready_pipe is not None:
|
|
logger.exception("WorkerMultiprocProc failed to start.")
|
|
# Send error status to parent before closing pipe
|
|
try:
|
|
traceback_str = get_exception_traceback()
|
|
ready_pipe.send({
|
|
"status": "ERROR",
|
|
"error": str(exc),
|
|
"traceback": traceback_str,
|
|
"rank": rank,
|
|
})
|
|
except Exception:
|
|
# If sending fails, at least log it
|
|
pass
|
|
else:
|
|
logger.exception("WorkerMultiprocProc failed.")
|
|
|
|
# The parent sends a SIGTERM to all worker processes if
|
|
# any worker dies. Set this value so we don't re-throw
|
|
# SystemExit() to avoid zmq exceptions in __del__.
|
|
shutdown_requested = True
|
|
traceback = get_exception_traceback()
|
|
logger.error("Worker %d hit an exception: %s", rank, traceback)
|
|
if parent_process:
|
|
parent_process.send_signal(signal.SIGQUIT)
|
|
|
|
finally:
|
|
if ready_pipe is not None:
|
|
ready_pipe.close()
|
|
# Clean up once worker exits busy loop
|
|
if worker is not None:
|
|
worker.shutdown()
|
|
|
|
@staticmethod
|
|
def wait_for_ready(unready_proc_handles: list[UnreadyWorkerProcHandle]) -> list[WorkerProcHandle]:
|
|
|
|
e = Exception("WorkerMultiprocProc initialization failed due to "
|
|
"an exception in a background process. "
|
|
"See stack trace for root cause.")
|
|
|
|
pipes = {handle.ready_pipe: handle for handle in unready_proc_handles}
|
|
ready_proc_handles: list[WorkerProcHandle | None] = [None] * len(unready_proc_handles)
|
|
worker_errors: list[str] = []
|
|
while pipes:
|
|
ready = mp.connection.wait(pipes.keys())
|
|
for pipe in ready:
|
|
assert isinstance(pipe, Connection)
|
|
try:
|
|
# Wait until the WorkerProc is ready.
|
|
unready_proc_handle = pipes.pop(pipe)
|
|
response: dict[str, Any] = pipe.recv()
|
|
if response["status"] == "ERROR":
|
|
# Worker sent error details
|
|
error_msg = response.get("error", "Unknown error")
|
|
traceback_str = response.get("traceback", "")
|
|
rank = response.get("rank", "unknown")
|
|
error_info = f"Worker {rank} error: {error_msg}"
|
|
if traceback_str:
|
|
error_info += f"\n{traceback_str}"
|
|
worker_errors.append(error_info)
|
|
# Log a concise error message (full traceback will be in the exception)
|
|
logger.error("Worker %s initialization failed: %s", rank, error_msg)
|
|
# Continue to check other workers, but we'll fail at the end
|
|
elif response["status"] != "READY":
|
|
worker_errors.append(f"Worker returned unexpected status: {response.get('status', 'unknown')}")
|
|
|
|
if response["status"] == "READY":
|
|
ready_proc_handles[unready_proc_handle.rank] = (
|
|
WorkerProcHandle.from_unready_handle(unready_proc_handle))
|
|
|
|
except EOFError:
|
|
# Pipe closed without sending status - worker crashed
|
|
worker_errors.append("Worker process crashed (pipe closed unexpectedly)")
|
|
e.__suppress_context__ = True
|
|
raise e from None
|
|
|
|
finally:
|
|
# Close connection.
|
|
pipe.close()
|
|
|
|
# If any workers failed, raise exception with details
|
|
if worker_errors:
|
|
error_msg = "WorkerMultiprocProc initialization failed due to exceptions in background processes:\n"
|
|
error_msg += "\n".join(f" - {err}" for err in worker_errors)
|
|
raise Exception(error_msg) from None
|
|
|
|
logger.info("%d workers ready", len(ready_proc_handles))
|
|
return cast(list[WorkerProcHandle], ready_proc_handles)
|
|
|
|
def shutdown(self) -> dict[str, Any]:
|
|
return self.worker.shutdown()
|
|
|
|
def worker_busy_loop(self) -> None:
|
|
"""Main busy loop for Multiprocessing Workers"""
|
|
while True:
|
|
logger.info("Worker %d starting event loop...", self.rank)
|
|
try:
|
|
rpc_call = self.pipe.recv()
|
|
method = rpc_call.get("method")
|
|
args = rpc_call.get("args", ())
|
|
kwargs = rpc_call.get("kwargs", {})
|
|
|
|
if isinstance(method, str):
|
|
if method == "shutdown":
|
|
response = self.shutdown()
|
|
with contextlib.suppress(Exception):
|
|
self.pipe.send(response)
|
|
break
|
|
if method == "set_log_queue":
|
|
self._set_log_queue(kwargs.get("log_queue"))
|
|
self.pipe.send({"status": "ok"})
|
|
continue
|
|
if method == "clear_log_queue":
|
|
self._clear_log_queue()
|
|
self.pipe.send({"status": "ok"})
|
|
continue
|
|
if method == "start_streaming_queue_loop":
|
|
self.pipe.send({"status": "streaming_queue_loop_started"})
|
|
self.streaming_queue_loop()
|
|
continue
|
|
if method == 'execute_forward':
|
|
forward_batch = kwargs['forward_batch']
|
|
fastvideo_args = kwargs['fastvideo_args']
|
|
output_batch = self.worker.execute_forward(forward_batch, fastvideo_args)
|
|
logging_info = None
|
|
if envs.FASTVIDEO_STAGE_LOGGING:
|
|
logging_info = output_batch.logging_info
|
|
# result tensor shared by CUDA IPC to avoid serialization overhead
|
|
result = output_batch.output
|
|
extra = output_batch.extra or {}
|
|
extra["peak_memory_mb"] = (torch.cuda.max_memory_allocated() / (1024 * 1024))
|
|
self.pipe.send({
|
|
"output_batch": result,
|
|
"logging_info": logging_info,
|
|
"extra": extra,
|
|
"trajectory_latents": output_batch.trajectory_latents,
|
|
"trajectory_timesteps": output_batch.trajectory_timesteps,
|
|
})
|
|
else:
|
|
result = self.worker.execute_method(method, *args, **kwargs)
|
|
self.pipe.send(result)
|
|
else:
|
|
result = self.worker.execute_method(method, *args, **kwargs)
|
|
self.pipe.send(result)
|
|
except EOFError:
|
|
logger.info("Worker %d RPC pipe closed; exiting event loop", self.rank)
|
|
break
|
|
except KeyboardInterrupt:
|
|
logger.error("Worker %d in loop received KeyboardInterrupt, aborting forward pass", self.rank)
|
|
try:
|
|
self.pipe.send({"error": "Operation aborted by KeyboardInterrupt"})
|
|
logger.info("Worker %d sent error response after interrupt", self.rank)
|
|
except Exception as e:
|
|
logger.error("Worker %d failed to send error response: %s", self.rank, str(e))
|
|
continue
|
|
except Exception as error:
|
|
logger.exception("Worker %d failed RPC without exiting", self.rank)
|
|
try:
|
|
self.pipe.send({
|
|
_RPC_ERROR_KEY: True,
|
|
"error": f"{type(error).__name__}: {error}",
|
|
"traceback": get_exception_traceback(),
|
|
})
|
|
except Exception:
|
|
logger.exception("Worker %d could not return its RPC error", self.rank)
|
|
break
|
|
|
|
def streaming_queue_loop(self) -> None:
|
|
if self.streaming_input_queue is None or self.streaming_output_queue is None:
|
|
logger.error("Worker %d: streaming queues not initialized", self.rank)
|
|
return
|
|
|
|
while True:
|
|
try:
|
|
task: StreamingTask = self.streaming_input_queue.get()
|
|
|
|
if task.task_type == StreamingTaskType.EXIT:
|
|
break
|
|
elif task.task_type == StreamingTaskType.RESET:
|
|
try:
|
|
self.worker.execute_streaming_reset(task.batch, task.fastvideo_args)
|
|
self.streaming_output_queue.put(StreamingResult(task_type=StreamingTaskType.RESET))
|
|
except Exception as e:
|
|
logger.error("Worker %d reset error: %s", self.rank, e)
|
|
self.streaming_output_queue.put(StreamingResult(task_type=StreamingTaskType.RESET, error=e))
|
|
elif task.task_type == StreamingTaskType.STEP:
|
|
try:
|
|
batch = self.worker.execute_streaming_step(task.keyboard_action, task.mouse_action)
|
|
self.streaming_output_queue.put(
|
|
StreamingResult(task_type=StreamingTaskType.STEP, output_batch=batch))
|
|
except Exception as e:
|
|
logger.error("Worker %d step error: %s", self.rank, e)
|
|
self.streaming_output_queue.put(StreamingResult(task_type=StreamingTaskType.STEP, error=e))
|
|
elif task.task_type == StreamingTaskType.CLEAR:
|
|
self.worker.execute_streaming_clear()
|
|
self.streaming_output_queue.put(StreamingResult(task_type=StreamingTaskType.CLEAR))
|
|
except Exception as e:
|
|
logger.error("Worker %d queue loop error: %s", self.rank, e)
|
|
self.streaming_output_queue.put(StreamingResult(task_type=StreamingTaskType.STEP, error=e))
|
|
|
|
_log_queue_handler: logging.Handler | None = None
|
|
|
|
def _set_log_queue(self, log_queue: Queue | None) -> None:
|
|
"""Add a handler that forwards fastvideo logs to the given queue."""
|
|
self._clear_log_queue()
|
|
if log_queue is None:
|
|
return
|
|
# Remove initial handler if present (from worker_main) to avoid duplicates
|
|
if self._initial_log_handler is not None:
|
|
logging.getLogger("fastvideo").removeHandler(self._initial_log_handler)
|
|
self._initial_log_handler = None
|
|
self._log_queue_handler = _make_queue_log_handler(log_queue)
|
|
logging.getLogger("fastvideo").addHandler(self._log_queue_handler)
|
|
|
|
def _clear_log_queue(self) -> None:
|
|
"""Remove the log queue handler."""
|
|
if self._log_queue_handler is not None:
|
|
logging.getLogger("fastvideo").removeHandler(self._log_queue_handler)
|
|
self._log_queue_handler = None
|
|
|
|
@staticmethod
|
|
def setup_proc_title_and_log_prefix() -> None:
|
|
dp_size = get_dp_group().world_size
|
|
dp_rank = get_dp_group().rank_in_group
|
|
tp_size = get_tp_group().world_size
|
|
tp_rank = get_tp_group().rank_in_group
|
|
process_name = "Worker"
|
|
if dp_size > 1:
|
|
process_name += f"_DP{dp_rank}"
|
|
if tp_size > 1:
|
|
process_name += f"_TP{tp_rank}"
|
|
decorate_logs(process_name)
|
|
|
|
|
|
def set_multiproc_executor_envs() -> None:
|
|
""" Set up environment variables that should be used when there are workers
|
|
in a multiprocessing environment. This should be called by the parent
|
|
process before worker processes are created"""
|
|
|
|
force_spawn()
|