vec-inf 0.7.2__py3-none-any.whl → 0.8.0__py3-none-any.whl

vec_inf/client/_helper.py

@@ -17,9 +17,10 @@ import requests
 import vec_inf.client._utils as utils
 from vec_inf.client._client_vars import (
     BATCH_MODE_REQUIRED_MATCHING_ARGS,
+    ENGINE_SHORT_TO_LONG_MAP,
     KEY_METRICS,
     SRC_DIR,
-    VLLM_SHORT_TO_LONG_MAP,
+    SUPPORTED_ENGINES,
 )
 from vec_inf.client._exceptions import (
     MissingRequiredFieldsError,
@@ -63,6 +64,7 @@ class ModelLauncher:
         self.slurm_job_id = ""
         self.slurm_script_path = Path("")
         self.model_config = self._get_model_configuration(self.kwargs.get("config"))
+        self.engine = ""
         self.params = self._get_launch_params()
 
     def _warn(self, message: str) -> None:
@@ -137,32 +139,38 @@ class ModelLauncher:
             f"not found at expected path '{model_weights_path}'"
         )
 
-    def _process_vllm_args(self, arg_string: str) -> dict[str, Any]:
-        """Process the vllm_args string into a dictionary.
+    def _process_engine_args(
+        self, arg_string: str, engine_choice: str
+    ) -> dict[str, Any]:
+        """Process the engine_args string into a dictionary.
 
         Parameters
         ----------
         arg_string : str
-            Comma-separated string of vLLM arguments
+            Comma-separated string of inference engine arguments
 
         Returns
         -------
         dict[str, Any]
-            Processed vLLM arguments as key-value pairs
+            Processed inference engine arguments as key-value pairs
         """
-        vllm_args: dict[str, str | bool] = {}
+        engine_args: dict[str, str | bool] = {}
+        engine_arg_map = ENGINE_SHORT_TO_LONG_MAP[engine_choice]
+
         for arg in arg_string.split(","):
             if "=" in arg:
                 key, value = arg.split("=")
-                if key.strip() in VLLM_SHORT_TO_LONG_MAP:
-                    key = VLLM_SHORT_TO_LONG_MAP[key.strip()]
-                vllm_args[key.strip()] = value.strip()
+                if key.strip() in engine_arg_map:
+                    key = engine_arg_map[key.strip()]
+                engine_args[key.strip()] = value.strip()
             elif "-O" in arg.strip():
-                key = VLLM_SHORT_TO_LONG_MAP["-O"]
-                vllm_args[key] = arg.strip()[2:].strip()
+                if engine_choice != "vllm":
+                    raise ValueError("-O is only supported for vLLM")
+                key = engine_arg_map["-O"]
+                engine_args[key] = arg.strip()[2:].strip()
             else:
-                vllm_args[arg.strip()] = True
-        return vllm_args
+                engine_args[arg.strip()] = True
+        return engine_args
 
     def _process_env_vars(self, env_arg: str) -> dict[str, str]:
         """Process the env string into a dictionary of environment variables.
@@ -196,6 +204,63 @@ class ModelLauncher:
                         print(f"WARNING: Could not parse env var: {line}")
         return env_vars
 
+    def _engine_check_override(self, params: dict[str, Any]) -> None:
+        """Check for engine override in CLI args and warn user.
+
+        Parameters
+        ----------
+        params : dict[str, Any]
+            Dictionary of launch parameters to check
+        """
+
+        def overwrite_engine_args(params: dict[str, Any]) -> None:
+            engine_args = self._process_engine_args(
+                self.kwargs[f"{self.engine}_args"], self.engine
+            )
+            for key, value in engine_args.items():
+                params["engine_args"][key] = value
+            del self.kwargs[f"{self.engine}_args"]
+
+        # Infer engine name from engine-specific args if provided
+        extracted_engine = ""
+        for engine in SUPPORTED_ENGINES:
+            if self.kwargs.get(f"{engine}_args"):
+                if not extracted_engine:
+                    extracted_engine = engine
+                else:
+                    raise ValueError(
+                        "Cannot provide engine-specific args for multiple engines, please choose one"
+                    )
+        # Check for mismatch between provided engine arg and engine-specific args
+        input_engine = self.kwargs.get("engine", "")
+
+        if input_engine and extracted_engine:
+            if input_engine != extracted_engine:
+                raise ValueError(
+                    f"Mismatch between provided engine '{input_engine}' and engine-specific args '{extracted_engine}'"
+                )
+            self.engine = input_engine
+            params["engine_args"] = params[f"{self.engine}_args"]
+            overwrite_engine_args(params)
+        elif input_engine:
+            # Only engine arg in CLI, use default engine args from config
+            self.engine = input_engine
+            params["engine_args"] = params[f"{self.engine}_args"]
+        elif extracted_engine:
+            # Only engine-specific args in CLI, infer engine and warn user
+            self.engine = extracted_engine
+            params["engine_inferred"] = True
+            params["engine_args"] = params[f"{self.engine}_args"]
+            overwrite_engine_args(params)
+        else:
+            # No engine-related args in CLI, use defaults from config
+            self.engine = params.get("engine", "vllm")
+            params["engine_args"] = params[f"{self.engine}_args"]
+
+        # Remove $ENGINE_NAME_args from params as they won't get populated to sjob json.
+        for engine in SUPPORTED_ENGINES:
+            del params[f"{engine}_args"]
+
     def _apply_cli_overrides(self, params: dict[str, Any]) -> None:
         """Apply CLI argument overrides to params.
 
@@ -204,11 +269,7 @@ class ModelLauncher:
         params : dict[str, Any]
             Dictionary of launch parameters to override
         """
