[pre-commit.ci] Add auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

tests/vec_inf/client/test_api.py

@@ -8,7 +8,7 @@
 from vec_inf.client._exceptions import (
     ModelConfigurationError,
     ServerError,
-    SlurmJobError
+    SlurmJobError,
 )
 
 
@@ -324,8 +324,7 @@ def test_batch_launch_models_with_config():
     client.batch_launch_models = lambda model_names, batch_config=None: mock_response
 
     result = client.batch_launch_models(
-        ["model1", "model2"],
-        batch_config="custom_config.yaml"
+        ["model1", "model2"], batch_config="custom_config.yaml"
     )
 
     assert result.slurm_job_id == "12345678"
@@ -420,11 +419,16 @@ def test_batch_launch_models_configuration_error():
 
     # Mock the batch launch method to raise a configuration error
     def mock_batch_launch(model_names, batch_config=None):
-        raise ModelConfigurationError("Model 'nonexistent-model' not found in configuration")
+        raise ModelConfigurationError(
+            "Model 'nonexistent-model' not found in configuration"
+        )
 
     client.batch_launch_models = mock_batch_launch
 
-    with pytest.raises(ModelConfigurationError, match="Model 'nonexistent-model' not found in configuration"):
+    with pytest.raises(
+        ModelConfigurationError,
+        match="Model 'nonexistent-model' not found in configuration",
+    ):
         client.batch_launch_models(["model1", "nonexistent-model"])
 
 
@@ -438,7 +442,9 @@ def mock_batch_launch(model_names, batch_config=None):
 
     client.batch_launch_models = mock_batch_launch
 
-    with pytest.raises(SlurmJobError, match="sbatch: error: Invalid partition specified"):
+    with pytest.raises(
+        SlurmJobError, match="sbatch: error: Invalid partition specified"
+    ):
         client.batch_launch_models(["model1", "model2"])
 
 
@@ -448,15 +454,18 @@ def test_batch_launch_models_integration():
 
     with (
         patch("vec_inf.client.api.BatchModelLauncher") as mock_launcher_class,
-        patch("vec_inf.client.api.run_bash_command", return_value=("Submitted batch job 12345678", ""))
+        patch(
+            "vec_inf.client.api.run_bash_command",
+            return_value=("Submitted batch job 12345678", ""),
+        ),
     ):
         # Mock the BatchModelLauncher instance
         mock_launcher = MagicMock()
         mock_launcher.launch.return_value = MagicMock(
             slurm_job_id="12345678",
             slurm_job_name="BATCH-model1-model2",
             model_names=["model1", "model2"],
-            config={"slurm_job_id": "12345678"}
+            config={"slurm_job_id": "12345678"},
         )
         mock_launcher_class.return_value = mock_launcher
 
@@ -478,25 +487,29 @@ def test_batch_launch_models_with_custom_config_integration():
 
     with (
         patch("vec_inf.client.api.BatchModelLauncher") as mock_launcher_class,
-        patch("vec_inf.client.api.run_bash_command", return_value=("Submitted batch job 12345678", ""))
+        patch(
+            "vec_inf.client.api.run_bash_command",
+            return_value=("Submitted batch job 12345678", ""),
+        ),
     ):
         # Mock the BatchModelLauncher instance
         mock_launcher = MagicMock()
         mock_launcher.launch.return_value = MagicMock(
             slurm_job_id="12345678",
             slurm_job_name="BATCH-model1-model2",
             model_names=["model1", "model2"],
-            config={"slurm_job_id": "12345678"}
+            config={"slurm_job_id": "12345678"},
         )
         mock_launcher_class.return_value = mock_launcher
 
         result = client.batch_launch_models(
-            ["model1", "model2"],
-            batch_config="custom_config.yaml"
+            ["model1", "model2"], batch_config="custom_config.yaml"
         )
 
         # Verify BatchModelLauncher was called with custom config
-        mock_launcher_class.assert_called_once_with(["model1", "model2"], "custom_config.yaml")
+        mock_launcher_class.assert_called_once_with(
+            ["model1", "model2"], "custom_config.yaml"
+        )
         mock_launcher.launch.assert_called_once()
 
         # Verify the response

tests/vec_inf/client/test_helper.py

@@ -7,26 +7,26 @@
 import pytest
 import requests
 
+from vec_inf.client._client_vars import SRC_DIR
 from vec_inf.client._exceptions import (
     MissingRequiredFieldsError,
     ModelConfigurationError,
     ModelNotFoundError,
     SlurmJobError,
 )
 from vec_inf.client._helper import (
+    BatchModelLauncher,
     ModelLauncher,
     ModelRegistry,
     ModelStatusMonitor,
     PerformanceMetricsCollector,
-    BatchModelLauncher,
 )
 from vec_inf.client.config import ModelConfig
 from vec_inf.client.models import (
     ModelStatus,
     ModelType,
     StatusResponse,
 )
-from vec_inf.client._client_vars import SRC_DIR
 
 
 class TestModelLauncher:
