[DP][ray] Support different VLLM_RAY_DP_PACK_STRATEGY

Signed-off-by: Rui Qiao <ruisearch42@gmail.com>

Author: ruisearch42

vllm/envs.py

@@ -133,6 +133,7 @@
     VLLM_DP_MASTER_PORT: int = 0
     VLLM_MOE_DP_CHUNK_SIZE: int = 256
     VLLM_RANDOMIZE_DP_DUMMY_INPUTS: bool = False
+    VLLM_RAY_DP_PACK_STRATEGY: str = "fill"
     VLLM_MARLIN_USE_ATOMIC_ADD: bool = False
     VLLM_MXFP4_USE_MARLIN: Optional[bool] = None
     VLLM_V0_USE_OUTLINES_CACHE: bool = False
@@ -997,10 +998,19 @@ def get_vllm_port() -> Optional[int]:
     # units.
     "VLLM_MOE_DP_CHUNK_SIZE": lambda: int(os.getenv("VLLM_MOE_DP_CHUNK_SIZE", "256")),
     # Randomize inputs during dummy runs when using Data Parallel
-    "VLLM_RANDOMIZE_DP_DUMMY_INPUTS": lambda: os.environ.get(
-        "VLLM_RANDOMIZE_DP_DUMMY_INPUTS", "0"
-    )
-    == "1",
+    "VLLM_RANDOMIZE_DP_DUMMY_INPUTS":
+    lambda: os.environ.get("VLLM_RANDOMIZE_DP_DUMMY_INPUTS", "0") == "1",
+
+    # Strategy to pack the data parallel ranks for Ray.
+    # Available options:
+    # - "fill":
+    #   for DP master node, allocate exactly data-parallel-size-local DP ranks,
+    #   for non-master nodes, allocate as many DP ranks as can fit;
+    # - "strict":
+    #   allocate exactly data-parallel-size-local DP ranks to each picked node;
+    # This environment variable is ignored if data-parallel-backend is not Ray.
+    "VLLM_RAY_DP_PACK_STRATEGY":
+    lambda: os.getenv("VLLM_RAY_DP_PACK_STRATEGY", "fill"),
     # Whether to use S3 path for model loading in CI via RunAI Streamer
     "VLLM_CI_USE_S3": lambda: os.environ.get("VLLM_CI_USE_S3", "0") == "1",
     # Use model_redirect to redirect the model name to a local folder.

vllm/v1/engine/utils.py

@@ -15,6 +15,7 @@
 import msgspec
 import zmq
 
+from vllm import envs
 from vllm.config import CacheConfig, ParallelConfig, VllmConfig
 from vllm.logger import init_logger
 from vllm.platforms import current_platform
@@ -335,9 +336,11 @@ def create_dp_placement_groups(
         from ray._private.state import available_resources_per_node
 
         logger.info("Creating placement groups for data parallel")
-        dp_master_ip = vllm_config.parallel_config.data_parallel_master_ip
-        num_pg_to_create = vllm_config.parallel_config.data_parallel_size
-        local_engine_count = vllm_config.parallel_config.data_parallel_size_local
+        dp_master_ip = \
+            vllm_config.parallel_config.data_parallel_master_ip
+        dp_size = vllm_config.parallel_config.data_parallel_size
+        dp_size_local = \
+            vllm_config.parallel_config.data_parallel_size_local
 
         available_resources = available_resources_per_node()
         world_size = vllm_config.parallel_config.world_size
@@ -349,53 +352,76 @@ def create_dp_placement_groups(
         )
         assert len(nodes) > 0, "No nodes with resources found in Ray cluster."
         assert dp_master_ip_key in nodes[0], (
-            "The DP master node (ip: %s) is missing or dead",
-            dp_master_ip,
-        )
-        device_str = current_platform.ray_device_key
+            "The DP master node (ip: %s) is missing or dead", dp_master_ip)
+
+        if envs.VLLM_RAY_DP_PACK_STRATEGY == "strict":
+            logger.info(
+                "Using strict local size packing strategy based "
+                "on VLLM_RAY_DP_PACK_STRATEGY (%s)",
+                envs.VLLM_RAY_DP_PACK_STRATEGY)
+            strict_local_size = True
+        elif (envs.VLLM_ALL2ALL_BACKEND == "deepep_high_throughput"
+              or envs.VLLM_ALL2ALL_BACKEND == "deepep_low_latency"):
+            logger.info(
+                "Using strict local size packing strategy based "
+                "on VLLM_ALL2ALL_BACKEND (%s)", envs.VLLM_ALL2ALL_BACKEND)
+            strict_local_size = True
+        else:
+            logger.info(
+                "Using fill packing strategy based "
+                "on VLLM_RAY_DP_PACK_STRATEGY (%s)",
+                envs.VLLM_RAY_DP_PACK_STRATEGY)
+            strict_local_size = False
         for node_resources in nodes:
