Skip to content

gpu_worker

Classes

fastvideo.worker.gpu_worker.Worker 

Worker(fastvideo_args: FastVideoArgs, local_rank: int, rank: int, distributed_init_method: str)
Source code in fastvideo/worker/gpu_worker.py 
27
28
29
30
31
def __init__(self, fastvideo_args: FastVideoArgs, local_rank: int, rank: int, distributed_init_method: str):
    self.fastvideo_args = fastvideo_args
    self.local_rank = local_rank
    self.rank = rank
    self.distributed_init_method = distributed_init_method

Methods:

fastvideo.worker.gpu_worker.Worker.init_device 
init_device() -> None

Initialize the device for the worker.

Source code in fastvideo/worker/gpu_worker.py 
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
def init_device(self) -> None:
    """Initialize the device for the worker."""

    # torch.distributed.all_reduce does not free the input tensor until
    # the synchronization point. This causes the memory usage to grow
    # as the number of all_reduce calls increases. This env var disables
    # this behavior.
    # Related issue:
    # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573
    os.environ["TORCH_NCCL_AVOID_RECORD_STREAMS"] = "1"
    # This env var set by Ray causes exceptions with graph building.
    os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)

    # Set environment variables BEFORE calling get_local_torch_device()
    # so that each worker uses the correct device
    # Both multiprocessing and Ray pass the worker-local rank explicitly.
    # Ray deliberately excludes LOCAL_RANK from the copied driver
    # environment and exposes all GPUs assigned to the node, so leaving an
    # inherited or missing value here would bind every Ray actor to cuda:0.
    os.environ["LOCAL_RANK"] = str(self.local_rank)
    os.environ["RANK"] = str(self.rank)
    os.environ["WORLD_SIZE"] = str(self.fastvideo_args.num_gpus)

    # Platform-agnostic device initialization
    self.device = get_local_torch_device()

    from fastvideo.platforms import current_platform

    # Set the CUDA device BEFORE any CUDA calls
    if current_platform.is_cuda_alike():
        torch.cuda.set_device(self.device)
        self.init_gpu_memory = torch.cuda.mem_get_info(self.device)[0]
        if current_platform.is_cuda():
            _log_cuda_device_uuid(self.rank, self.device)
    else:
        # For MPS, we can't get memory info the same way
        self.init_gpu_memory = 0

    # CUDA's unified-memory classification reads runtime device
    # properties, so make this decision only after this worker has bound
    # its own device. The worker-local args object is what every loader and
    # pipeline stage below will consume.
    device_id = self.device.index if self.device.index is not None else 0
    self.fastvideo_args.finalize_device_offload_policy(device_id)

    # Initialize the distributed environment.
    maybe_init_distributed_environment_and_model_parallel(self.fastvideo_args.tp_size, self.fastvideo_args.sp_size,
                                                          self.distributed_init_method)

    self.pipeline = build_pipeline(self.fastvideo_args)
fastvideo.worker.gpu_worker.Worker.shutdown 
shutdown() -> dict[str, Any]

Gracefully shut down the worker process

Source code in fastvideo/worker/gpu_worker.py 
106
107
108
109
110
111
112
113
114
115
116
117
118
def shutdown(self) -> dict[str, Any]:
    """Gracefully shut down the worker process"""
    logger.info("Worker %d shutting down...", self.rank, local_main_process_only=False)
    # Clean up resources
    if hasattr(self, 'pipeline') and self.pipeline is not None:
        # Clean up pipeline resources if needed
        pass

    # Destroy the distributed environment
    cleanup_dist_env_and_memory(shutdown_ray=False)

    logger.info("Worker %d shutdown complete", self.rank, local_main_process_only=False)
    return {"status": "shutdown_complete"}

Functions: