Fix tests

Author: XkunW

tests/vec_inf/cli/test_cli.py

@@ -26,22 +26,23 @@ def test_launch_command_success(runner):
         # for Rich table rendering
         mock_response = MagicMock()
         mock_response.config = {
-            "slurm_job_id": "14933053",
-            "model_name": "Meta-Llama-3.1-8B",
-            "model_type": "LLM",
-            "log_dir": "/tmp/test_logs",
-            "partition": "gpu",
-            "qos": "normal",
-            "time": "1:00:00",
-            "num_nodes": "1",
-            "gpus_per_node": "1",
-            "cpus_per_task": "8",
-            "mem_per_node": "32G",
-            "model_weights_parent_dir": "/model-weights",
-            "vocab_size": "128000",
-            "vllm_args": {"max_model_len": 8192},
-            "env": {"CACHE": "/cache"},
-        }
+                "slurm_job_id": "14933053",
+                "model_name": "Meta-Llama-3.1-8B",
+                "model_type": "LLM",
+                "log_dir": "/tmp/test_logs",
+                "partition": "gpu",
+                "qos": "normal",
+                "time": "1:00:00",
+                "num_nodes": "1",
+                "gpus_per_node": "1",
+                "cpus_per_task": "8",
+                "mem_per_node": "32G",
+                "model_weights_parent_dir": "/model-weights",
+                "vocab_size": "128000",
+                "venv": "/path/to/venv",
+                "vllm_args": {"max_model_len": 8192},
+                "env": {"CACHE": "/cache"},
+            }
         mock_client.launch_model.return_value = mock_response
 
         result = runner.invoke(cli, ["launch", "Meta-Llama-3.1-8B"])

tests/vec_inf/cli/test_helper.py

@@ -35,6 +35,7 @@ def test_format_table_output(self):
             "mem_per_node": "32G",
             "model_weights_parent_dir": "/model-weights",
             "log_dir": "/tmp/logs",
+            "venv": "/path/to/venv",
             "vllm_args": {"max_model_len": 8192, "enable_prefix_caching": True},
             "env": {"CACHE": "/cache"},
         }
@@ -63,6 +64,7 @@ def test_format_table_output_with_minimal_params(self):
             "mem_per_node": "16G",
             "model_weights_parent_dir": "/weights",
             "log_dir": "/logs",
+            "venv": "/path/to/venv",
             "vllm_args": {},
             "env": {},
         }

tests/vec_inf/client/test_slurm_script_generator.py

@@ -53,7 +53,7 @@ def singularity_params(self, basic_params):
         singularity = basic_params.copy()
         singularity.update(
             {
-                "venv": "singularity",
+                "venv": "apptainer",
                 "bind": "/scratch:/scratch,/data:/data",
                 "env": {
                     "CACHE_DIR": "/cache",
@@ -109,7 +109,7 @@ def test_init_singularity(self, singularity_params):
     def test_init_singularity_no_bind(self, basic_params):
         """Test Singularity initialization without additional binds."""
         params = basic_params.copy()
-        params["venv"] = "singularity"
+        params["venv"] = "apptainer"
         generator = SlurmScriptGenerator(params)
 
         assert generator.params == params
@@ -185,7 +185,7 @@ def test_generate_launch_cmd_singularity(self, singularity_params):
         generator = SlurmScriptGenerator(singularity_params)
         launch_cmd = generator._generate_launch_cmd()
 
-        assert "exec --nv" in 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
@@ -306,9 +306,9 @@ def batch_params(self):
     def batch_singularity_params(self, batch_params):
         """Generate batch SLURM configuration parameters with Singularity."""
         singularity_params = batch_params.copy()
-        singularity_params["venv"] = "singularity"  # Set top-level venv to singularity
+        singularity_params["venv"] = "apptainer"  # Set top-level venv to apptainer
         for model_name in singularity_params["models"]:
-            singularity_params["models"][model_name]["venv"] = "singularity"
+            singularity_params["models"][model_name]["venv"] = "apptainer"
             singularity_params["models"][model_name]["bind"] = (
                 "/scratch:/scratch,/data:/data"
             )
@@ -341,9 +341,9 @@ def test_init_singularity(self, batch_singularity_params):
     def test_init_singularity_no_bind(self, batch_params):
         """Test Singularity initialization without additional binds."""
         params = batch_params.copy()
-        params["venv"] = "singularity"  # Set top-level venv to singularity
+        params["venv"] = "apptainer"  # Set top-level venv to apptainer
         for model_name in params["models"]:
-            params["models"][model_name]["venv"] = "singularity"
+            params["models"][model_name]["venv"] = "apptainer"
 
         generator = BatchSlurmScriptGenerator(params)
 

vec_inf/cli/_helper.py

@@ -76,7 +76,8 @@ def format_table_output(self) -> Table:
         table.add_row("Memory/Node", self.params["mem_per_node"])
 
         # Add job config details
-        table.add_row("Virtual Environment", self.params["venv"])
+        if self.params.get("venv"):
+            table.add_row("Virtual Environment", self.params["venv"])
         table.add_row(
             "Model Weights Directory",
             str(Path(self.params["model_weights_parent_dir"], self.model_name)),

vec_inf/client/_slurm_script_generator.py

@@ -193,9 +193,7 @@ class BatchSlurmScriptGenerator:
     def __init__(self, params: dict[str, Any]):
         self.params = params
         self.script_paths: list[Path] = []
-        self.use_container = (
-            self.params["venv"] == "singularity" or self.params["venv"] == "apptainer"
-        )
+        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"):