@@ -197,6 +197,7 @@ def test_launch_with_slurm_error(
         with pytest.raises(SlurmJobError):
             launcher.launch()
 
+
 class TestBatchModelLauncher:
     """Tests for the BatchModelLauncher class."""
 
@@ -252,7 +253,9 @@ def test_init_with_valid_configs(self, mock_load_config, batch_model_configs):
         assert "family2-variant1" in launcher.model_configs
 
     @patch("vec_inf.client._helper.utils.load_config")
-    def test_init_with_missing_model_config(self, mock_load_config, batch_model_configs):
+    def test_init_with_missing_model_config(
+        self, mock_load_config, batch_model_configs
+    ):
         """Test error is raised when one of the models is missing from config."""
         mock_load_config.return_value = batch_model_configs
 
@@ -266,9 +269,14 @@ def test_init_with_missing_model_config(self, mock_load_config, batch_model_conf
     def test_get_slurm_job_name(self, mock_load_config, batch_model_configs):
         """Test SLURM job name is constructed correctly from model names."""
         mock_load_config.return_value = batch_model_configs
-        launcher = BatchModelLauncher(["family1-variant1", "family2-variant1", "family1-variant2"])
+        launcher = BatchModelLauncher(
+            ["family1-variant1", "family2-variant1", "family1-variant2"]
+        )
 
-        assert launcher.slurm_job_name == "BATCH-family1-variant1-family2-variant1-family1-variant2"
+        assert (
+            launcher.slurm_job_name
+            == "BATCH-family1-variant1-family2-variant1-family1-variant2"
+        )
 
     @patch("vec_inf.client._helper.utils.load_config")
     @patch("pathlib.Path.mkdir")
@@ -278,14 +286,19 @@ def test_get_launch_params_creates_log_dirs(
         """Test launch parameters preparation creates log directories."""
         mock_load_config.return_value = batch_model_configs
 
-        launcher = BatchModelLauncher(["family1-variant1", "family2-variant1", "family1-variant2"])
+        launcher = BatchModelLauncher(
+            ["family1-variant1", "family2-variant1", "family1-variant2"]
+        )
         params = launcher.params
 
         assert "models" in params
         assert "family1-variant1" in params["models"]
         assert "family2-variant1" in params["models"]
         assert "family1-variant2" in params["models"]
-        assert params["slurm_job_name"] == "BATCH-family1-variant1-family2-variant1-family1-variant2"
+        assert (
+            params["slurm_job_name"]
+            == "BATCH-family1-variant1-family2-variant1-family1-variant2"
+        )
         assert params["src_dir"] == str(SRC_DIR)
 
         # Check that log directories are created
@@ -318,7 +331,7 @@ def test_get_launch_params_with_non_power_of_two_gpus(
             batch_model_configs[0].model_copy(
                 update={
                     "gpus_per_node": 3,
-                    "vllm_args": {"--tensor-parallel-size": "3"}
+                    "vllm_args": {"--tensor-parallel-size": "3"},
                 }
             ),
             batch_model_configs[1],
@@ -345,7 +358,7 @@ def test_get_launch_params_with_mismatched_batch_args(
                 update={
                     "gpus_per_node": 1,
                     "num_nodes": 2,  # This will cause the mismatch
-                    "vllm_args": {"--tensor-parallel-size": "1"}
+                    "vllm_args": {"--tensor-parallel-size": "1"},
                 }
             ),
         ]
@@ -450,13 +463,31 @@ def test_launch_params_log_file_paths(self, mock_load_config, batch_model_config
         params = launcher.params
 
         # Check individual model log files
-        assert "family1-variant1.%j.out" in params["models"]["family1-variant1"]["out_file"]
-        assert "family1-variant1.%j.err" in params["models"]["family1-variant1"]["err_file"]
-        assert "family1-variant1.$SLURM_JOB_ID.json" in params["models"]["family1-variant1"]["json_file"]
+        assert (
+            "family1-variant1.%j.out"
+            in params["models"]["family1-variant1"]["out_file"]
+        )
+        assert (
+            "family1-variant1.%j.err"
+            in params["models"]["family1-variant1"]["err_file"]
+        )
+        assert (
+            "family1-variant1.$SLURM_JOB_ID.json"
+            in params["models"]["family1-variant1"]["json_file"]
+        )
 
-        assert "family2-variant1.%j.out" in params["models"]["family2-variant1"]["out_file"]
-        assert "family2-variant1.%j.err" in params["models"]["family2-variant1"]["err_file"]
-        assert "family2-variant1.$SLURM_JOB_ID.json" in params["models"]["family2-variant1"]["json_file"]
+        assert (
+            "family2-variant1.%j.out"
+            in params["models"]["family2-variant1"]["out_file"]
+        )
+        assert (
+            "family2-variant1.%j.err"
+            in params["models"]["family2-variant1"]["err_file"]
+        )
+        assert (
+            "family2-variant1.$SLURM_JOB_ID.json"
+            in params["models"]["family2-variant1"]["json_file"]
+        )
 
         # Check batch-level log files
         assert "BATCH-family1-variant1-family2-variant1.%j.out" in params["out_file"]
@@ -467,7 +498,9 @@ def test_init_with_batch_config(self, mock_load_config, batch_model_configs):
         """Test launcher initializes correctly with custom batch config."""
         mock_load_config.return_value = batch_model_configs
 
-        launcher = BatchModelLauncher(["family1-variant1", "family2-variant1"], batch_config="custom_config.yaml")
+        launcher = BatchModelLauncher(
+            ["family1-variant1", "family2-variant1"], batch_config="custom_config.yaml"
+        )
 
         assert launcher.batch_config == "custom_config.yaml"
         # Verify load_config was called with the custom config
@@ -855,4 +888,4 @@ def test_get_single_model_config_not_found(self, mock_load_config, mock_configs)
         registry = ModelRegistry()
 
         with pytest.raises(ModelNotFoundError):
-            registry.get_single_model_config("nonexistent_model")
+            registry.get_single_model_config("nonexistent_model")