@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 5df2da5b97

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 18:30 UTC

[Misc] Better RayExecutor and multiprocessing compatibility (#14705)

Signed-off-by: Cody Yu <hao.yu.cody@gmail.com>

Cody Yu committed Mar 20, 2025 at 19:27 UTC 5df2da5b974d577cb98944c057de435dfd615cac

4 files changed +67 -21

vllm/engine/arg_utils.py

+14 -1

@@ -26,7 +26,7 @@ from vllm.plugins import load_general_plugins

26 from vllm.test_utils import MODEL_WEIGHTS_S3_BUCKET, MODELS_ON_S3

27 from vllm.transformers_utils.utils import check_gguf_file

28 from vllm.usage.usage_lib import UsageContext

29 -from vllm.utils import FlexibleArgumentParser, StoreBoolean

29 +from vllm.utils import FlexibleArgumentParser, StoreBoolean, is_in_ray_actor

31 if TYPE_CHECKING:

32 from vllm.transformers_utils.tokenizer_group import BaseTokenizerGroup

@@ -1245,6 +1245,18 @@ class EngineArgs:

1245 cpu_offload_gb=self.cpu_offload_gb,

1246 calculate_kv_scales=self.calculate_kv_scales,

1247 )

1248 +

1249 + # Get the current placement group if Ray is initialized and

1250 + # we are in a Ray actor. If so, then the placement group will be

1251 + # passed to spawned processes.

1252 + placement_group = None

1253 + if is_in_ray_actor():

1254 + import ray

1255 +

1256 + # This call initializes Ray automatically if it is not initialized,

1257 + # but we should not do this here.

1258 + placement_group = ray.util.get_current_placement_group()

1259 +

1260 parallel_config = ParallelConfig(

1261 pipeline_parallel_size=self.pipeline_parallel_size,

1262 tensor_parallel_size=self.tensor_parallel_size,

@@ -1257,6 +1269,7 @@ class EngineArgs:

1269 self.tokenizer_pool_extra_config,

1270 ),

1271 ray_workers_use_nsight=self.ray_workers_use_nsight,

1272 + placement_group=placement_group,

1273 distributed_executor_backend=self.distributed_executor_backend,

1274 worker_cls=self.worker_cls,

1275 worker_extension_cls=self.worker_extension_cls,

vllm/executor/multiproc_worker_utils.py

+2 -2

@@ -16,7 +16,7 @@ import torch

17 from vllm.config import VllmConfig

18 from vllm.logger import init_logger

19 -from vllm.utils import _check_multiproc_method, get_mp_context, run_method

19 +from vllm.utils import _maybe_force_spawn, get_mp_context, run_method

21 logger = init_logger(__name__)

@@ -291,7 +291,7 @@ def set_multiprocessing_worker_envs(parallel_config):

291 in a multiprocessing environment. This should be called by the parent

