diff --git a/benchmark/locust/README.md b/benchmark/locust/README.md new file mode 100644 index 0000000000..1fdc4e4205 --- /dev/null +++ b/benchmark/locust/README.md @@ -0,0 +1,111 @@ +# Locust Load Testing for NeMo Guardrails + +This directory contains a Locust-based load testing framework for the NeMo Guardrails OpenAI-compatible server. + +## Introduction + +The [Locust](https://locust.io/) stress-testing tool ramps up concurrent users making API calls to the `/v1/chat/completions` endpoint of an OpenAI-compatible LLM with configurable parameters. +This complements [ai-perf](https://github.com/ai-dynamo/aiperf), which measures steady-state performance. Locust instead focuses on ramping up load potentially beyond what a system can handle, and measure how gracefully it degrades under higher-than-expected load. + +## Getting Started + +### Prerequisites + +These steps have been tested with Python 3.11.11. + +1. **Create a virtual environment in which to install Locust and other benchmarking tools** + + ```bash + $ mkdir ~/env + $ python -m venv ~/env/benchmark_env + ``` + +2. **Activate environment and install dependencies in the virtual environment** + + ```bash + $ source ~/env/benchmark_env/bin/activate + (benchmark_env) $ pip install -r benchmark/requirements.txt + ``` + +## Running Benchmarks + +The Locust benchmarks uses YAML configuration file to configure load-testing parameters. +To get started and load-test a model hosted at `http://localhost:8000`, use the following command. +Set `headless: false` in your YAML config to use Locust's interactive web UI. Then open http://localhost:8089 to control the test and view real-time metrics. + + ```bash + (benchmark_env) $ python -m benchmark.locust benchmark/locust/configs/local.yaml + ``` + +### CLI Options + +The `benchmark.locust` CLI supports the following options: + +```bash +python -m benchmark.locust [OPTIONS] CONFIG_FILE +``` + +**Arguments:** +- `CONFIG_FILE`: Path to YAML configuration file (required) + +**Options:** +- `--dry-run`: Print commands without executing them +- `--verbose`: Enable verbose logging and debugging information + +## Configuration Options + +All configuration is done via YAML files. The following fields are supported: + +### Required Fields + +- `config_id`: Guardrails configuration ID to use +- `model`: Model name to send in requests + +### Optional Fields + +- `host`: Server base URL (default: `http://localhost:8000`) +- `users`: Maximum concurrent users (default: `256`, minimum: `1`) +- `spawn_rate`: Users spawned per second (default: `10`, minimum: `0.1`) +- `run_time`: Test duration in seconds (default: `60`, minimum: `1`) +- `message`: Message content to send (default: `"Hello, what can you do?"`) +- `headless`: Run without web UI (default: `true`) +- `output_base_dir`: Directory for test results (default: `"locust_results"`) + +## Load Test Behavior + +- **Request Type**: 100% POST `/v1/chat/completions` requests +- **Wait Time**: Zero wait time between requests (continuous hammering) +- **Ramp-up**: Users spawn gradually at the specified `spawn_rate` +- **Message Content**: Static message content (configurable via `message` field) + +## Output + +### Headless Mode + +When run in headless mode, results are saved to timestamped directories: + +``` +locust_results/ +└── YYYYMMDD_HHMMSS/ + ├── report.html # HTML report with charts + ├── run_metadata.json # Test configuration metadata + ├── stats.csv # Request statistics + ├── stats_failures.csv # Failure statistics + └── stats_history.csv # Statistics over time +``` + +### Web UI Mode + +Real-time metrics are displayed in the web interface at http://localhost:8089, including: +- Requests per second (RPS) +- Response time percentiles (50th, 95th, 99th) +- Failure rate +- Number of users + +### Troubleshooting + +If you see validation errors: +- Ensure all required fields (`config_id`, `model`) are present in your YAML config +- Check that the `config_id` matches a configuration on your server +- Verify that numeric values meet minimum requirements (e.g., `users >= 1`, `spawn_rate >= 0.1`) +- Ensure `host` starts with `http://` or `https://` diff --git a/benchmark/locust/__init__.py b/benchmark/locust/__init__.py new file mode 100644 index 0000000000..35669c048a --- /dev/null +++ b/benchmark/locust/__init__.py @@ -0,0 +1,15 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/benchmark/locust/__main__.py b/benchmark/locust/__main__.py new file mode 100644 index 0000000000..569b86c235 --- /dev/null +++ b/benchmark/locust/__main__.py @@ -0,0 +1,22 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Entry point for running the Locust load test CLI as a module: python -m benchmark.locust""" + +from benchmark.locust.run_locust import app + +if __name__ == "__main__": + app() diff --git a/benchmark/locust/configs/local.yaml b/benchmark/locust/configs/local.yaml new file mode 100644 index 0000000000..b916d998fb --- /dev/null +++ b/benchmark/locust/configs/local.yaml @@ -0,0 +1,18 @@ +# Example Locust load test configuration for NeMo Guardrails + +# Server details +host: "http://localhost:8000" +config_id: "my-guardrails-config" +model: "meta/llama-3.3-70b-instruct" + +# Load test parameters +users: 1024 # Maximum number of concurrent users +spawn_rate: 16 # Users spawned per second +run_time: 120 # Test duration in seconds + +# Request configuration +message: "Hello, what can you do?" + +# Output configuration +headless: true # Set to true for headless mode, false for web UI +output_base_dir: "locust_results" # Directory for test results diff --git a/benchmark/locust/locust_models.py b/benchmark/locust/locust_models.py new file mode 100644 index 0000000000..a43ce74a8c --- /dev/null +++ b/benchmark/locust/locust_models.py @@ -0,0 +1,82 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Pydantic models for Locust load test configuration validation. +""" + +from pathlib import Path + +from pydantic import BaseModel, Field, field_validator + + +class LocustConfig(BaseModel): + """Configuration for a Locust load-test run""" + + # Server details + host: str = Field( + default="http://localhost:8000", + description="Base URL of the NeMo Guardrails server to test", + ) + config_id: str = Field(..., description="Guardrails configuration ID to use") + model: str = Field(..., description="Model name to use in requests") + + # Load test parameters + users: int = Field( + default=256, + ge=1, + description="Maximum number of concurrent users", + ) + spawn_rate: float = Field( + default=10, + ge=0.1, + description="Rate at which users are spawned (users/second)", + ) + run_time: int = Field( + default=60, + ge=1, + description="Test duration in seconds", + ) + + # Request configuration + message: str = Field( + default="Hello, what can you do?", + description="Message content to send in chat completion requests", + ) + + # Output configuration + headless: bool = Field( + default=True, + description="Run in headless mode without web UI", + ) + + output_base_dir: str = Field( + default="locust_results", + description="Base directory for load test results", + ) + + @field_validator("host") + @classmethod + def validate_host(cls, v: str) -> str: + """Ensure host starts with http:// or https://""" + if not v.startswith(("http://", "https://")): + raise ValueError("Host must start with http:// or https://") + # Remove trailing slash if present + return v.rstrip("/") + + def get_output_base_path(self) -> Path: + """Get the base output directory as a Path object.""" + return Path(self.output_base_dir) diff --git a/benchmark/locust/locustfile.py b/benchmark/locust/locustfile.py new file mode 100644 index 0000000000..d1d7638d38 --- /dev/null +++ b/benchmark/locust/locustfile.py @@ -0,0 +1,73 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Locust load test file for NeMo Guardrails OpenAI-compatible server. + +This file defines the load test behavior. It can be run directly with: + locust -f locustfile.py --host http://localhost:8000 + +Or via the Typer CLI wrapper: + python -m benchmark.locust run --config-file config.yaml +""" + +import os + +from locust import HttpUser, constant, task + + +class GuardrailsUser(HttpUser): + """ + Simulated user that continuously sends chat completion requests to the + NeMo Guardrails server. + + Each user will continuously send requests with no wait time between them + (continuous hammering). The load is distributed such that 99% of requests + go to chat completions. + """ + + # No wait time between requests (continuous hammering) + wait_time = constant(0) + + def on_start(self): + """Called when a simulated user starts. + Uses environment variables to pass the Guardrails config_id, model, and message""" + # Get configuration from environment variables set by the CLI wrapper + self.config_id = os.getenv("LOCUST_CONFIG_ID", "default") + self.model = os.getenv("LOCUST_MODEL", "mock-llm") + self.message = os.getenv("LOCUST_MESSAGE", "Hello, what can you do?") + + @task + def chat_completion(self): + """ + Send a Guardrails chat completion request (/v1/chat/completions) + """ + payload = { + "model": self.model, + "messages": [{"role": "user", "content": self.message}], + "guardrails": {"config_id": self.config_id}, + } + + with self.client.post( + "/v1/chat/completions", + json=payload, + catch_response=True, + name="/v1/chat/completions", + ) as response: + if response.status_code == 200: + response.success() + else: + response.failure(f"Got status code {response.status_code}: {response.text}") diff --git a/benchmark/locust/run_locust.py b/benchmark/locust/run_locust.py new file mode 100644 index 0000000000..888927b8ee --- /dev/null +++ b/benchmark/locust/run_locust.py @@ -0,0 +1,256 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Typer CLI wrapper for running Locust load tests against NeMo Guardrails server. + +This module provides a command-line interface for running load tests, supporting +both direct CLI arguments and YAML configuration files. +""" + +import json +import logging +import os +import subprocess +import sys +from datetime import datetime +from pathlib import Path +from typing import Optional + +import httpx +import typer +import yaml +from pydantic import ValidationError + +from benchmark.locust.locust_models import LocustConfig + +log = logging.getLogger(__name__) +log.setLevel(logging.INFO) + +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) + +log.addHandler(console_handler) + +app = typer.Typer( + help="Locust load testing application for NeMo Guardrails", + add_completion=False, +) + + +class LocustRunner: + """Run Locust load tests against NeMo Guardrails server.""" + + def __init__(self, config: LocustConfig): + self.config = config + self.locustfile_path = Path(__file__).parent / "locustfile.py" + + def _check_service(self) -> None: + """Check if the NeMo Guardrails server is up before running tests.""" + url = f"{self.config.host}/health" + log.debug("Checking service is up at %s", url) + + try: + # Try a simple request to verify the server is accessible + response = httpx.get(url, timeout=5) + except httpx.ConnectError as e: + raise RuntimeError(f"ConnectError accessing {url}: {e}") + except httpx.TimeoutException as e: + raise RuntimeError(f"HTTP Timeout accessing {url}: {e}") + + if response.is_error: + raise RuntimeError(f"Error {response.status_code} connecting to {url}: {response.text}") + + try: + if response.json().get("status") != "healthy": + raise RuntimeError(f"Service at {url} is unhealthy: {response.text}") + except json.decoder.JSONDecodeError as e: + raise RuntimeError(f"Error: response {response.text} couldn't be parsed as JSON: {e}") + + log.info("Successfully connected to server at %s", self.config.host) + + def _build_locust_command(self, output_dir: Optional[Path] = None) -> list[str]: + """Build the Locust command with all parameters.""" + cmd = ["locust", "-f", str(self.locustfile_path)] + + # Host + cmd.extend(["--host", self.config.host]) + + # User and spawn rate + cmd.extend(["--users", str(self.config.users)]) + cmd.extend(["--spawn-rate", str(self.config.spawn_rate)]) + cmd.extend(["--run-time", f"{self.config.run_time}s"]) + + # Headless mode + if self.config.headless: + cmd.append("--headless") + cmd.append("--only-summary") # only print last latency table + + # Add output files for headless mode + if output_dir: + html_file = output_dir / "report.html" + csv_prefix = output_dir / "stats" + cmd.extend(["--html", str(html_file)]) + cmd.extend(["--csv", str(csv_prefix)]) + + log.debug("Locust command: %s", " ".join(cmd)) + return cmd + + def _save_run_metadata(self, output_dir: Path, command: list[str], start_time: datetime) -> None: + """Save metadata about the load test run.""" + metadata = { + "start_time": start_time.isoformat(), + "config": self.config.model_dump(), + "command": " ".join([str(c) for c in command]), + } + + metadata_file = output_dir / "run_metadata.json" + with open(metadata_file, "w", encoding="utf-8") as f: + json.dump(metadata, f, indent=2) + + log.debug("Saved run metadata to %s", metadata_file) + + def _create_output_path(self, base_dir: str) -> Path: + """Create timestamped output directory for test results.""" + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + output_path = Path(base_dir) / Path(timestamp) + output_path.mkdir(parents=True, exist_ok=True) + return output_path + + def run(self, dry_run: bool) -> int: + """Run the Locust load test.""" + # Check service availability + try: + self._check_service() + except RuntimeError as e: + log.error(str(e)) + return 1 + + # Build command + output_path = self._create_output_path(self.config.output_base_dir) + command = self._build_locust_command(output_path) + + # Save metadata + start_time = datetime.now() + self._save_run_metadata(output_path, command, start_time) + log.info("Saving metadata to: %s", output_path) + + # Set environment variables for the locustfile + env = os.environ.copy() + env["LOCUST_CONFIG_ID"] = self.config.config_id + env["LOCUST_MODEL"] = self.config.model + env["LOCUST_MESSAGE"] = self.config.message + + # Log test configuration + log.info("Starting Locust load test") + log.info("Config: %s", self.config.model_dump_json()) + + rampup_seconds = min(int(self.config.users / self.config.spawn_rate), self.config.run_time) + steady_state_seconds = self.config.run_time - rampup_seconds + log.info("Duration: rampup: %is, steady-state %is", rampup_seconds, steady_state_seconds) + + if not self.config.headless: + log.info("Web UI will be available at: http://localhost:8089") + + try: + # For dry-run, just print out the command + if dry_run: + log.info("Dry run mode. Command: %s", " ".join(command)) + return 0 + + result = subprocess.run(command, env=env, check=False) + + if result.returncode == 0: + log.info("Load test completed successfully") + if output_path: + log.info("Results saved to: %s", output_path) + else: + log.error("Load test failed with exit code %s", result.returncode) + + return result.returncode + + except KeyboardInterrupt: + log.warning("Load test interrupted by user") + return 130 + except Exception as e: + log.error("Error running load test: %s", e) + return 1 + + +def _load_config_from_yaml(config_file: Path) -> LocustConfig: + """Load and validate configuration from YAML file.""" + try: + with open(config_file, "r", encoding="utf-8") as f: + config_data = yaml.safe_load(f) + + config = LocustConfig(**config_data) + return config + + except FileNotFoundError: + log.error("Configuration file not found: %s", config_file) + sys.exit(1) + except yaml.YAMLError as e: + log.error("Error parsing YAML configuration: %s", e) + sys.exit(1) + except ValidationError as e: + log.error("Configuration validation error:\n%s", e) + sys.exit(1) + except Exception as e: + log.error("Unexpected error loading configuration: %s", e) + sys.exit(1) + + +@app.command() +def run( + config_file: Path = typer.Argument( + help="Path to YAML configuration file", + exists=True, + file_okay=True, + dir_okay=False, + readable=True, + ), + dry_run: bool = typer.Option( + False, + "--dry-run", + help="Print commands without executing them", + ), + verbose: bool = typer.Option( + False, + "--verbose", + help="Print additional debugging information during run", + ), +): + """ + Run Locust load test using provided config file + """ + if verbose: + log.setLevel(logging.DEBUG) + + # Load config from file if provided + if config_file: + locust_config = _load_config_from_yaml(config_file) + + # Create and run the test + runner = LocustRunner(locust_config) + exit_code = runner.run(dry_run) + + raise typer.Exit(code=exit_code) + + +if __name__ == "__main__": + app() diff --git a/benchmark/requirements.txt b/benchmark/requirements.txt index a86e8ec3c5..8e93ae9182 100644 --- a/benchmark/requirements.txt +++ b/benchmark/requirements.txt @@ -19,3 +19,6 @@ numpy>=2.3.2 httpx>=0.24.1 typer>=0.8 pyyaml>=6.0 + +# --- locust load testing dependencies --- +locust>=2.0.0 diff --git a/benchmark/tests/test_locust_models.py b/benchmark/tests/test_locust_models.py new file mode 100644 index 0000000000..d9838fd029 --- /dev/null +++ b/benchmark/tests/test_locust_models.py @@ -0,0 +1,161 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Tests for Locust load test configuration models. +""" + +from pathlib import Path + +import pytest +from pydantic import ValidationError + +from benchmark.locust.locust_models import LocustConfig + + +class TestLocustConfig: + """Test the LocustConfig model.""" + + def test_config_minimal_valid_with_defaults(self): + """Test creating LocustConfig with minimal required fields and verify all defaults.""" + config = LocustConfig( + config_id="test-config", + model="test-model", + ) + + # Verify required fields + assert config.config_id == "test-config" + assert config.model == "test-model" + + # Verify all defaults + assert config.host == "http://localhost:8000" + assert config.users == 256 + assert config.spawn_rate == 10 + assert config.run_time == 60 + assert config.message == "Hello, what can you do?" + assert config.headless is True + assert config.output_base_dir == "locust_results" + + def test_config_with_all_fields(self): + """Test creating LocustBaseConfig with all fields specified.""" + config = LocustConfig( + host="http://example.com:9000", + config_id="my-config", + model="my-model", + users=100, + spawn_rate=5.5, + run_time=120, + message="Custom message", + headless=True, + output_base_dir="/tmp/locust", + ) + assert config.host == "http://example.com:9000" + assert config.config_id == "my-config" + assert config.model == "my-model" + assert config.users == 100 + assert config.spawn_rate == 5.5 + assert config.run_time == 120 + assert config.message == "Custom message" + assert config.headless is True + + def test_config_missing_required_fields(self): + """Test that missing required fields raise validation error.""" + with pytest.raises(ValidationError) as exc_info: + LocustConfig( + host="http://localhost:8000", + # Missing config_id and model + ) + errors = exc_info.value.errors() + error_fields = {err["loc"][0] for err in errors} + assert "config_id" in error_fields + assert "model" in error_fields + + def test_config_host_without_protocol(self): + """Test that host without http:// or https:// raises validation error.""" + with pytest.raises(ValidationError) as exc_info: + LocustConfig( + host="localhost:8000", # Missing http:// + config_id="test-config", + model="test-model", + ) + error_msg = str(exc_info.value) + assert "Host must start with http:// or https://" in error_msg + + def test_config_host_with_https(self): + """Test that host with https:// is valid.""" + config = LocustConfig( + host="https://secure.example.com", + config_id="test-config", + model="test-model", + ) + assert config.host == "https://secure.example.com" + + def test_config_host_trailing_slash_removed(self): + """Test that trailing slash in host is removed.""" + config = LocustConfig( + host="http://localhost:8000/", + config_id="test-config", + model="test-model", + ) + assert config.host == "http://localhost:8000" + + def test_config_host_multiple_trailing_slashes(self): + """Test that multiple trailing slashes are removed.""" + config = LocustConfig( + host="http://localhost:8000///", + config_id="test-config", + model="test-model", + ) + assert config.host == "http://localhost:8000" + + +class TestLocustConfigHelpers: + """Test helper methods on LocustConfig model.""" + + @pytest.fixture + def config(self) -> LocustConfig: + """Helper to get a valid base config.""" + return LocustConfig( + config_id="test-config", + model="test-model", + ) + + def test_locust_config_get_output_base_path(self, config): + """Test get_output_base_path method.""" + config.output_base_dir = "custom_results" + + path = config.get_output_base_path() + assert isinstance(path, Path) + assert str(path) == "custom_results" + + def test_locust_config_get_output_base_path_default(self, config): + """Test get_output_base_path method with default output_base_dir.""" + path = config.get_output_base_path() + assert isinstance(path, Path) + assert str(path) == "locust_results" + + def test_locust_config_with_dict(self): + """Test creating LocustConfig with dict base_config.""" + config = LocustConfig( + **{ + "config_id": "test-config", + "model": "test-model", + "users": 100, + } + ) + assert config.config_id == "test-config" + assert config.model == "test-model" + assert config.users == 100 diff --git a/benchmark/tests/test_run_locust.py b/benchmark/tests/test_run_locust.py new file mode 100644 index 0000000000..0aaf1c061f --- /dev/null +++ b/benchmark/tests/test_run_locust.py @@ -0,0 +1,474 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Tests for Locust load test CLI runner. +""" + +import json +from datetime import datetime +from json.decoder import JSONDecodeError +from pathlib import Path +from typing import Any, Dict, Optional +from unittest.mock import Mock, patch +from urllib.parse import urljoin + +import httpx +import pytest +import yaml +from typer.testing import CliRunner + +from benchmark.locust.locust_models import LocustConfig +from benchmark.locust.run_locust import LocustRunner, _load_config_from_yaml, app + + +@pytest.fixture +def create_config_data(tmp_path): + """Returns a function with sample basic config, and allows mutation of fields to cover + more cases or add extra fields""" + + def _create_config( + config_id="test-config", + model="test-model", + host="http://localhost:8000", + users=256, + spawn_rate=10, + run_time=60, + message="Hello, what can you do?", + headless=False, + output_base_dir=str(tmp_path), + **extra_config, + ): + config_data = { + "host": host, + "config_id": config_id, + "model": model, + "users": users, + "spawn_rate": spawn_rate, + "run_time": run_time, + "message": message, + "headless": headless, + "output_base_dir": output_base_dir, + } + + # Merge any extra config parameters + if extra_config: + config_data.update(extra_config) + + return config_data + + return _create_config + + +@pytest.fixture +def create_config_file(tmp_path, create_config_data): + """Fixture to write config data to a file and return the path.""" + + def _write_config_file( + extra_base_config: Optional[Dict[str, Any]] = None, + filename: Optional[str] = "config.yml", + ) -> Path: + """Apply extra base config to config data, write to file and return the path.""" + + # Unpack extra_base_config as kwargs if provided + if extra_base_config: + config_data = create_config_data(**extra_base_config) + else: + config_data = create_config_data() + + config_file = tmp_path / filename + config_file.write_text(yaml.dump(config_data)) + return config_file + + return _write_config_file + + +class TestLocustRunner: + """Test LocustRunner class.""" + + @pytest.fixture + def valid_config(self): + """Get a valid LocustConfig for testing.""" + return LocustConfig( + host="http://localhost:8000", + config_id="test-config", + model="test-model", + users=10, + spawn_rate=2, + run_time=30, + ) + + @pytest.fixture + def runner(self, valid_config): + """Get a LocustRunner instance for testing.""" + return LocustRunner(valid_config) + + def _service_health_endpoint(self, runner: LocustRunner): + """The endpoint used ot check if the service is healthy""" + return urljoin(runner.config.host, "health") + + def test_runner_init(self, valid_config): + """Test LocustRunner initialization.""" + runner = LocustRunner(valid_config) + assert runner.config == valid_config + assert runner.locustfile_path.exists() + assert runner.locustfile_path.name == "locustfile.py" + + def test_check_service_success(self, runner): + """Test _check_service with successful connection.""" + with patch("httpx.get") as mock_get: + mock_response = Mock() + mock_response.is_error = False + mock_response.json.return_value = {"status": "healthy", "timestamp": 1770675471} + mock_get.return_value = mock_response + + # Should not raise + runner._check_service() + mock_get.assert_called_once_with(self._service_health_endpoint(runner), timeout=5) + + def test_check_service_connection_error(self, runner): + """Test _check_service with httpx.ConnectError""" + with patch("httpx.get") as mock_get: + mock_get.side_effect = httpx.ConnectError("Connection refused") + + with pytest.raises(RuntimeError) as exc_info: + runner._check_service() + + mock_get.assert_called_once_with(self._service_health_endpoint(runner), timeout=5) + assert ( + exc_info.value.args[0] + == f"ConnectError accessing {self._service_health_endpoint(runner)}: Connection refused" + ) + + def test_check_service_timeout_error(self, runner): + """Test _check_service when httpx.get times out""" + with patch("httpx.get") as mock_get: + mock_get.side_effect = httpx.TimeoutException("httpx.ConnectTimeout: The connection operation timed out") + + with pytest.raises(RuntimeError) as exc_info: + runner._check_service() + + mock_get.assert_called_once_with(self._service_health_endpoint(runner), timeout=5) + assert ( + exc_info.value.args[0] + == f"HTTP Timeout accessing {self._service_health_endpoint(runner)}: httpx.ConnectTimeout: The connection operation timed out" + ) + + def test_check_service_error_response(self, runner): + """Test _check_service with non-200 response code""" + with patch("httpx.get") as mock_get: + mock_response = Mock() + mock_response.is_error = True + mock_response.status_code = 404 + mock_response.text = '{"detail":"Not Found"}' + mock_response.json.return_value = json.dumps(mock_response.text) + mock_get.return_value = mock_response + + with pytest.raises(RuntimeError) as exc_info: + runner._check_service() + + mock_get.assert_called_once_with(self._service_health_endpoint(runner), timeout=5) + assert ( + exc_info.value.args[0] + == f"Error {mock_response.status_code} connecting to {self._service_health_endpoint(runner)}: {mock_response.text}" + ) + + def test_check_service_unhealthy_response(self, runner): + """Test _check_service with 200 response from an unhealthy service""" + with patch("httpx.get") as mock_get: + mock_response = Mock() + mock_response.is_error = False + mock_response.status_code = 200 # Successful HTTP request .. + mock_response.text = ( + '{"status":"unhealthy","timestamp":1770677847}' # .. but the application itself is unhealthy + ) + mock_response.json.return_value = json.loads(mock_response.text) + mock_get.return_value = mock_response + + with pytest.raises(RuntimeError) as exc_info: + runner._check_service() + + mock_get.assert_called_once_with(self._service_health_endpoint(runner), timeout=5) + assert ( + exc_info.value.args[0] + == f"Service at {self._service_health_endpoint(runner)} is unhealthy: {mock_response.text}" + ) + + def test_check_service_invalid_json(self, runner): + """Test _check_service with an invalid JSON response""" + with patch("httpx.get") as mock_get: + mock_response = Mock() + mock_response.is_error = False + mock_response.status_code = 200 + mock_response.text = "{'key': 'value'}" + mock_response.json.side_effect = JSONDecodeError( + "Expecting property name enclosed in double quotes: line 1 column 2 (char 1)", "{'key': 'value'}", 1 + ) + mock_get.return_value = mock_response + + with pytest.raises(RuntimeError) as exc_info: + runner._check_service() + + mock_get.assert_called_once_with(self._service_health_endpoint(runner), timeout=5) + assert ( + exc_info.value.args[0] + == "Error: response {'key': 'value'} couldn't be parsed as JSON: Expecting property name enclosed in double quotes: line 1 column 2 (char 1): line 1 column 2 (char 1)" + ) + + def test_build_locust_command_basic(self, runner): + """Test building basic Locust command.""" + cmd = runner._build_locust_command() + assert cmd[0] == "locust" + + cmd_string = " ".join(cmd) + assert "--host http://localhost:8000 --users 10 --spawn-rate 2.0 --run-time 30s --headless" in cmd_string + + def test_build_locust_command_headless(self, runner, tmp_path): + """Test building Locust command in headless mode.""" + runner.config.headless = True + output_dir = tmp_path / "output" + output_dir.mkdir() + + cmd = runner._build_locust_command(output_dir) + + assert "--headless" in cmd + assert "--html" in cmd + assert "--csv" in cmd + + def test_save_run_metadata(self, runner, tmp_path): + """Test saving run metadata to file.""" + output_dir = tmp_path / "output" + output_dir.mkdir() + start_time = datetime.now() + command = ["locust", "-f", "locustfile.py"] + + runner._save_run_metadata(output_dir, command, start_time) + + metadata_file = output_dir / "run_metadata.json" + assert metadata_file.exists() + + with open(metadata_file) as f: + metadata = json.load(f) + + assert "start_time" in metadata + assert "config" in metadata + assert "command" in metadata + assert metadata["config"]["config_id"] == "test-config" + assert metadata["config"]["model"] == "test-model" + + def test_create_output_dir(self, runner, tmp_path): + """Test creating timestamped output directory.""" + base_dir = str(tmp_path) + "results" + + output_dir = runner._create_output_path(base_dir) + + assert output_dir.exists() + assert output_dir.is_dir() + assert output_dir.parent == Path(base_dir) + # Check that directory name looks like a timestamp + assert len(output_dir.name) == len("20250101_120000") + + def test_run_success_headless(self, runner, tmp_path): + """Test successful run in headless mode.""" + runner.config.headless = True + + with patch.object(runner, "_check_service"), patch("subprocess.run") as mock_run: + mock_result = Mock() + mock_result.returncode = 0 + mock_run.return_value = mock_result + + exit_code = runner.run(dry_run=False) + + assert exit_code == 0 + mock_run.assert_called_once() + + # Check that command was built correctly + call_args = mock_run.call_args + assert call_args[0][0][0] == "locust" + assert "--headless" in call_args[0][0] + + # Check that env variables were set + env = call_args[1]["env"] + assert env["LOCUST_CONFIG_ID"] == "test-config" + assert env["LOCUST_MODEL"] == "test-model" + assert env["LOCUST_MESSAGE"] == "Hello, what can you do?" + + def test_run_success_web_ui(self, runner): + """Test successful run in web UI mode.""" + runner.config.headless = False + + with patch.object(runner, "_check_service"), patch("subprocess.run") as mock_run: + mock_result = Mock() + mock_result.returncode = 0 + mock_run.return_value = mock_result + + exit_code = runner.run(dry_run=False) + + assert exit_code == 0 + mock_run.assert_called_once() + + # Check that command was built correctly + call_args = mock_run.call_args + assert call_args[0][0][0] == "locust" + assert "--headless" not in call_args[0][0] + + def test_run_failure(self, runner): + """Test run with command failure.""" + with patch.object(runner, "_check_service"), patch("subprocess.run") as mock_run: + mock_result = Mock() + mock_result.returncode = 1 + mock_run.return_value = mock_result + + exit_code = runner.run(dry_run=False) + + assert exit_code == 1 + + def test_run_keyboard_interrupt(self, runner): + """Test run interrupted by user.""" + with patch.object(runner, "_check_service"), patch("subprocess.run") as mock_run: + mock_run.side_effect = KeyboardInterrupt() + + exit_code = runner.run(dry_run=False) + + assert exit_code == 130 + + def test_run_service_check_failure(self, runner): + """Test run when service check fails.""" + with patch.object(runner, "_check_service") as mock_check: + mock_check.side_effect = RuntimeError("Service unavailable") + + exit_code = runner.run(dry_run=False) + + assert exit_code == 1 + + def test_run_exception(self, runner): + """Test run with unexpected exception.""" + with patch.object(runner, "_check_service"), patch("subprocess.run") as mock_run: + mock_run.side_effect = Exception("Unexpected error") + + exit_code = runner.run(dry_run=False) + + assert exit_code == 1 + + +class TestLoadConfigFromYaml: + """Test _load_config_from_yaml function.""" + + def test_load_valid_config(self, create_config_file): + """Test loading a valid config file.""" + config_file = create_config_file() + + config = _load_config_from_yaml(config_file) + + assert isinstance(config, LocustConfig) + assert config.config_id == "test-config" + assert config.model == "test-model" + + def test_load_config_file_not_found(self, tmp_path): + """Test loading non-existent config file.""" + config_file = tmp_path / "nonexistent.yml" + + with pytest.raises(SystemExit) as exc_info: + _load_config_from_yaml(config_file) + + assert exc_info.value.code == 1 + + def test_load_config_invalid_yaml(self, tmp_path): + """Test loading file with invalid YAML.""" + config_file = tmp_path / "invalid.yml" + config_file.write_text("invalid: yaml: content: [") + + with pytest.raises(SystemExit) as exc_info: + _load_config_from_yaml(config_file) + + assert exc_info.value.code == 1 + + def test_load_config_validation_error(self, tmp_path): + """Test loading file with validation errors.""" + config_file = tmp_path / "invalid.yml" + config_data = { + "batch_name": "test", + "base_config": { + "config_id": "test-config", + # Missing required model field + }, + } + config_file.write_text(yaml.dump(config_data)) + + with pytest.raises(SystemExit) as exc_info: + _load_config_from_yaml(config_file) + + assert exc_info.value.code == 1 + + def test_load_config_unexpected_error(self, tmp_path): + """Test loading config with unexpected error.""" + config_file = tmp_path / "config.yml" + config_file.write_text("valid_yaml: true") + + with patch("yaml.safe_load") as mock_load: + mock_load.side_effect = Exception("Unexpected error") + + with pytest.raises(SystemExit) as exc_info: + _load_config_from_yaml(config_file) + + assert exc_info.value.code == 1 + + +class TestCLI: + """Test CLI commands.""" + + @pytest.fixture + def cli_runner(self): + """Get a Typer CLI test runner.""" + return CliRunner() + + def test_run_command_missing_config_file(self, cli_runner): + """Test run command without required config file.""" + result = cli_runner.invoke(app, []) + + assert result.exit_code != 0 # Should fail + # Check that error message mentions missing argument + assert "missing" in result.stdout.lower() or result.exit_code == 2 + + def test_run_command_config_file_not_found(self, cli_runner, tmp_path): + """Test run command with non-existent config file.""" + nonexistent_file = tmp_path / "nonexistent.yaml" + result = cli_runner.invoke(app, [str(nonexistent_file)]) + + assert result.exit_code != 0 # Should fail + + def test_run_command_with_config_file(self, cli_runner, create_config_file): + """Test run command with YAML config file.""" + config_file = create_config_file() + + with patch("benchmark.locust.run_locust.LocustRunner") as mock_runner_class: + mock_runner = Mock() + mock_runner.run.return_value = 0 + mock_runner_class.return_value = mock_runner + + result = cli_runner.invoke( + app, + [str(config_file)], + catch_exceptions=False, + ) + + assert result.exit_code == 0, f"Output: {result.stdout}" + mock_runner_class.assert_called_once() + config = mock_runner_class.call_args[0][0] + assert config.config_id == "test-config" + assert config.model == "test-model" + # Verify run was called with dry_run parameter + mock_runner.run.assert_called_once_with(False)