Merge pull request #160 from VectorInstitute/feature/misc-fixes

* Move all bindings into BINDPATH env var
* Fix required field checking for batch launch

Author: XkunW

tests/vec_inf/client/test_slurm_script_generator.py

@@ -99,7 +99,7 @@ def test_init_singularity(self, singularity_params):
         assert generator.params == singularity_params
         assert generator.use_container
         assert not generator.is_multinode
-        assert generator.additional_binds == " --bind /scratch:/scratch,/data:/data"
+        assert generator.additional_binds == ",/scratch:/scratch,/data:/data"
         assert generator.model_weights_path == "/path/to/model_weights/test-model"
         assert (
             generator.env_str
@@ -185,8 +185,6 @@ def test_generate_launch_cmd_singularity(self, singularity_params):
         launch_cmd = generator._generate_launch_cmd()
 
         assert "apptainer exec --nv" in launch_cmd
-        assert "--bind /path/to/model_weights/test-model" in launch_cmd
-        assert "--bind /scratch:/scratch,/data:/data" in launch_cmd
         assert "source" not in launch_cmd
 
     def test_generate_launch_cmd_boolean_args(self, basic_params):
@@ -327,14 +325,17 @@ def test_init_singularity(self, batch_singularity_params):
         """Test initialization with Singularity configuration."""
         generator = BatchSlurmScriptGenerator(batch_singularity_params)
 
+        print(generator.params["models"]["model1"]["additional_binds"])
+        print(generator.params["models"]["model2"]["additional_binds"])
+
         assert generator.use_container
         assert (
             generator.params["models"]["model1"]["additional_binds"]
-            == " --bind /scratch:/scratch,/data:/data"
+            == ",/scratch:/scratch,/data:/data"
         )
         assert (
             generator.params["models"]["model2"]["additional_binds"]
-            == " --bind /scratch:/scratch,/data:/data"
+            == ",/scratch:/scratch,/data:/data"
         )
 
     def test_init_singularity_no_bind(self, batch_params):

vec_inf/client/_helper.py

@@ -469,16 +469,15 @@ def _get_launch_params(
             If required fields are missing or tensor parallel size is not specified
             when using multiple GPUs
         """
-        params: dict[str, Any] = {
-            "models": {},
+        common_params: dict[str, Any] = {
             "slurm_job_name": self.slurm_job_name,
             "src_dir": str(SRC_DIR),
             "account": account,
             "work_dir": work_dir,
         }
 
-        # Check for required fields without default vals, will raise an error if missing
-        utils.check_required_fields(params)
+        params: dict[str, Any] = common_params.copy()
+        params["models"] = {}
 
         for i, (model_name, config) in enumerate(self.model_configs.items()):
             params["models"][model_name] = config.model_dump(exclude_none=True)
@@ -555,6 +554,10 @@ def _get_launch_params(
                     raise ValueError(
                         f"Mismatch found for {arg}: {params[arg]} != {params['models'][model_name][arg]}, check your configuration"
                     )
+            # Check for required fields, will raise an error if missing any
+            utils.check_required_fields(
+                {**params["models"][model_name], **common_params}
+            )
 
         return params
 

vec_inf/client/_slurm_script_generator.py

@@ -34,9 +34,9 @@ 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 +107,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["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 +130,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 +167,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 +218,11 @@ 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['models'][model_name]['bind']}"
+                if self.params["models"][model_name].get("bind")
+                else ""
+            )
             self.params["models"][model_name]["model_weights_path"] = str(
                 Path(
                     self.params["models"][model_name]["model_weights_parent_dir"],
@@ -259,7 +262,12 @@ 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 +285,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 : 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: 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} \\",