Rebase onto develop and re-run ruff formatter

Author: tgasser-nv

aiperf/aiperf_models.py

@@ -33,22 +33,16 @@ class BaseConfig(BaseModel):
         description="Optional tokenizer Huggingface name, or local directory",
     )
     url: str = Field(..., description="Model base URL")
-    endpoint: str = Field(
-        default="/v1/chat/completions", description="API endpoint path"
-    )
+    endpoint: str = Field(default="/v1/chat/completions", description="API endpoint path")
     endpoint_type: Literal["chat", "completions"] = Field(
         default="chat",
         description="Type of endpoint (chat or completions)",
     )
-    api_key_env_var: Optional[str] = Field(
-        default=None, description="API key environment variable"
-    )
+    api_key_env_var: Optional[str] = Field(default=None, description="API key environment variable")
     streaming: Optional[bool] = Field(default=False, description="Streaming mode")
 
     # Load generation settings
-    warmup_request_count: int = Field(
-        description="Requests to send before beginning performance-test"
-    )
+    warmup_request_count: int = Field(description="Requests to send before beginning performance-test")
     benchmark_duration: int = Field(description="Benchmark duration in seconds")
     concurrency: int = Field(description="Number of concurrent requests")
     request_rate: Optional[float] = Field(
@@ -61,9 +55,7 @@ class BaseConfig(BaseModel):
     )
 
     # Synthetic data generation