-        if self.kwargs.get("vllm_args"):
-            vllm_args = self._process_vllm_args(self.kwargs["vllm_args"])
-            for key, value in vllm_args.items():
-                params["vllm_args"][key] = value
-            del self.kwargs["vllm_args"]
+        self._engine_check_override(params)
 
         if self.kwargs.get("env"):
             env_vars = self._process_env_vars(self.kwargs["env"])
@@ -241,7 +302,7 @@ class ModelLauncher:
         """
         if (
             int(params["gpus_per_node"]) > 1
-            and params["vllm_args"].get("--tensor-parallel-size") is None
+            and params["engine_args"].get("--tensor-parallel-size") is None
         ):
             raise MissingRequiredFieldsError(
                 "--tensor-parallel-size is required when gpus_per_node > 1"
@@ -252,8 +313,8 @@ class ModelLauncher:
             raise ValueError("Total number of GPUs requested must be a power of two")
 
         total_parallel_sizes = int(
-            params["vllm_args"].get("--tensor-parallel-size", "1")
-        ) * int(params["vllm_args"].get("--pipeline-parallel-size", "1"))
+            params["engine_args"].get("--tensor-parallel-size", "1")
+        ) * int(params["engine_args"].get("--pipeline-parallel-size", "1"))
         if total_gpus_requested != total_parallel_sizes:
             raise ValueError(
                 "Mismatch between total number of GPUs requested and parallelization settings"
@@ -312,7 +373,8 @@ class ModelLauncher:
 
         # Convert path to string for JSON serialization
         for field in params:
-            if field in ["vllm_args", "env"]:
+            # Keep structured fields (dicts/bools) intact
+            if field in ["engine_args", "env", "engine_inferred"]:
                 continue
             params[field] = str(params[field])
 
@@ -370,7 +432,7 @@ class ModelLauncher:
 
         # Replace venv with image path if using container
         if self.params["venv"] == CONTAINER_MODULE_NAME:
-            self.params["venv"] = IMAGE_PATH
+            self.params["venv"] = IMAGE_PATH[self.params["engine"]]
 
         with job_json.open("w") as file:
             json.dump(self.params, file, indent=4)
@@ -453,6 +515,53 @@ class BatchModelLauncher:
 
         return model_configs_dict
 
+    def _validate_resource_and_parallel_settings(
+        self,
+        config: ModelConfig,
+        model_engine_args: dict[str, Any] | None,
+        model_name: str,
+    ) -> None:
+        """Validate resource allocation and parallelization settings for each model.
+
+        Parameters
+        ----------
+        config : ModelConfig
+            Configuration of the model to validate
+        model_engine_args : dict[str, Any] | None
+            Inference engine arguments of the model to validate
+        model_name : str
+            Name of the model to validate
+
+        Raises
+        ------
+        MissingRequiredFieldsError
+            If tensor parallel size is not specified when using multiple GPUs
+        ValueError
+            If total # of GPUs requested is not a power of two
+            If mismatch between total # of GPUs requested and parallelization settings
+        """
+        if (
+            int(config.gpus_per_node) > 1
+            and (model_engine_args or {}).get("--tensor-parallel-size") is None
+        ):
+            raise MissingRequiredFieldsError(
+                f"--tensor-parallel-size is required when gpus_per_node > 1, check your configuration for {model_name}"
+            )
+
+        total_gpus_requested = int(config.gpus_per_node) * int(config.num_nodes)
+        if not utils.is_power_of_two(total_gpus_requested):
+            raise ValueError(
+                f"Total number of GPUs requested must be a power of two, check your configuration for {model_name}"
+            )
+
+        total_parallel_sizes = int(
+            (model_engine_args or {}).get("--tensor-parallel-size", "1")
+        ) * int((model_engine_args or {}).get("--pipeline-parallel-size", "1"))
+        if total_gpus_requested != total_parallel_sizes:
+            raise ValueError(
+                f"Mismatch between total number of GPUs requested and parallelization settings, check your configuration for {model_name}"
+            )
+
     def _get_launch_params(
         self, account: Optional[str] = None, work_dir: Optional[str] = None
     ) -> dict[str, Any]:
@@ -469,43 +578,29 @@ class BatchModelLauncher:
             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)
             params["models"][model_name]["het_group_id"] = i
 
-            # Validate resource allocation and parallelization settings
-            if (
-                int(config.gpus_per_node) > 1
-                and (config.vllm_args or {}).get("--tensor-parallel-size") is None
-            ):
-                raise MissingRequiredFieldsError(
-                    f"--tensor-parallel-size is required when gpus_per_node > 1, check your configuration for {model_name}"
-                )
+            model_engine_args = getattr(config, f"{config.engine}_args", None)
+            params["models"][model_name]["engine_args"] = model_engine_args
+            for engine in SUPPORTED_ENGINES:
+                del params["models"][model_name][f"{engine}_args"]
 
-            total_gpus_requested = int(config.gpus_per_node) * int(config.num_nodes)
-            if not utils.is_power_of_two(total_gpus_requested):
-                raise ValueError(
-                    f"Total number of GPUs requested must be a power of two, check your configuration for {model_name}"
-                )
-
-            total_parallel_sizes = int(
-                (config.vllm_args or {}).get("--tensor-parallel-size", "1")
-            ) * int((config.vllm_args or {}).get("--pipeline-parallel-size", "1"))
-            if total_gpus_requested != total_parallel_sizes:
-                raise ValueError(
-                    f"Mismatch between total number of GPUs requested and parallelization settings, check your configuration for {model_name}"
-                )
+            # Validate resource allocation and parallelization settings
+            self._validate_resource_and_parallel_settings(
+                config, model_engine_args, model_name
+            )
 
             # Convert gpus_per_node and resource_type to gres
             params["models"][model_name]["gres"] = (
@@ -555,6 +650,16 @@ class BatchModelLauncher:
                     raise ValueError(
                         f"Mismatch found for {arg}: {params[arg]} != {params['models'][model_name][arg]}, check your configuration"
                     )
+            # Check for required fields and return environment variable overrides
+            env_overrides = utils.check_required_fields(
+                {**params["models"][model_name], **common_params}
+            )
+
+            for arg, value in env_overrides.items():
+                if arg in common_params:
+                    params[arg] = value
+                else:
+                    params["models"][model_name][arg] = value
 
         return params
 
@@ -718,7 +823,7 @@ class ModelStatusMonitor:
             Basic status information for the job
         """
         try:
-            job_name = self.job_status["JobName"]
+            job_name = self.job_status["JobName"].removesuffix("-vec-inf")
             job_state = self.job_status["JobState"]
         except KeyError:
             job_name = "UNAVAILABLE"

vec_inf/client/_slurm_script_generator.py

@@ -1,7 +1,7 @@
-"""Class for generating Slurm scripts to run vLLM servers.
+"""Class for generating Slurm scripts to run inference servers.
 
-This module provides functionality to generate Slurm scripts for running vLLM servers
-in both single-node and multi-node configurations.
+This module provides functionality to generate Slurm scripts for running inference
+servers in both single-node and multi-node configurations.
 """
 
 from datetime import datetime
@@ -14,11 +14,11 @@ from vec_inf.client._slurm_templates import (
     BATCH_SLURM_SCRIPT_TEMPLATE,
     SLURM_SCRIPT_TEMPLATE,
 )
-from vec_inf.client._slurm_vars import CONTAINER_MODULE_NAME
+from vec_inf.client._slurm_vars import CONTAINER_MODULE_NAME, IMAGE_PATH
 
 
 class SlurmScriptGenerator:
-    """A class to generate Slurm scripts for running vLLM servers.
+    """A class to generate Slurm scripts for running inference servers.
 
     This class handles the generation of Slurm scripts for both single-node and
     multi-node configurations, supporting different virtualization environments
@@ -32,11 +32,12 @@ class SlurmScriptGenerator:
 
     def __init__(self, params: dict[str, Any]):
         self.params = params
+        self.engine = params.get("engine", "vllm")
         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"])
         )
@@ -89,6 +90,8 @@ class SlurmScriptGenerator:
         for arg, value in SLURM_JOB_CONFIG_ARGS.items():
             if self.params.get(value):
                 shebang.append(f"#SBATCH --{arg}={self.params[value]}")
+            if value == "model_name":
+                shebang[-1] += "-vec-inf"
         if self.is_multinode:
             shebang += SLURM_SCRIPT_TEMPLATE["shebang"]["multinode"]
         return "\n".join(shebang)
@@ -107,7 +110,12 @@ class SlurmScriptGenerator:
         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"])
@@ -116,17 +124,17 @@ class SlurmScriptGenerator:
         server_script.append(
             SLURM_SCRIPT_TEMPLATE["imports"].format(src_dir=self.params["src_dir"])
         )
-        if self.is_multinode:
+
+        if self.is_multinode and self.engine == "vllm":
             server_setup_str = "\n".join(
-                SLURM_SCRIPT_TEMPLATE["server_setup"]["multinode"]
+                SLURM_SCRIPT_TEMPLATE["server_setup"]["multinode_vllm"]
             ).format(gpus_per_node=self.params["gpus_per_node"])
             if self.use_container:
                 server_setup_str = server_setup_str.replace(
                     "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,
+                        image_path=IMAGE_PATH[self.engine],
                     ),
                 )
             else:
@@ -134,12 +142,16 @@ class SlurmScriptGenerator:
                     "CONTAINER_PLACEHOLDER",
                     "\\",
                 )
+        elif self.is_multinode and self.engine == "sglang":
+            server_setup_str = "\n".join(
+                SLURM_SCRIPT_TEMPLATE["server_setup"]["multinode_sglang"]
+            )
         else:
             server_setup_str = "\n".join(
                 SLURM_SCRIPT_TEMPLATE["server_setup"]["single_node"]
             )
         server_script.append(server_setup_str)
-        server_script.append("\n".join(SLURM_SCRIPT_TEMPLATE["find_vllm_port"]))
+        server_script.append("\n".join(SLURM_SCRIPT_TEMPLATE["find_server_port"]))
         server_script.append(
             "\n".join(SLURM_SCRIPT_TEMPLATE["write_to_json"]).format(
                 log_dir=self.params["log_dir"], model_name=self.params["model_name"]
@@ -148,39 +160,85 @@ class SlurmScriptGenerator:
         return "\n".join(server_script)
 
     def _generate_launch_cmd(self) -> str:
-        """Generate the vLLM server launch command.
+        """Generate the inference server launch command.
 
-        Creates the command to launch the vLLM server, handling different virtualization
-        environments (venv or singularity/apptainer).
+        Creates the command to launch the inference server, handling different
+        virtualization environments (venv or singularity/apptainer).
 
         Returns
         -------
         str
             Server launch command.
         """
