Add --account slurm argument for launch

Author: XkunW

vec_inf/cli/_cli.py

@@ -62,6 +62,11 @@ def cli() -> None:
     type=int,
     help="Number of GPUs/node to use, default to suggested resource allocation for model",
 )
+@click.option(
+    "--account",
+    type=str,
+    help="Charge resources used by this job to specified account.",
+)
 @click.option(
     "--qos",
     type=str,

vec_inf/client/_client_vars.py

@@ -56,6 +56,7 @@
 SLURM_JOB_CONFIG_ARGS = {
     "job-name": "model_name",
     "partition": "partition",
+    "account": "account",
     "qos": "qos",
     "time": "time",
     "nodes": "num_nodes",

vec_inf/client/_helper.py

@@ -175,7 +175,7 @@ def _get_launch_params(self) -> dict[str, Any]:
             If required fields are missing or tensor parallel size is not specified
             when using multiple GPUs
         """
-        params = self.model_config.model_dump()
+        params = self.model_config.model_dump(exclude_none=True)
 
         # Override config defaults with CLI arguments
         if self.kwargs.get("vllm_args"):

vec_inf/client/_slurm_script_generator.py

@@ -68,7 +68,8 @@ def _generate_shebang(self) -> str:
         """
         shebang = [SLURM_SCRIPT_TEMPLATE["shebang"]["base"]]
         for arg, value in SLURM_JOB_CONFIG_ARGS.items():
-            shebang.append(f"#SBATCH --{arg}={self.params[value]}")
+            if self.params.get(value):
+                shebang.append(f"#SBATCH --{arg}={self.params[value]}")
         if self.is_multinode:
             shebang += SLURM_SCRIPT_TEMPLATE["shebang"]["multinode"]
         return "\n".join(shebang)

vec_inf/client/config.py

@@ -47,6 +47,8 @@ class ModelConfig(BaseModel):
         Memory allocation per node in GB format (e.g., '32G')
     vocab_size : int
         Size of the model's vocabulary (1-1,000,000)
+    account : Optional[str], optional
+        Charge resources used by this job to specified account.
     qos : Union[QOS, str], optional
         Quality of Service tier for job scheduling
     time : str, optional
@@ -92,6 +94,9 @@ class ModelConfig(BaseModel):
         description="Memory per node",
     )
     vocab_size: int = Field(..., gt=0, le=1_000_000)
+    account: Optional[str] = Field(
+        default=None, description="Account name for job scheduling"
+    )
     qos: Union[QOS, str] = Field(
         default=cast(str, DEFAULT_ARGS["qos"]), description="Quality of Service tier"
     )

vec_inf/client/models.py

@@ -164,6 +164,8 @@ class LaunchOptions:
         Number of nodes to allocate
     gpus_per_node : int, optional
         Number of GPUs per node
+    account : str, optional
+        Account name for job scheduling
     qos : str, optional
         Quality of Service level
     time : str, optional
@@ -187,6 +189,7 @@ class LaunchOptions:
     partition: Optional[str] = None
     num_nodes: Optional[int] = None
     gpus_per_node: Optional[int] = None
+    account: Optional[str] = None
     qos: Optional[str] = None
     time: Optional[str] = None
     vocab_size: Optional[int] = None