292 process before worker processes are created"""

293

294 - _check_multiproc_method()

294 + _maybe_force_spawn()

295

296 # Configure thread parallelism if OMP_NUM_THREADS isn't set

297 #

vllm/executor/ray_utils.py

+13 -8

@@ -284,8 +284,9 @@ def initialize_ray_cluster(

284 assert_ray_available()

285 from vllm.platforms import current_platform

286

287 - # Connect to a ray cluster.

288 - if current_platform.is_rocm() or current_platform.is_xpu():

287 + if ray.is_initialized():

288 + logger.info("Ray is already initialized. Skipping Ray initialization.")

289 + elif current_platform.is_rocm() or current_platform.is_xpu():

290 # Try to connect existing ray instance and create a new one if not found

291 try:

292 ray.init("auto", ignore_reinit_error=True)

@@ -299,19 +300,21 @@ def initialize_ray_cluster(

300 else:

301 ray.init(address=ray_address, ignore_reinit_error=True)

302

302 - if parallel_config.placement_group:

303 - # Placement group is already set.

304 - return

305 -

303 device_str = current_platform.ray_device_key

304 if not device_str:

305 raise ValueError(

306 f"current platform {current_platform.device_name} does not "

307 "support ray.")

308

312 - # Create placement group for worker processes

313 - current_placement_group = ray.util.get_current_placement_group()

309 + # Create or get the placement group for worker processes

310 + if parallel_config.placement_group:

311 + current_placement_group = parallel_config.placement_group

312 + else:

313 + current_placement_group = ray.util.get_current_placement_group()

314 +

315 if current_placement_group:

316 + logger.info("Using the existing placement group")

317 +

318 # We are in a placement group

319 bundles = current_placement_group.bundle_specs

320 # Verify that we can use the placement group.

@@ -331,6 +334,8 @@ def initialize_ray_cluster(

334 f"Required number of devices: {parallel_config.world_size}. "

335 f"Total number of devices: {device_bundles}.")

336 else:

337 + logger.info("No current placement group found. "

338 + "Creating a new placement group.")

339 num_devices_in_cluster = ray.cluster_resources().get(device_str, 0)

340 # Log a warning message and delay resource allocation failure response.

341 # Avoid immediate rejection to allow user-initiated placement group

vllm/utils.py

+38 -10

@@ -2147,20 +2147,48 @@ def zmq_socket_ctx(path: str, socket_type: Any) -> Iterator[zmq.Socket]:

2147 ctx.destroy(linger=0)

2148

2149

2150 -def _check_multiproc_method():

2151 - if (cuda_is_initialized()

2152 - and os.environ.get("VLLM_WORKER_MULTIPROC_METHOD") != "spawn"):

2153 - logger.warning("CUDA was previously initialized. We must use "

2154 - "the `spawn` multiprocessing start method. Setting "

2155 - "VLLM_WORKER_MULTIPROC_METHOD to 'spawn'. "

2156 - "See https://docs.vllm.ai/en/latest/getting_started/"

2157 - "troubleshooting.html#python-multiprocessing "

2158 - "for more information.")

2150 +def is_in_ray_actor():

2151 + """Check if we are in a Ray actor."""

2152 +

2153 + try:

2154 + import ray

2155 + return (ray.is_initialized()

2156 + and ray.get_runtime_context().get_actor_id() is not None)

2157 + except ImportError:

2158 + return False

2159 +

2160 +

2161 +def _maybe_force_spawn():

2162 + """Check if we need to force the use of the `spawn` multiprocessing start

2163 + method.

2164 + """

2165 + if os.environ.get("VLLM_WORKER_MULTIPROC_METHOD") == "spawn":

2166 + return

2167 +

2168 + reason = None

2169 + if cuda_is_initialized():

2170 + reason = "CUDA is initialized"

2171 + elif is_in_ray_actor():

2172 + reason = "In a Ray actor and can only be spawned"

2173 +

2174 + if reason is not None:

2175 + logger.warning(

2176 + "We must use the `spawn` multiprocessing start method. "

2177 + "Overriding VLLM_WORKER_MULTIPROC_METHOD to 'spawn'. "

2178 + "See https://docs.vllm.ai/en/latest/getting_started/"

2179 + "troubleshooting.html#python-multiprocessing "

2180 + "for more information. Reason: %s", reason)

2181 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"

2182

2183

2184 def get_mp_context():

2163 - _check_multiproc_method()

2185 + """Get a multiprocessing context with a particular method (spawn or fork).

2186 + By default we follow the value of the VLLM_WORKER_MULTIPROC_METHOD to

2187 + determine the multiprocessing method (default is fork). However, under

2188 + certain conditions, we may enforce spawn and override the value of

2189 + VLLM_WORKER_MULTIPROC_METHOD.

2190 + """

2191 + _maybe_force_spawn()

2192 mp_method = envs.VLLM_WORKER_MULTIPROC_METHOD

2193 return multiprocessing.get_context(mp_method)

2194