fix: cluster multiple ray instance support (RLinf#376)

andylin-hao · web-flow · commit 5ce3aa252a05 · 2025-12-03T21:17:24.000+08:00
Signed-off-by: Hao Lin &lt;linhaomails@gmail.com&gt;
diff --git a/examples/coding_online_rl/main_coding_online_rl.py b/examples/coding_online_rl/main_coding_online_rl.py
@@ -40,7 +40,7 @@ def main(cfg) -> None:
     cfg = validate_cfg(cfg)
     print(json.dumps(OmegaConf.to_container(cfg, resolve=True), indent=2))
 
-    cluster = Cluster(num_nodes=cfg.cluster.num_nodes)
+    cluster = Cluster(cluster_cfg=cfg.cluster)
     component_placement = ModelParallelComponentPlacement(cfg, cluster)
 
     singleton_placement_strategy = PackedPlacementStrategy(
diff --git a/examples/coding_online_rl/main_coding_rl_llm_judge.py b/examples/coding_online_rl/main_coding_rl_llm_judge.py
@@ -40,7 +40,7 @@ def main(cfg) -> None:
     cfg = validate_cfg(cfg)
     print(json.dumps(OmegaConf.to_container(cfg, resolve=True), indent=2))
 
-    cluster = Cluster(num_nodes=cfg.cluster.num_nodes)
+    cluster = Cluster(cluster_cfg=cfg.cluster)
     component_placement = ModelParallelComponentPlacement(cfg, cluster)
 
     rollout_worker_cls = get_rollout_backend_worker(cfg, component_placement)
diff --git a/examples/embodiment/eval_embodied_agent.py b/examples/embodiment/eval_embodied_agent.py
@@ -32,7 +32,7 @@ def main(cfg) -> None:
     cfg = validate_cfg(cfg)
     cfg.runner.only_eval = True
 
-    cluster = Cluster(num_nodes=cfg.cluster.num_nodes)
+    cluster = Cluster(cluster_cfg=cfg.cluster)
     component_placement = HybridComponentPlacement(cfg, cluster)
 
     # Create rollout worker group
diff --git a/examples/embodiment/train_embodied_agent.py b/examples/embodiment/train_embodied_agent.py
@@ -36,7 +36,7 @@ def main(cfg) -> None:
     cfg = validate_cfg(cfg)
     print(json.dumps(OmegaConf.to_container(cfg, resolve=True), indent=2))
 
-    cluster = Cluster(num_nodes=cfg.cluster.num_nodes)
+    cluster = Cluster(cluster_cfg=cfg.cluster)
     component_placement = HybridComponentPlacement(cfg, cluster)
 
     # Create actor worker group
diff --git a/examples/multiturn_demo/main_mcp_with_session.py b/examples/multiturn_demo/main_mcp_with_session.py
@@ -43,7 +43,7 @@ def main(cfg) -> None:
     cfg = validate_cfg(cfg)
     print(json.dumps(OmegaConf.to_container(cfg, resolve=True), indent=2))
 
-    cluster = Cluster(num_nodes=cfg.cluster.num_nodes)
+    cluster = Cluster(cluster_cfg=cfg.cluster)
     component_placement = ModelParallelComponentPlacement(cfg, cluster)
 
     # Generator group
diff --git a/examples/multiturn_demo/main_tool.py b/examples/multiturn_demo/main_tool.py
@@ -43,7 +43,7 @@ def main(cfg) -> None:
     cfg = validate_cfg(cfg)
     print(json.dumps(OmegaConf.to_container(cfg, resolve=True), indent=2))
 
-    cluster = Cluster(num_nodes=cfg.cluster.num_nodes)
+    cluster = Cluster(cluster_cfg=cfg.cluster)
     component_placement = ModelParallelComponentPlacement(cfg, cluster)
 
     # Generator group
diff --git a/examples/reasoning/main_grpo.py b/examples/reasoning/main_grpo.py
@@ -41,7 +41,7 @@ def main(cfg) -> None:
     cfg = validate_cfg(cfg)
     print(json.dumps(OmegaConf.to_container(cfg, resolve=True), indent=2))
 
-    cluster = Cluster(num_nodes=cfg.cluster.num_nodes)
+    cluster = Cluster(cluster_cfg=cfg.cluster)
     component_placement = ModelParallelComponentPlacement(cfg, cluster)
 
     rollout_worker_cls = get_rollout_backend_worker(cfg, component_placement)
diff --git a/rlinf/scheduler/cluster/cluster.py b/rlinf/scheduler/cluster/cluster.py
@@ -77,6 +77,9 @@ class Cluster:
         ClusterEnvVar.COMM_NET_DEVICES: None,
     }
 
+    class NamespaceConflictError(Exception):
+        """Raised when there is a namespace conflict in Ray initialization."""
+
     @classmethod
     def find_free_port(cls):
         """Find a free port on the node."""
@@ -109,16 +112,26 @@ def __init__(
         """
         if self._has_initialized:
             return
+        self._setup_logger()
         if num_nodes is not None or cluster_cfg is not None:
             self._ray_instance_count = 0
-            self._init_and_launch_managers(num_nodes, cluster_cfg)
+            while True:
+                try:
+                    self._init_and_launch_managers(num_nodes, cluster_cfg)
+                    break
+                except Cluster.NamespaceConflictError:
+                    # Switch the namespace when multiple ray instances are created in the same node
+                    self._ray_instance_count += 1
+                    self._logger.info(
+                        f"Ray namespace conflict detected. Retrying to initialize Cluster with a new namespace (attempt {self._ray_instance_count})."
+                    )
+                    Cluster.NAMESPACE = f"{Cluster.SYS_NAME}_{self._ray_instance_count}"
+                    continue
         else:
             self._init_from_existing_managers()
         self._has_initialized = True
 
-    def _init_and_launch_managers(
-        self, num_nodes: int, cluster_cfg: Optional[DictConfig]
-    ):
+    def _setup_logger(self):
         # Add logger
         self._logger = logging.getLogger(Cluster.SYS_NAME)
         self._logger.setLevel(Cluster.LOGGING_LEVEL)
@@ -133,6 +146,9 @@ def _init_and_launch_managers(
         handler.setFormatter(formatter)
         self._logger.addHandler(handler)
 
+    def _init_and_launch_managers(
+        self, num_nodes: int, cluster_cfg: Optional[DictConfig]
+    ):
         if ray.is_initialized():
             if self._ray_instance_count > 0:
                 # For reinit Ray to switch namespace
@@ -235,10 +251,7 @@ def _init_and_launch_managers(
                 .remote()
             )
         except ValueError:
-            # If the WorkerManager is already running, we need to switch the namespace
-            self._ray_instance_count += 1
-            Cluster.NAMESPACE = f"RLinf_{self._ray_instance_count}"
-            return self._init_and_launch_managers(num_nodes)
+            raise Cluster.NamespaceConflictError
 
         def signal_handler(sig, frame):
             # Exit the main process if SIGUSR1 is received, which is sent by the worker group when an exception occurs.
diff --git a/rlinf/scheduler/cluster/config.py b/rlinf/scheduler/cluster/config.py
@@ -273,10 +273,14 @@ def from_dict_cfg(cfg_dict: DictConfig) -> "ClusterConfig":
         Returns:
             ClusterConfig: The created ClusterConfig instance.
         """
-        dataclass_arg_check(
-            ClusterConfig, cfg_dict, error_suffix="in cluster yaml config"
+        _, _, valid_args = dataclass_arg_check(
+            ClusterConfig,
+            cfg_dict,
+            no_check_unknown=True,
+            error_suffix="in cluster yaml config",
         )
-        return ClusterConfig(**cfg_dict)
+        valid_cfg_dict = {key: cfg_dict[key] for key in valid_args if key in cfg_dict}
+        return ClusterConfig(**valid_cfg_dict)
 
     def get_node_labels_by_rank(self, node_rank: int) -> list[str]:
         """Get the node group labels for a given node rank.
diff --git a/rlinf/scheduler/cluster/node.py b/rlinf/scheduler/cluster/node.py
@@ -284,12 +284,15 @@ def __init__(self, cluster_num_nodes: int, cluster_cfg: Optional[ClusterConfig])
         num_nodes = len(node_infos)
         for node_info in node_infos:
             node_ray_id = node_info["NodeID"]
-            probe = _RemoteNodeProbe.options(
-                scheduling_strategy=ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy(
-                    node_id=node_ray_id, soft=False
-                ),
-                name=f"NodeProbe_{node_ray_id}",
-            ).remote(node_info, num_nodes, cluster_cfg, sys.executable)
+            try:
+                probe = _RemoteNodeProbe.options(
+                    scheduling_strategy=ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy(
+                        node_id=node_ray_id, soft=False
+                    ),
+                    name=f"NodeProbe_{node_ray_id}",
+                ).remote(node_info, num_nodes, cluster_cfg, sys.executable)
+            except ValueError:
+                raise Cluster.NamespaceConflictError
             self._probes.append(probe)
 
         handles = []
diff --git a/rlinf/scheduler/cluster/utils.py b/rlinf/scheduler/cluster/utils.py
@@ -89,13 +89,17 @@ def parse_rank_config(
 
 
 def dataclass_arg_check(
-    dataclass: DataclassProtocol, kwargs: dict, error_suffix: str = ""
+    dataclass: DataclassProtocol,
+    kwargs: dict,
+    no_check_unknown: bool = False,
+    error_suffix: str = "",
 ):
     """Check if the kwargs contain only valid fields for the given dataclass.
 
     Args:
         dataclass (DataclassProtocol): The dataclass to check against.
         kwargs (dict): The keyword arguments to check.
+        no_check_unknown (bool): Whether to skip checking for unknown fields.
         error_suffix (str): Additional error message suffix.
     """
     args = set(kwargs.keys())
@@ -115,8 +119,9 @@ def dataclass_arg_check(
     assert not missing_required_args, (
         f"Missing fields '{missing_required_args}' detected {error_suffix}. Only got: {kwargs.keys()}."
     )
-    assert not unknown_args, (
-        f"Unknown fields '{unknown_args}' detected {error_suffix}. Valid fields are: {valid_args}."
-    )
+    if not no_check_unknown:
+        assert not unknown_args, (
+            f"Unknown fields '{unknown_args}' detected {error_suffix}. Valid fields are: {valid_args}."
+        )
 
     return missing_required_args, unknown_args, valid_args
diff --git a/rlinf/scheduler/hardware/robots/franka.py b/rlinf/scheduler/hardware/robots/franka.py
@@ -14,6 +14,7 @@
 
 import importlib
 import ipaddress
+import warnings
 from dataclasses import dataclass
 from typing import Optional
 
@@ -95,10 +96,19 @@ def enumerate(
                         count=cls.ROBOT_PING_COUNT,
                         timeout=cls.ROBOT_PING_TIMEOUT,
                     )
-                except Exception as e:
+                except ConnectionError as e:
                     raise ConnectionError(
                         f"Cannot reach Franka robot at IP {config.robot_ip} from node rank {node_rank}. Error: {e}"
                     )
+                except PermissionError as e:
+                    warnings.warn(
+                        f"Permission denied when trying to ping Franka robot at IP {config.robot_ip} from node rank {node_rank}. "
+                        f"This may be due to insufficient permissions to send ICMP packets. Ignoring the ping test. Error: {e}"
+                    )
+                except Exception as e:
+                    warnings.warn(
+                        f"An unexpected error occurred while pinging Franka robot at IP {config.robot_ip} from node rank {node_rank}. Ignoring the ping test. Error: {e}"
+                    )
                 if not response.is_alive:
                     raise ConnectionError(
                         f"Cannot reach Franka robot at IP {config.robot_ip} from node rank {node_rank}."
diff --git a/rlinf/scheduler/placement/flexible.py b/rlinf/scheduler/placement/flexible.py
@@ -199,7 +199,9 @@ def get_placement(
                     accelerator_type=cluster.get_node_info(
                         cluster_node_rank
                     ).accelerator_type,
-                    local_accelerator_rank=local_accel_ranks[0],
+                    local_accelerator_rank=local_accel_ranks[0]
+                    if len(local_accel_ranks) > 0
+                    else -1,
                     local_rank=-1,
                     local_world_size=0,
                     visible_accelerators=visible_accelerators,
diff --git a/rlinf/scheduler/placement/packed.py b/rlinf/scheduler/placement/packed.py
@@ -221,7 +221,9 @@ def get_placement(
                     accelerator_type=cluster.get_node_info(
                         cluster_node_rank
                     ).accelerator_type,
-                    local_accelerator_rank=local_accel_ranks[0],
+                    local_accelerator_rank=local_accel_ranks[0]
+                    if len(local_accel_ranks) > 0
+                    else -1,
                     local_rank=local_rank,
                     local_world_size=0,
                     visible_accelerators=visible_accelerators,
