ruff and mypy fixes

Author: XkunW

vec_inf/client/_client_vars.py

@@ -129,6 +129,8 @@ class SlurmScriptTemplate(TypedDict):
         Commands for Singularity container setup
     imports : str
         Import statements and source commands
+    env_vars : list[str]
+        Environment variables to set
     singularity_command : str
         Template for Singularity execution command
     activate_venv : str
@@ -146,6 +148,7 @@ class SlurmScriptTemplate(TypedDict):
     shebang: ShebangConfig
     singularity_setup: list[str]
     imports: str
+    env_vars: list[str]
     singularity_command: str
     activate_venv: str
     server_setup: ServerSetupConfig

vec_inf/client/_helper.py

@@ -5,7 +5,6 @@
 """
 
 import json
-import os
 import time
 import warnings
 from pathlib import Path

vec_inf/client/_slurm_script_generator.py

@@ -89,12 +89,8 @@ def _generate_server_setup(self) -> str:
         """
         server_script = ["\n"]
         if self.use_singularity:
-            server_script.append(
-                "\n".join(SLURM_SCRIPT_TEMPLATE["singularity_setup"])
-            )
-        server_script.append(
-            "\n".join(SLURM_SCRIPT_TEMPLATE["env_vars"])
-        )
+            server_script.append("\n".join(SLURM_SCRIPT_TEMPLATE["singularity_setup"]))
+        server_script.append("\n".join(SLURM_SCRIPT_TEMPLATE["env_vars"]))
         server_script.append(
             SLURM_SCRIPT_TEMPLATE["imports"].format(src_dir=self.params["src_dir"])
         )

vec_inf/client/config.py

@@ -109,7 +109,8 @@ class ModelConfig(BaseModel):
         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"
+        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"