Move all bindings into BINDPATH env var

Author: XkunW

vec_inf/client/_slurm_script_generator.py

@@ -34,9 +34,7 @@ def __init__(self, params: dict[str, Any]):
         self.params = params
         self.is_multinode = int(self.params["num_nodes"]) > 1
         self.use_container = self.params["venv"] == CONTAINER_MODULE_NAME
-        self.additional_binds = self.params.get("bind", "")
-        if self.additional_binds:
-            self.additional_binds = f" --bind {self.additional_binds}"
+        self.additional_binds = f",{self.params['bind']}" if self.params.get("bind") else ""
         self.model_weights_path = str(
             Path(self.params["model_weights_parent_dir"], self.params["model_name"])
         )
@@ -107,7 +105,7 @@ 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["container_env_vars"]))
+            server_script.append(SLURM_SCRIPT_TEMPLATE["bind_path"].format(model_weights_path=self.model_weights_path, additional_binds=self.additional_binds))
         else:
             server_script.append(
                 SLURM_SCRIPT_TEMPLATE["activate_venv"].format(venv=self.params["venv"])
@@ -125,7 +123,6 @@ def _generate_server_setup(self) -> str:
                     "CONTAINER_PLACEHOLDER",
                     SLURM_SCRIPT_TEMPLATE["container_command"].format(
                         model_weights_path=self.model_weights_path,
-                        additional_binds=self.additional_binds,
                         env_str=self.env_str,
                     ),
                 )
@@ -163,7 +160,6 @@ def _generate_launch_cmd(self) -> str:
             launcher_script.append(
                 SLURM_SCRIPT_TEMPLATE["container_command"].format(
                     model_weights_path=self.model_weights_path,
-                    additional_binds=self.additional_binds,
                     env_str=self.env_str,
                 )
             )
@@ -215,11 +211,7 @@ def __init__(self, params: dict[str, Any]):
         self.script_paths: list[Path] = []
         self.use_container = self.params["venv"] == CONTAINER_MODULE_NAME
         for model_name in self.params["models"]:
-            self.params["models"][model_name]["additional_binds"] = ""
-            if self.params["models"][model_name].get("bind"):
-                self.params["models"][model_name]["additional_binds"] = (
-                    f" --bind {self.params['models'][model_name]['bind']}"
-                )
+            self.params["models"][model_name]["additional_binds"] = f",{self.params['bind']}" if self.params.get("bind") else ""
             self.params["models"][model_name]["model_weights_path"] = str(
                 Path(
                     self.params["models"][model_name]["model_weights_parent_dir"],
@@ -259,7 +251,10 @@ def _generate_model_launch_script(self, model_name: str) -> Path:
         script_content.append(BATCH_MODEL_LAUNCH_SCRIPT_TEMPLATE["shebang"])
         if self.use_container:
             script_content.append(BATCH_MODEL_LAUNCH_SCRIPT_TEMPLATE["container_setup"])
-        script_content.append("\n".join(BATCH_MODEL_LAUNCH_SCRIPT_TEMPLATE["env_vars"]))
+        script_content.append(BATCH_MODEL_LAUNCH_SCRIPT_TEMPLATE["bind_path"].format(
+            model_weights_path=model_params["model_weights_path"],
+            additional_binds=model_params["additional_binds"],
+        ))
         script_content.append(
             "\n".join(
                 BATCH_MODEL_LAUNCH_SCRIPT_TEMPLATE["server_address_setup"]
@@ -277,7 +272,6 @@ def _generate_model_launch_script(self, model_name: str) -> Path:
             script_content.append(
                 BATCH_MODEL_LAUNCH_SCRIPT_TEMPLATE["container_command"].format(
                     model_weights_path=model_params["model_weights_path"],
-                    additional_binds=model_params["additional_binds"],
                 )
             )
         script_content.append(

vec_inf/client/_slurm_templates.py

@@ -57,6 +57,8 @@ class SlurmScriptTemplate(TypedDict):
         Commands for container setup
     imports : str
         Import statements and source commands
+    bind_path : str
+        Bind path environment variable for the container
     container_command : str
         Template for container execution command
     activate_venv : str
@@ -74,7 +76,7 @@ class SlurmScriptTemplate(TypedDict):
     shebang: ShebangConfig
     container_setup: list[str]
     imports: str
-    container_env_vars: list[str]
+    bind_path: str
     container_command: str
     activate_venv: str
     server_setup: ServerSetupConfig
@@ -96,10 +98,8 @@ class SlurmScriptTemplate(TypedDict):
         f"{CONTAINER_MODULE_NAME} exec {IMAGE_PATH} ray stop",
     ],
     "imports": "source {src_dir}/find_port.sh",
-    "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} \\",
+    "bind_path": f"export {CONTAINER_MODULE_NAME.upper()}_BINDPATH=${CONTAINER_MODULE_NAME.upper()}_BINDPATH,/dev,/tmp,{{model_weights_path}}{{additional_binds}}",
+    "container_command": f"{CONTAINER_MODULE_NAME} exec --nv {{env_str}} --containall {IMAGE_PATH} \\",
     "activate_venv": "source {venv}/bin/activate",
     "server_setup": {
         "single_node": [
@@ -215,8 +215,8 @@ class BatchModelLaunchScriptTemplate(TypedDict):
         Shebang line for the script
     container_setup : list[str]
         Commands for container setup
-    env_vars : list[str]
-        Environment variables to set
+    bind_path : list[str]
+        Bind path environment variable for the container
     server_address_setup : list[str]
         Commands to setup the server address
     launch_cmd : list[str]
@@ -227,7 +227,7 @@ class BatchModelLaunchScriptTemplate(TypedDict):
 
     shebang: str
     container_setup: str
-    env_vars: list[str]
+    bind_path: list[str]
     server_address_setup: list[str]
     write_to_json: list[str]
     launch_cmd: list[str]
@@ -237,9 +237,7 @@ class BatchModelLaunchScriptTemplate(TypedDict):
 BATCH_MODEL_LAUNCH_SCRIPT_TEMPLATE: BatchModelLaunchScriptTemplate = {
     "shebang": "#!/bin/bash\n",
     "container_setup": f"{CONTAINER_LOAD_CMD}\n",
-    "env_vars": [
-        f"export {CONTAINER_MODULE_NAME}_BINDPATH=${CONTAINER_MODULE_NAME}_BINDPATH,$(echo /dev/infiniband* | sed -e 's/ /,/g')"
-    ],
+    "bind_path": f"export {CONTAINER_MODULE_NAME.upper()}_BINDPATH=${CONTAINER_MODULE_NAME.upper()}_BINDPATH,/dev,/tmp,{{model_weights_path}}{{additional_binds}}",
     "server_address_setup": [
         "source {src_dir}/find_port.sh",
         "head_node_ip=${{SLURMD_NODENAME}}",
@@ -255,7 +253,7 @@ class BatchModelLaunchScriptTemplate(TypedDict):
         '    "$json_path" > temp_{model_name}.json \\',
         '    && mv temp_{model_name}.json "$json_path"\n',
     ],
-    "container_command": f"{CONTAINER_MODULE_NAME} exec --nv --bind {{model_weights_path}}{{additional_binds}} --containall {IMAGE_PATH} \\",
+    "container_command": f"{CONTAINER_MODULE_NAME} exec --nv --containall {IMAGE_PATH} \\",
     "launch_cmd": [
         "vllm serve {model_weights_path} \\",
         "    --served-model-name {model_name} \\",