-        launcher_script = ["\n"]
+        if self.is_multinode and self.engine == "sglang":
+            return self._generate_multinode_sglang_launch_cmd()
+
+        launch_cmd = ["\n"]
         if self.use_container:
-            launcher_script.append(
+            launch_cmd.append(
                 SLURM_SCRIPT_TEMPLATE["container_command"].format(
-                    model_weights_path=self.model_weights_path,
-                    additional_binds=self.additional_binds,
                     env_str=self.env_str,
+                    image_path=IMAGE_PATH[self.engine],
                 )
             )
 
-        launcher_script.append(
-            "\n".join(SLURM_SCRIPT_TEMPLATE["launch_cmd"]).format(
+        launch_cmd.append(
+            "\n".join(SLURM_SCRIPT_TEMPLATE["launch_cmd"][self.engine]).format(  # type: ignore[literal-required]
                 model_weights_path=self.model_weights_path,
                 model_name=self.params["model_name"],
             )
         )
 
-        for arg, value in self.params["vllm_args"].items():
+        for arg, value in self.params["engine_args"].items():
             if isinstance(value, bool):
-                launcher_script.append(f"    {arg} \\")
+                launch_cmd.append(f"    {arg} \\")
             else:
-                launcher_script.append(f"    {arg} {value} \\")
-        return "\n".join(launcher_script)
+                launch_cmd.append(f"    {arg} {value} \\")
+
+        # A known bug in vLLM requires setting backend to ray for multi-node
+        # Remove this when the bug is fixed
+        if self.is_multinode:
+            launch_cmd.append("    --distributed-executor-backend ray \\")
+
+        return "\n".join(launch_cmd).rstrip(" \\")
+
+    def _generate_multinode_sglang_launch_cmd(self) -> str:
+        """Generate the launch command for multi-node sglang setup.
+
+        Returns
+        -------
+        str
+            Multi-node sglang launch command.
+        """
+        launch_cmd = "\n" + "\n".join(
+            SLURM_SCRIPT_TEMPLATE["launch_cmd"]["sglang_multinode"]
+        ).format(
+            num_nodes=self.params["num_nodes"],
+            model_weights_path=self.model_weights_path,
+            model_name=self.params["model_name"],
+        )
+
+        container_placeholder = "\\"
+        if self.use_container:
+            container_placeholder = SLURM_SCRIPT_TEMPLATE["container_command"].format(
+                env_str=self.env_str,
+                image_path=IMAGE_PATH[self.engine],
+            )
+        launch_cmd = launch_cmd.replace(
+            "CONTAINER_PLACEHOLDER",
+            container_placeholder,
+        )
+
+        engine_arg_str = ""
+        for arg, value in self.params["engine_args"].items():
+            if isinstance(value, bool):
+                engine_arg_str += f"            {arg} \\\n"
+            else:
+                engine_arg_str += f"            {arg} {value} \\\n"
+
+        return launch_cmd.replace(
+            "SGLANG_ARGS_PLACEHOLDER", engine_arg_str.rstrip("\\\n")
+        )
 
     def write_to_log_dir(self) -> Path:
         """Write the generated Slurm script to the log directory.
@@ -207,7 +265,7 @@ class BatchSlurmScriptGenerator:
     """A class to generate Slurm scripts for batch mode.
 
     This class handles the generation of Slurm scripts for batch mode, which
-    launches multiple vLLM servers with different configurations in parallel.
+    launches multiple inference servers with different configurations in parallel.
     """
 
     def __init__(self, params: dict[str, Any]):
@@ -215,11 +273,11 @@ class BatchSlurmScriptGenerator:
         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"],
@@ -241,7 +299,7 @@ class BatchSlurmScriptGenerator:
         return script_path
 
     def _generate_model_launch_script(self, model_name: str) -> Path:
-        """Generate the bash script for launching individual vLLM servers.
+        """Generate the bash script for launching individual inference servers.
 
         Parameters
         ----------
@@ -251,7 +309,7 @@ class BatchSlurmScriptGenerator:
         Returns
         -------
         Path
-            The bash script path for launching the vLLM server.
+            The bash script path for launching the inference server.
         """
         # Generate the bash script content
         script_content = []
@@ -259,7 +317,12 @@ class BatchSlurmScriptGenerator:
         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"]
@@ -276,22 +339,23 @@ class BatchSlurmScriptGenerator:
         if self.use_container:
             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"],
+                    image_path=IMAGE_PATH[model_params["engine"]],
                 )
             )
         script_content.append(
-            "\n".join(BATCH_MODEL_LAUNCH_SCRIPT_TEMPLATE["launch_cmd"]).format(
+            "\n".join(
+                BATCH_MODEL_LAUNCH_SCRIPT_TEMPLATE["launch_cmd"][model_params["engine"]]
+            ).format(
                 model_weights_path=model_params["model_weights_path"],
                 model_name=model_name,
             )
         )
-        for arg, value in model_params["vllm_args"].items():
+        for arg, value in model_params["engine_args"].items():
             if isinstance(value, bool):
                 script_content.append(f"    {arg} \\")
             else:
                 script_content.append(f"    {arg} {value} \\")
-        script_content[-1] = script_content[-1].replace("\\", "")
+        script_content[-1] = script_content[-1].rstrip(" \\")
         # Write the bash script to the log directory
         launch_script_path = self._write_to_log_dir(
             script_content, f"launch_{model_name}.sh"
@@ -321,6 +385,8 @@ class BatchSlurmScriptGenerator:
                 model_params = self.params["models"][model_name]
                 if model_params.get(value) and value not in ["out_file", "err_file"]:
                     shebang.append(f"#SBATCH --{arg}={model_params[value]}")
+                if value == "model_name":
+                    shebang[-1] += "-vec-inf"
             shebang[-1] += "\n"
             shebang.append(BATCH_SLURM_SCRIPT_TEMPLATE["hetjob"])
         # Remove the last hetjob line
@@ -328,12 +394,12 @@ class BatchSlurmScriptGenerator:
         return "\n".join(shebang)
 
     def generate_batch_slurm_script(self) -> Path:
-        """Generate the Slurm script for launching multiple vLLM servers in batch mode.
+        """Generate the Slurm script for launching multiple inference servers in batch.
 
         Returns
         -------
         Path
-            The Slurm script for launching multiple vLLM servers in batch mode.
+            The Slurm script for launching multiple inference servers in batch.
         """
         script_content = []