63 lines
2.4 KiB
Python
63 lines
2.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Process and platform setup for isolated collective test workers."""
|
|
|
|
from collections.abc import Callable
|
|
from multiprocessing.context import BaseContext
|
|
|
|
import torch
|
|
|
|
|
|
def collective_context(device_kind: str) -> BaseContext:
|
|
"""Reuse imports in a clean CPU forkserver; CUDA workers use spawn."""
|
|
if device_kind == "cpu" and "forkserver" in torch.multiprocessing.get_all_start_methods():
|
|
torch.multiprocessing.set_forkserver_preload(["tests.helpers.cpu_collective_server"])
|
|
return torch.multiprocessing.get_context("forkserver")
|
|
return torch.multiprocessing.get_context("spawn")
|
|
|
|
|
|
def start_collective_workers(
|
|
worker: Callable[..., None], *, args: tuple[object, ...], nprocs: int, device_kind: str
|
|
) -> None:
|
|
context = collective_context(device_kind)
|
|
if context.get_start_method() == "spawn":
|
|
torch.multiprocessing.spawn(worker, args=args, nprocs=nprocs)
|
|
else:
|
|
torch.multiprocessing.start_processes(worker, args=args, nprocs=nprocs, start_method=context.get_start_method())
|
|
|
|
|
|
def configure_cpu_collective_worker() -> None:
|
|
"""Select CPU platforms inside a spawned worker that only uses Gloo.
|
|
|
|
Upstream coordinators resolve local ranks through the selected
|
|
platform's visible devices. CPU ranks can outnumber visible GPUs, so they
|
|
must use the CPU platform even when the test host has CUDA installed.
|
|
The worker still creates real process groups and runs real collectives.
|
|
"""
|
|
import vllm.platforms
|
|
from vllm.platforms.cpu import CpuPlatform
|
|
|
|
from vllm_omni.diffusion.distributed import comm, group_coordinator, parallel_state
|
|
from vllm_omni.platforms.interface import UnspecifiedOmniPlatform
|
|
|
|
class CpuCollectivePlatform(UnspecifiedOmniPlatform):
|
|
dist_backend = "gloo"
|
|
|
|
@classmethod
|
|
def get_device_count(cls) -> int:
|
|
# CPU tensors share one logical device regardless of process rank.
|
|
return 1
|
|
|
|
@classmethod
|
|
def set_device(cls, device: torch.device) -> None:
|
|
assert device.type == "cpu"
|
|
|
|
@classmethod
|
|
def synchronize(cls) -> None:
|
|
pass
|
|
|
|
vllm.platforms.current_platform = CpuPlatform()
|
|
platform = CpuCollectivePlatform()
|
|
for module in (comm, group_coordinator, parallel_state):
|
|
module.current_omni_platform = platform
|