[Feat] support using external ray pg with bundles (#3850)

Author: CyCle1024

lmdeploy/pytorch/engine/executor/ray_executor.py

@@ -524,12 +524,19 @@ def get_priority(ip):
         sorted_workers = [item[0] for item in sorted_worker_ip_map]
         return sorted_workers
 
+    def _valid_bundle_id(self, bundle_id: int):
+        """Check if a bundle is valid only when self.use_external_ray=True."""
+        if (not self.ray_ctx.owned_pg and _envs.ray_external_pg_bundles
+                and bundle_id not in _envs.ray_external_pg_bundles):
+            return False
+        return True
+
     def _init_workers_ray(self, placement_group: PlacementGroup, worker_kwargs: dict):
         """Init worker ray."""
         device_str = get_device_str()
         bundle_indices = []
         for bundle_id, bundle in enumerate(placement_group.bundle_specs):
-            if bundle.get(device_str, 0):
+            if bundle.get(device_str, 0) and self._valid_bundle_id(bundle_id):
                 bundle_indices.append(bundle_id)
         bundle_indices = bundle_indices[:self.world_size]
 
@@ -556,7 +563,7 @@ def _init_workers_ray(self, placement_group: PlacementGroup, worker_kwargs: dict
                 worker = ray.remote(
                     num_cpus=0,
                     num_gpus=0,
-                    resources={device_str: 1.0},
+                    resources={device_str: 0.01},
                     scheduling_strategy=scheduling_strategy,
                 )(RayWorkerWrapper).remote(**worker_kwargs)
             workers.append(worker)

lmdeploy/pytorch/envs.py

@@ -40,6 +40,22 @@ def env_to_int(
     return value
 
 
+def env_to_list_int(
+    env_var: str,
+    default: list[int] = None,
+):
+    """Env to list of int."""
+    default_ = default if default is not None else []
+    value = os.getenv(env_var)
+    if value is None:
+        return default_
+    try:
+        value = [int(x) for x in value.split(',')]
+    except Exception:
+        value = default_
+    return value
+
+
 _ENVS = dict()
 
 
@@ -91,6 +107,10 @@ def _patched_get_env(
     ray_timeline_enable = env_to_bool('LMDEPLOY_RAY_TIMELINE_ENABLE', False)
     ray_timeline_output_path = os.getenv('LMDEPLOY_RAY_TIMELINE_OUT_PATH', 'ray_timeline.json')
 
+    # ray external placement group bundles
+    # only used when lmdeploy is initialized inside a Ray Actor with pg allocated
+    ray_external_pg_bundles = env_to_list_int('LMDEPLOY_RAY_EXTERNAL_PG_BUNDLES', [])
+
     # dist
     dist_master_addr = os.getenv('LMDEPLOY_DIST_MASTER_ADDR', None)
     dist_master_port = os.getenv('LMDEPLOY_DIST_MASTER_PORT', None)