Update slurm template for RDMA setup and binding to resolve Ray issues

Author: XkunW

vec_inf/client/_slurm_script_generator.py

@@ -40,15 +40,29 @@ def __init__(self, params: dict[str, Any]):
         self.model_weights_path = str(
             Path(self.params["model_weights_parent_dir"], self.params["model_name"])
         )
+        self.env_str = self._generate_env_str()
+
+    def _generate_env_str(self) -> str:
+        """Generate the environment variables string for the Slurm script.
+        
+        Returns
+        -------
+        str
+            Formatted environment variables string for container or shell export commands.
+        """
         env_dict: dict[str, str] = self.params.get("env", {})
-        # Create string of environment variables
-        self.env_str = ""
-        for key, val in env_dict.items():
-            if len(self.env_str) == 0:
-                self.env_str = "--env "
-            else:
-                self.env_str += ","
-            self.env_str += key + "=" + val
+        
+        if not env_dict:
+            return ""
+        
+        if self.use_container:
+            # Format for container: --env KEY1=VAL1,KEY2=VAL2
+            env_pairs = [f"{key}={val}" for key, val in env_dict.items()]
+            return f"--env {','.join(env_pairs)}"
+        else:
+            # Format for shell: export KEY1=VAL1\nexport KEY2=VAL2
+            export_lines = [f"export {key}={val}" for key, val in env_dict.items()]
+            return "\n".join(export_lines)
 
     def _generate_script_content(self) -> str:
         """Generate the complete Slurm script content.
@@ -94,7 +108,12 @@ def _generate_server_setup(self) -> str:
         server_script = ["\n"]
         if self.use_container:
             server_script.append("\n".join(SLURM_SCRIPT_TEMPLATE["container_setup"]))
-        server_script.append("\n".join(SLURM_SCRIPT_TEMPLATE["env_vars"]))
+            server_script.append("\n".join(SLURM_SCRIPT_TEMPLATE["container_env_vars"]))
+        else:
+            server_script.append(
+                SLURM_SCRIPT_TEMPLATE["activate_venv"].format(venv=self.params["venv"])
+            )
+            server_script.append(self.env_str)
         server_script.append(
             SLURM_SCRIPT_TEMPLATE["imports"].format(src_dir=self.params["src_dir"])
         )
@@ -111,6 +130,11 @@ def _generate_server_setup(self) -> str:
                         env_str=self.env_str,
                     ),
                 )
+            else:
+                server_setup_str = server_setup_str.replace(
+                    "CONTAINER_PLACEHOLDER",
+                    "\\",
+                )
         else:
             server_setup_str = "\n".join(
                 SLURM_SCRIPT_TEMPLATE["server_setup"]["single_node"]
@@ -144,10 +168,7 @@ def _generate_launch_cmd(self) -> str:
                     env_str=self.env_str,
                 )
             )
-        else:
-            launcher_script.append(
-                SLURM_SCRIPT_TEMPLATE["activate_venv"].format(venv=self.params["venv"])
-            )
+
         launcher_script.append(
             "\n".join(SLURM_SCRIPT_TEMPLATE["launch_cmd"]).format(
                 model_weights_path=self.model_weights_path,

vec_inf/client/_slurm_templates.py

@@ -96,8 +96,8 @@ class SlurmScriptTemplate(TypedDict):
         f"{CONTAINER_MODULE_NAME} exec {IMAGE_PATH} ray stop",
     ],
     "imports": "source {src_dir}/find_port.sh",
-    "env_vars": [
-        f"export {CONTAINER_MODULE_NAME.upper()}_BINDPATH=${CONTAINER_MODULE_NAME.upper()}_BINDPATH,$(echo /dev/infiniband* | sed -e 's/ /,/g')"
+    "container_env_vars": [
+        f"export {CONTAINER_MODULE_NAME.upper()}_BINDPATH=${CONTAINER_MODULE_NAME.upper()}_BINDPATH,/dev,/tmp"
     ],
     "container_command": f"{CONTAINER_MODULE_NAME} exec --nv {{env_str}} --bind {{model_weights_path}}{{additional_binds}} --containall {IMAGE_PATH} \\",
     "activate_venv": "source {venv}/bin/activate",
@@ -112,6 +112,23 @@ class SlurmScriptTemplate(TypedDict):
             "nodes_array=($nodes)",
             "head_node=${{nodes_array[0]}}",
             'head_node_ip=$(srun --nodes=1 --ntasks=1 -w "$head_node" hostname --ip-address)',
+            "\n# Check for RDMA devices and set environment variable accordingly",
+            "if ! command -v ibv_devices >/dev/null 2>&1; then",
+            "   echo \"ibv_devices not found; forcing TCP. (No RDMA userland on host?)\"",
+            "   export NCCL_IB_DISABLE=1",
+            "   export NCCL_ENV_ARG=\"--env NCCL_IB_DISABLE=1\"",
+            "else",
+            "   # Pick GID index based on link layer (IB vs RoCE)",
+            "   if ibv_devinfo 2>/dev/null | grep -q \"link_layer:.*Ethernet\"; then",
+            "       # RoCEv2 typically needs a nonzero GID index; 3 is common, try 2 if your fabric uses it",
+            "       export NCCL_IB_GID_INDEX={{NCCL_IB_GID_INDEX:-3}}",
+            "       export NCCL_ENV_ARG=\"--env NCCL_IB_GID_INDEX={{NCCL_IB_GID_INDEX:-3}}\"",
+            "   else",
+            "       # Native InfiniBand => GID 0",
+            "       export NCCL_IB_GID_INDEX={{NCCL_IB_GID_INDEX:-0}}",
+            "       export NCCL_ENV_ARG=\"--env NCCL_IB_GID_INDEX={{NCCL_IB_GID_INDEX:-0}}\"",
+            "   fi",
+            "fi",
             "\n# Start Ray head node",
             "head_node_port=$(find_available_port $head_node_ip 8080 65535)",
             "ray_head=$head_node_ip:$head_node_port",