Add options for additional bindings, node exclusion and node list for model launch

Author: XkunW

vec_inf/client/_client_vars.py

@@ -65,6 +65,8 @@
     "qos": "qos",
     "time": "time",
     "nodes": "num_nodes",
+    "exclude": "exclude",
+    "nodelist": "node_list",
     "gpus-per-node": "gpus_per_node",
     "cpus-per-task": "cpus_per_task",
     "mem": "mem_per_node",
@@ -164,7 +166,7 @@ class SlurmScriptTemplate(TypedDict):
         f"export LD_LIBRARY_PATH={LD_LIBRARY_PATH}",
         f"export VLLM_NCCL_SO_PATH={VLLM_NCCL_SO_PATH}",
     ],
-    "singularity_command": f"singularity exec --nv --bind {{model_weights_path}}:{{model_weights_path}} --containall {SINGULARITY_IMAGE}",
+    "singularity_command": f"singularity exec --nv --bind {{model_weights_path}}:{{model_weights_path}}{{additional_binds}} --containall {SINGULARITY_IMAGE}",
     "activate_venv": "source {venv}/bin/activate",
     "server_setup": {
         "single_node": [

vec_inf/client/_slurm_script_generator.py

@@ -39,6 +39,9 @@ def __init__(self, params: dict[str, Any]):
         self.params = params
         self.is_multinode = int(self.params["num_nodes"]) > 1
         self.use_singularity = self.params["venv"] == "singularity"
+        self.additional_binds = self.params.get("binds", "")
+        if self.additional_binds:
+            self.additional_binds = f" --bind {self.additional_binds}"
         self.model_weights_path = str(
             Path(params["model_weights_parent_dir"], params["model_name"])
         )
@@ -104,6 +107,7 @@ def _generate_server_setup(self) -> str:
                     "SINGULARITY_PLACEHOLDER",
                     SLURM_SCRIPT_TEMPLATE["singularity_command"].format(
                         model_weights_path=self.model_weights_path,
+                        additional_binds=self.additional_binds,
                     ),
                 )
         else:
@@ -135,6 +139,7 @@ def _generate_launch_cmd(self) -> str:
             launcher_script.append(
                 SLURM_SCRIPT_TEMPLATE["singularity_command"].format(
                     model_weights_path=self.model_weights_path,
+                    additional_binds=self.additional_binds,
                 )
                 + " \\"
             )

vec_inf/client/config.py

@@ -108,6 +108,15 @@ class ModelConfig(BaseModel):
     partition: Union[PARTITION, str] = Field(
         default=cast(str, DEFAULT_ARGS["partition"]), description="GPU partition type"
     )
+    exclude: Optional[str] = Field(
+        default=None, description="Exclude certain nodes from the resources granted to the job"
+    )
+    node_list: Optional[str] = Field(
+        default=None, description="Request a specific list of nodes for deployment"
+    )
+    binds: Optional[str] = Field(
+        default=None, description="Additional binds for the singularity container"
+    )
     venv: str = Field(
         default="singularity", description="Virtual environment/container system"
     )