-    random_seed: Optional[int] = Field(
-        default=None, description="Random seed for reproducibility"
-    )
+    random_seed: Optional[int] = Field(default=None, description="Random seed for reproducibility")
     prompt_input_tokens_mean: Optional[int] = Field(
         default=None,
         description="Mean number of input tokens",
@@ -85,26 +77,20 @@ class BaseConfig(BaseModel):
 class AIPerfConfig(BaseModel):
     """Main configuration model for AIPerf benchmark runner."""
 
-    batch_name: str = Field(
-        default="benchmark", description="Name for this batch of benchmarks"
-    )
+    batch_name: str = Field(default="benchmark", description="Name for this batch of benchmarks")
     output_base_dir: str = Field(
         default="aiperf_results",
         description="Base directory for benchmark results",
     )
-    base_config: BaseConfig = Field(
-        ..., description="Base configuration applied to all benchmark runs"
-    )
+    base_config: BaseConfig = Field(..., description="Base configuration applied to all benchmark runs")
     sweeps: Optional[Dict[str, List[Union[int, str]]]] = Field(
         default=None,
         description="Parameter sweeps. Key is the parameter to change, value is a list of values to use",
     )
 
     @field_validator("sweeps")
     @classmethod
-    def validate_sweeps(
-        cls, v: Optional[Dict[str, List[Any]]]
-    ) -> Optional[Dict[str, List[Any]]]:
+    def validate_sweeps(cls, v: Optional[Dict[str, List[Any]]]) -> Optional[Dict[str, List[Any]]]:
         """Validate that sweep values are lists of ints or strings."""
         if v is None:
             return v

aiperf/run_aiperf.py

@@ -25,7 +25,7 @@
 from datetime import datetime
 from pathlib import Path
 from subprocess import CompletedProcess
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Union
 
 import httpx
 import typer
@@ -37,9 +37,7 @@
 log = logging.getLogger(__name__)
 log.setLevel(logging.INFO)
 
-formatter = logging.Formatter(
-    "%(asctime)s %(levelname)s: %(message)s", datefmt="%Y-%m-%d %H:%M:%S"
-)
+formatter = logging.Formatter("%(asctime)s %(levelname)s: %(message)s", datefmt="%Y-%m-%d %H:%M:%S")
 console_handler = logging.StreamHandler()
 console_handler.setLevel(logging.DEBUG)
 console_handler.setFormatter(formatter)
@@ -143,9 +141,7 @@ def _sanitize_command_for_logging(cmd: List[str]) -> str:
 
         return " ".join(sanitized)
 
-    def _build_command(
-        self, sweep_params: Optional[Dict[str, Union[str, int]]], output_dir: Path
-    ) -> List[str]:
+    def _build_command(self, sweep_params: Optional[Dict[str, Union[str, int]]], output_dir: Path) -> List[str]:
         """Create a list of strings with the aiperf command and arguments to execute"""
 
         # Run aiperf in profile mode: `aiperf profile`
@@ -239,9 +235,7 @@ def _save_run_metadata(
             json.dump(metadata, f, indent=2)
 
     @staticmethod
-    def _save_subprocess_result_json(
-        output_dir: Path, result: CompletedProcess
-    ) -> None:
+    def _save_subprocess_result_json(output_dir: Path, result: CompletedProcess) -> None:
         """Save the subprocess result to the given filename"""
 
         process_result_file = output_dir / "process_result.json"
@@ -252,15 +246,11 @@ def _save_subprocess_result_json(
                 json.dump(save_data, f, indent=2)
 
         except (IOError, OSError) as e:
-            log.error(
-                "Could not write %s to file %s: %s", save_data, process_result_file, e
-            )
+            log.error("Could not write %s to file %s: %s", save_data, process_result_file, e)
             raise
 
         except TypeError as e:
-            log.error(
-                "Couldn't serialize %s to %s: %s", save_data, process_result_file, e
-            )
+            log.error("Couldn't serialize %s to %s: %s", save_data, process_result_file, e)
             raise
 
     def _check_service(self, endpoint: Optional[str] = "/v1/models") -> None:
@@ -357,9 +347,7 @@ def run_single_benchmark(
             log.info("Run completed successfully")
             self._save_subprocess_result_json(run_output_dir, result)
             run_completed = 1 if result.returncode == 0 else 0
-            return AIPerfSummary(
-                total=1, completed=run_completed, failed=1 - run_completed
-            )
+            return AIPerfSummary(total=1, completed=run_completed, failed=1 - run_completed)
 
         except subprocess.CalledProcessError as e:
             log.error("Run failed with exit code %s", e.returncode)
@@ -379,9 +367,7 @@ def run_batch_benchmarks(
         # Generate all sweep combinations
         combinations = self._get_sweep_combinations()
         if not combinations:
-            raise RuntimeError(
-                f"Can't generate sweep combinations from {self.config.sweeps}"
-            )
+            raise RuntimeError(f"Can't generate sweep combinations from {self.config.sweeps}")
 
         num_combinations = len(combinations)
         log.info("Running %s benchmarks", num_combinations)

nemoguardrails/cli/__init__.py

@@ -25,7 +25,6 @@
 
 from nemoguardrails import __version__
 from nemoguardrails.actions_server import actions_server
-from nemoguardrails.benchmark.aiperf.run_aiperf import app as aiperf_app
 from nemoguardrails.cli.chat import run_chat
 from nemoguardrails.cli.migration import migrate
 from nemoguardrails.cli.providers import _list_providers, select_provider_with_type

nemoguardrails/server/api.py

@@ -475,9 +475,7 @@ async def chat_completion(body: RequestBody, request: Request):
 
     except Exception as ex:
         log.exception(ex)
-        return ResponseBody(
-            messages=[{"role": "assistant", "content": "Internal server error."}]
-        )
+        return ResponseBody(messages=[{"role": "assistant", "content": "Internal server error."}])
 
 
 # By default, there are no challenges

tests/benchmark/test_aiperf_models.py

@@ -236,11 +236,7 @@ def test_aiperf_config_sweep_invalid_value_type_dict(self, valid_base_config):
         error_msg = str(exc_info.value)
         # Pydantic catches this during type validation
         assert "sweeps.concurrency" in error_msg
-        assert (
-            "must be int or str" in error_msg
-            or "int_type" in error_msg
-            or "string_type" in error_msg
-        )
+        assert "must be int or str" in error_msg or "int_type" in error_msg or "string_type" in error_msg
 
     def test_aiperf_config_sweep_invalid_value_type_list(self, valid_base_config):
         """Test that list values in sweeps raise validation error."""
@@ -254,11 +250,7 @@ def test_aiperf_config_sweep_invalid_value_type_list(self, valid_base_config):
         error_msg = str(exc_info.value)
         # Pydantic catches this during type validation
         assert "sweeps.concurrency" in error_msg
-        assert (
-            "must be int or str" in error_msg
-            or "int_type" in error_msg
-            or "string_type" in error_msg
-        )
+        assert "must be int or str" in error_msg or "int_type" in error_msg or "string_type" in error_msg
 
     def test_aiperf_config_sweep_empty_list(self, valid_base_config):
         """Test that empty sweep list raises validation error."""
@@ -302,9 +294,7 @@ def test_aiperf_config_multiple_invalid_sweep_keys(self, valid_base_config):
 
     def test_aiperf_config_get_output_base_path(self, valid_base_config):
         """Test get_output_base_path method."""
-        config = AIPerfConfig(
-            output_base_dir="custom_results", base_config=valid_base_config
-        )
+        config = AIPerfConfig(output_base_dir="custom_results", base_config=valid_base_config)
         path = config.get_output_base_path()
         assert isinstance(path, Path)
         assert str(path) == "custom_results"