-            if device_str not in node_resources:
-                continue
-            # For now, each DP rank can only be assigned to one node
-            # TODO(rui): support allocating a single DP rank
-            # to multiple nodes
-            available_engine_count = int(node_resources[device_str]) // world_size
-            if dp_master_ip_key in node_resources:
-                assert available_engine_count >= local_engine_count, (
-                    "Not enough resources to allocate DP ranks "
-                    f"on DP master node {dp_master_ip}"
-                )
-                for i in range(local_engine_count):
-                    bundles = [
-                        {device_str: 1.0, "node:" + dp_master_ip: 0.001}
-                    ] * world_size + [{"CPU": 1.0}]
-                    pg = ray.util.placement_group(
-                        name=f"dp_rank_{len(placement_groups)}",
-                        strategy="STRICT_PACK",
-                        bundles=bundles,
-                    )
-                    placement_groups.append(pg)
-                    local_dp_ranks.append(i)
+            node_ip_keys = [
+                key for key in node_resources if key.startswith('node:')
+            ]
+            assert len(node_ip_keys) == 1, (
+                "Zero or multiple node IP keys found in node resources: %s",
+                node_ip_keys)
+            node_ip_key = node_ip_keys[0]
+            node_ip = node_ip_key.split(":")[1]
+            dp_size_available = int(node_resources["GPU"]) // world_size
+            if strict_local_size:
+                if dp_size_available < dp_size_local:
+                    if node_ip == dp_master_ip:
+                        raise ValueError(
+                            "Not enough resources to allocate %s DP ranks "
+                            "on DP master node %s, possible to fit %s DP ranks",
+                            dp_size_local, dp_master_ip, dp_size_available)
+                    else:
+                        logger.info(
+                            "Skipping node %s as %s DP ranks could not fit, "
+                            "possible to fit %s DP ranks", node_ip,
+                            dp_size_local, dp_size_available)
+                        continue
+                dp_size_to_allocate = dp_size_local
+            elif node_ip == dp_master_ip:
+                dp_size_to_allocate = dp_size_local
             else:
-                for i in range(available_engine_count):
-                    if len(placement_groups) == num_pg_to_create:
-                        break
-                    bundles = [{device_str: 1.0}] * world_size + [{"CPU": 1.0}]
-                    pg = ray.util.placement_group(
-                        name=f"dp_rank_{len(placement_groups)}",
-                        strategy="STRICT_PACK",
-                        bundles=bundles,
-                    )
-                    placement_groups.append(pg)
-                    local_dp_ranks.append(i)
-        if len(placement_groups) < num_pg_to_create:
-            raise ValueError(
-                f"Not enough resources to allocate {num_pg_to_create} "
-                "placement groups, only created "
-                f"{len(placement_groups)} placement groups. "
-                "Available resources: "
-                f"{available_resources}"
-            )
+                dp_size_to_allocate = dp_size_available
+
+            for i in range(dp_size_to_allocate):
+                bundles = [{
+                    "GPU": 1.0,
+                    "node:" + node_ip: 0.001
+                }] * world_size + [{
+                    "CPU": 1.0
+                }]
+                pg = ray.util.placement_group(
+                    name=f"dp_rank_{len(placement_groups)}",
+                    strategy="STRICT_PACK",
+                    bundles=bundles,
+                )
+                placement_groups.append(pg)
+                local_dp_ranks.append(i)
+
+        if len(placement_groups) < dp_size:
+            raise ValueError(f"Not enough resources to allocate {dp_size} "
+                             "placement groups, only created "
+                             f"{len(placement_groups)} placement groups. "
+                             "Available resources: "
+                             f"{available_resources}")
         return placement_groups, local_dp_ranks
 
     @staticmethod