154 lines
5.6 KiB
Python
154 lines
5.6 KiB
Python
import atexit
|
|
import contextlib
|
|
import multiprocessing as mp
|
|
import time
|
|
from multiprocessing.process import BaseProcess
|
|
from typing import Any, Callable, List, Optional, Union, cast
|
|
|
|
from fastvideo.v1.fastvideo_args import FastVideoArgs
|
|
from fastvideo.v1.logger import init_logger
|
|
from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
|
|
from fastvideo.v1.worker.executor import Executor
|
|
from fastvideo.v1.worker.gpu_worker import run_worker_process
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class MultiprocExecutor(Executor):
|
|
|
|
def _init_executor(self) -> None:
|
|
self.world_size = self.fastvideo_args.num_gpus
|
|
self.shutting_down = False
|
|
|
|
# this will force the use of the `spawn` multiprocessing start if cuda
|
|
# is initialized
|
|
mp.set_start_method("spawn", force=True)
|
|
|
|
self.workers: List[BaseProcess] = []
|
|
self.worker_pipes = []
|
|
|
|
# Create pipes and start workers
|
|
for rank in range(self.world_size):
|
|
executor_pipe, worker_pipe = mp.Pipe(duplex=True)
|
|
self.worker_pipes.append(executor_pipe)
|
|
|
|
worker = mp.Process(target=run_worker_process,
|
|
name=f"FVWorkerProc-{rank}",
|
|
kwargs=dict(fastvideo_args=self.fastvideo_args,
|
|
local_rank=rank,
|
|
rank=rank,
|
|
pipe=worker_pipe))
|
|
worker.start()
|
|
self.workers.append(worker)
|
|
logger.info("Workers: %s", self.workers)
|
|
|
|
# Wait for all workers to be ready
|
|
for idx, pipe in enumerate(self.worker_pipes):
|
|
data = pipe.recv()
|
|
if data["status"] != "ready" or data["local_rank"] != idx:
|
|
raise RuntimeError(f"Worker {idx} failed to start")
|
|
logger.info("%d workers ready", self.world_size)
|
|
|
|
# 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
|
|
})
|
|
return cast(ForwardBatch, responses[0]["output_batch"])
|
|
|
|
def collective_rpc(self,
|
|
method: Union[str, Callable],
|
|
timeout: Optional[float] = None,
|
|
args: tuple = (),
|
|
kwargs: Optional[dict] = None) -> list[Any]:
|
|
kwargs = kwargs or {}
|
|
|
|
try:
|
|
for pipe in self.worker_pipes:
|
|
pipe.send({"method": method, "args": args, "kwargs": kwargs})
|
|
|
|
responses = []
|
|
for pipe in self.worker_pipes:
|
|
response = pipe.recv()
|
|
responses.append(response)
|
|
return responses
|
|
except TimeoutError as e:
|
|
raise TimeoutError(f"RPC call to {method} timed out.") from e
|
|
except Exception as e:
|
|
# Re-raise any other exceptions
|
|
raise e
|
|
|
|
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...")
|
|
self.shutting_down = True
|
|
|
|
# First try gentle termination
|
|
try:
|
|
# Send termination message to all workers
|
|
for pipe in self.worker_pipes:
|
|
with contextlib.suppress(Exception):
|
|
pipe.send({"method": "shutdown", "args": (), "kwargs": {}})
|
|
|
|
# Give workers some time to exit gracefully
|
|
start_time = time.time()
|
|
while time.time() - start_time < 5.0: # 5 seconds timeout
|
|
if all(not worker.is_alive() for worker in self.workers):
|
|
break
|
|
time.sleep(0.1)
|
|
|
|
# Force terminate any remaining workers
|
|
for worker in self.workers:
|
|
if worker.is_alive():
|
|
worker.terminate()
|
|
|
|
# Final timeout for terminate
|
|
start_time = time.time()
|
|
while time.time() - start_time < 2.0: # 2 seconds timeout
|
|
if all(not worker.is_alive() for worker in self.workers):
|
|
break
|
|
time.sleep(0.1)
|
|
|
|
# Kill if still alive
|
|
for worker in self.workers:
|
|
if worker.is_alive():
|
|
worker.kill()
|
|
worker.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.is_alive():
|
|
worker.kill()
|
|
|
|
# Clean up pipes
|
|
for pipe in self.worker_pipes:
|
|
with contextlib.suppress(Exception):
|
|
pipe.close()
|
|
|
|
self.workers = []
|
|
self.worker_pipes = []
|
|
logger.info("MultiprocExecutor shutdown complete")
|
|
|
|
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()
|