Metadata-Version: 2.3
Name: ml-loadtest
Version: 1.12.2
Summary: Adaptive load testing tool for ML inference APIs with dynamic scaling and regression detection
License: MIT
Keywords: load-testing,locust,ml,api-testing,performance,inference,benchmarking
Requires-Python: >=3.10
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Software Development :: Testing
Classifier: Topic :: System :: Benchmark
Requires-Dist: dacite (>=1.8,<2.0)
Requires-Dist: locust (>=2.20,<3.0)
Requires-Dist: notion-client (>=2.0,<3.0)
Requires-Dist: numpy (>=1.23.5,<2.0.0)
Requires-Dist: psutil (>=5.9)
Project-URL: Homepage, https://github.com/IncodeTechnologies/ml-load-testing-tool
Project-URL: Issues, https://github.com/IncodeTechnologies/ml-load-testing-tool/issues
Project-URL: Repository, https://github.com/IncodeTechnologies/ml-load-testing-tool
Description-Content-Type: text/markdown

# ML Load Testing Tool

Adaptive load testing tool for ML inference APIs. Uses Locust to dynamically scale concurrent users based on P99 latency targets, then analyzes results for regression detection and rate limit recommendations.

## Features

- **Adaptive Scaling**: Automatically adjusts concurrent users based on P99 latency targets
- **Multi-Mode Testing**: Isolated endpoint latency, individual capacity, production mix, and spike modes
- **Regression Detection**: Compare test results against baselines to catch performance degradations
- **Rate Limit Recommendations**: Calculates rate limits directly from measured endpoint capacity
- **Notion Integration**: Sync results to Notion for tracking and visualization

## Installation

### From GitHub

```bash
pip install git+https://github.com/IncodeTechnologies/ml-load-testing-tool.git
```

### From Source

```bash
git clone https://github.com/IncodeTechnologies/ml-load-testing-tool.git
cd ml-load-testing-tool
poetry install
```

## Quick Start

The tool requires a weights module that defines which endpoints to test. Use bundled examples or create your own:

```bash
# Test with bundled example TaskSets
locust -f $(ml-loadtest-file) \
  --host http://api:8000 \
  --weights-module loadtest.distribution_weights

# Run headless with specific parameters
locust -f $(ml-loadtest-file) \
  --host http://api:8000 \
  --weights-module loadtest.distribution_weights \
  --users 32 \
  --spawn-rate 4 \
  --run-time 60s \
  --headless
```

The `ml-loadtest-file` command prints the path to the installed locustfile, giving you full access to all Locust CLI parameters.

**Note:** The `--weights-module` parameter is required. It specifies a Python module containing a `production_weights` dictionary that maps TaskSet classes to their relative weights.

## Usage

### Two Usage Patterns

**Pattern 1: Direct Path (Recommended for CLI)**

```bash
# Get full locust CLI access with installed package
locust -f $(ml-loadtest-file) --host http://api:8000 [any locust params]
```

**Pattern 2: Local Import (Recommended for Customization)**

Create a local `locustfile.py` in your project:

```python
# Import everything from the installed package
from ml_loadtest.locustfile import *

# Optionally override settings or add custom logic here
```

Then run:

```bash
locust -f locustfile.py --host http://api:8000
```

### Load Testing

The tool supports four test modes via the `--mode` parameter. You can run a single mode or multiple modes in sequence (space-separated). Mode names are case-insensitive. The default is `ENDPOINT_LATENCY INDIVIDUAL PRODUCTION`, in that order.

#### 1. Isolated Endpoint Latency

Measure client-observed HTTP latency with no concurrent load:

```bash
locust -f $(ml-loadtest-file) \
  --host http://api:8000 \
  --weights-module loadtest.distribution_weights \
  --mode ENDPOINT_LATENCY
```

Each configured endpoint runs sequentially with exactly one Locust user. By default, requests have a 10-second timeout, the first successful request is discarded as warmup, and the next 20 successful requests are recorded for NumPy p50, p95, and p99 calculation. Failed requests are excluded; a success resets the consecutive-failure counter. By default, three consecutive failures mark that endpoint incomplete, leave its percentiles unavailable, and advance to the next endpoint. All four protocol values are configurable with the options below.

This measures elapsed time observed by Locust from HTTP dispatch through receipt of the response body. Task payload preparation and response validation are outside the measured HTTP duration. With only 20 observations, p99 is effectively a near-maximum sample.

#### 2. Individual Endpoint Testing

Test each endpoint separately to find individual capacity:

```bash
locust -f $(ml-loadtest-file) \
  --host http://api:8000 \
  --weights-module loadtest.distribution_weights \
  --mode INDIVIDUAL \
  --target-p99-ms 500 \
  --max-users 100
```

#### 3. Production Mix Testing

Test with production traffic distribution:

```bash
locust -f $(ml-loadtest-file) \
  --host http://api:8000 \
  --weights-module ml_loadtest.examples.distribution_weights \
  --mode PRODUCTION \
  --target-p99-ms 500
```

#### 4. Spike Testing

Test sudden traffic spikes:

```bash
locust -f $(ml-loadtest-file) \
  --host http://api:8000 \
  --weights-module ml_loadtest.examples.distribution_weights \
  --mode SPIKE \
  --spike-target-rps 1000 \
  --spike-duration 30
```

#### 5. Multiple Modes

Run multiple test modes in sequence:

```bash
# Run the default isolated-latency, individual, and production tests
locust -f $(ml-loadtest-file) \
  --host http://api:8000 \
  --weights-module loadtest.distribution_weights

# Run all test modes
locust -f $(ml-loadtest-file) \
  --host http://api:8000 \
  --weights-module loadtest.distribution_weights \
  --mode ENDPOINT_LATENCY INDIVIDUAL PRODUCTION SPIKE
```

When combined with load modes, all `ENDPOINT_LATENCY` endpoint measurements run first. They bypass adaptive scaling and the inter-load drain delay; selected load modes retain their existing behavior afterward.

### Adaptive Scaling

In `INDIVIDUAL` and `PRODUCTION` modes the tool runs a closed-loop controller that adjusts the number
of concurrent users to keep P99 latency around the target. It uses a **deadband**
defined by `--target-p99-ms` and `--tolerance`:

- **lower bound** = `target × (1 − tolerance)`
- **upper bound** = `target × (1 + tolerance)`

Every `--check-interval` seconds the controller compares the current windowed P99
against the band and takes one of three actions:

| Condition | Action |
| --- | --- |
| P99 **above** upper bound | Remove `--step` users (additive decrease) |
| P99 **below** lower bound | Add `--step` users (additive increase) |
| P99 **inside** the band | Hold steady |

Because it holds steady anywhere inside the band, the controller **converges on a
sustainable operating point** (settling near the lower bound) rather than
oscillating or continuously probing for the absolute breaking point. Both
directions move by the same fixed `--step`, so behavior is symmetric — raise
`--step` for faster convergence, lower it for finer-grained control.

After every user-count change the controller waits `--settle-periods ×
--check-interval` seconds before measuring again, and only computes P99 over
samples gathered since the change, so the transient from the previous load level
never drives a decision.

Each configuration finishes (results are recorded and the run advances to the
next config) on a per-mode condition:

- **INDIVIDUAL**: when the service is overloaded (P99 past the upper bound or
  `--max-users` reached) **or** the run's max duration elapses
  (`--individual-run-max-duration`). Because the controller holds steady once it
  settles inside the band, the duration cap is what bounds this capacity-search run.
- **PRODUCTION**: when `--production-run-duration` elapses — the controller keeps
  adjusting for the full duration and is not cut short by a transient spike.
- **SPIKE**: when `--spike-duration` elapses.

### Key Configuration Options

All standard Locust parameters are available, plus:

- `--target-p99-ms`: Target P99 latency in milliseconds (default: 1000)
- `--max-users`: Maximum concurrent users (default: 32)
- `--min-users`: Minimum concurrent users (default: 1)
- `--mode`: Test modes to run — ENDPOINT_LATENCY, INDIVIDUAL, PRODUCTION, or SPIKE. Space-separated for multiple (default: ENDPOINT_LATENCY INDIVIDUAL PRODUCTION)
- `--step`: Users added/removed per adjustment — additive increase below the lower bound, additive decrease above the upper bound (default: 1). See [Adaptive Scaling](#adaptive-scaling).
- `--tolerance`: Half-width of the deadband around the target, as a fraction (default: 0.1 = ±10%)
- `--check-interval`: Seconds between scaling checks (default: 30)
- `--settle-periods`: Check intervals to wait after a user-count change before adjusting again, so the new load level can stabilise and P99 is measured only over post-change samples (default: 1)
- `--production-run-duration`: Duration in seconds for a PRODUCTION run (default: 1200)
- `--individual-run-max-duration`: Max duration in seconds for an INDIVIDUAL run (default: 600)
- `--endpoint-latency-warmup-successes`: Successful warmup requests discarded per endpoint (default: 1)
- `--endpoint-latency-sample-count`: Successful measured requests retained per endpoint (default: 20)
- `--endpoint-latency-request-timeout-seconds`: HTTP timeout during isolated measurements (default: 10)
- `--endpoint-latency-consecutive-failure-limit`: Consecutive failures before an endpoint is incomplete (default: 3)
- `--output-file`: Output filename prefix (default: "report_loadtest_results")
- `--weights-module`: Python module with custom production_weights (required)
- `--spike-target-rps`: Target RPS for spike mode (default: 100.0)
- `--spike-duration`: Duration in seconds for spike mode (default: 30)

Full list of Locust parameters: https://docs.locust.io/en/stable/configuration.html

### Configuration File

The package includes a `locust.conf` configuration file that provides default settings for load tests. This allows you to avoid repeating common parameters on the command line.

**What is locust.conf?**

A Locust configuration file that sets default values for both standard Locust parameters and custom ml-loadtest parameters.

**Configuration example:**

```ini
; Locust configuration file
host = http://api:8000
headless
only-summary
run-time = 2h
loglevel = INFO
csv = report
html = report.html

; Load test settings
target-p99-ms = 1000
min-users = 1
max-users = 32
step = 1
check-interval = 30
settle-periods = 1
tolerance = 0.1
production-run-duration = 1200
weights-module = loadtest.distribution_weights
```

**How to use it:**

```bash
# Use the bundled config file (from installed package location)
locust -f $(ml-loadtest-file) \
  --config locust.conf

# Override specific settings from config file
locust -f $(ml-loadtest-file) \
  --config locust.conf \
  --max-users 64
```

**Note:** Command-line arguments always override config file settings.

### Analysis

After running tests, analyze results for regressions and get rate limit recommendations:

```bash
# Basic analysis
python -m ml_loadtest.analyze

# Update baseline after confirming results are good
python -m ml_loadtest.analyze --update-baseline

# Custom input/output files
python -m ml_loadtest.analyze \
  --input-file custom_report_loadtest_results.json \
  --baseline-file my_baseline.json \
  --output-file analysis_output.txt
```

The analyzer will:
- Compare current results against baseline
- Detect performance regressions (default 15% relative threshold plus 20ms absolute threshold for latency)
- Compare isolated p50, p95, and p99 when both current and baseline measurements are complete
- Recommend per-endpoint rate limits from measured individual capacity
- Generate detailed reports with statistics

Raw reports retain the existing `metrics` object and add isolated results at the top level:

```json
{
  "metrics": {},
  "endpoint_latency": {
    "/endpoint": {
      "p50_ms": 12.3,
      "p95_ms": 18.4,
      "p99_ms": 19.1,
      "sample_count": 20,
      "target_sample_count": 20,
      "failure_count": 0,
      "complete": true
    }
  }
}
```

Incomplete measurements retain their sample and failure counts but serialize all three percentiles as `null`. Reports and baselines created before this field existed remain readable; their isolated-latency mapping is treated as empty. If a load mode was not selected, its analyzer and Notion values are shown as unavailable rather than as zero.

### Notion Integration

Sync test results to Notion for tracking:

```bash
# Set Notion credentials (environment variables)
export NOTION_TOKEN="notion-integration-token"
export NOTION_TEST_RESULTS_DATABASE_ID="test-results-database-id"
export NOTION_ENDPOINT_DATABASE_ID="endpoint-database-id"

python -m ml_loadtest.notion_sync "service-name" "v1.0.0" \
    --report-file report_loadtest_results.json \
    --baseline-file baseline.json
```

## CI: Load-test Helm values

The reusable workflow `.github/workflows/run-load-test.yml` deploys the service to the
staging cluster through ArgoCD before driving load at it. By default the Helm values for
that deployment come from the shared **K8S-main** repo, at
`<helm_app_path>/env/values-loadtest.yaml` — so tuning replicas, node selectors, CPU
pinning or HPA/rate-limit toggles means opening a PR against a repo owned by another team.

The `values_file` input lets a service own its load-test topology instead:

```yaml
  run_load_test:
    uses: IncodeTechnologies/ml-load-testing-tool/.github/workflows/run-load-test.yml@v1.12.2
    with:
      helm_app_path: my-service-modular
      api_image_tag: ${{ needs.build_images.outputs.api_image_tag }}
      gpu_image_tag: ${{ needs.build_images.outputs.gpu_image_tag }}
      test_image_tag: ${{ needs.build_images.outputs.test_image_tag }}
      values_file: tests/loadtest/values.yaml   # path within THIS repo
    secrets:
      ARGOCD_AUTH_TOKEN: ${{ secrets.ARGOCD_AUTH_TOKEN }}
```

The path is resolved in the **calling** repo, at the ref that triggered the workflow. The
file is passed to `argocd app create --values-literal-file`, which inlines its contents
into the ArgoCD Application's `spec.source.helm.values`.

Leaving `values_file` empty keeps the current behaviour exactly: K8S-main's
`env/values-loadtest.yaml` is used.

### It is a replacement, not an overlay

When `values_file` is set, `--values ./env/values-loadtest.yaml` is **not** passed at all.
Anything your file omits falls through to the chart's own `values.yaml` defaults, which are
**production-shaped**: HPA enabled on both API and Triton, rate limits enabled, PDBs
enabled, no zone pinning, and the *production* image names. A thin file is not a safe file.

Keys that are required, because the chart default is empty or wrong for a load test:

| Key | Why |
| --- | --- |
| the whole `loadtest:` block | Chart default is `{ }`. Omit it and no load-test Job is created, and the workflow blocks on `argocd app wait` until `loadtest_timeout_seconds` (default 2h) expires. |
| `loadtest.bucket` | Must stay `ml-load-test-artifacts-incode-staging` — the `process_results` job downloads results from that bucket, hard-coded, at prefix `<helm_app_path>/<run_id>`. |
| `generalConfig.envName: loadtest` | Chart default is `{ }`. The Grafana metrics and Loki log links filter on `var-environment=k8s-loadtest`. |
| `api.image.name`, `triton.image.name`, `loadtest.image.name` | Chart defaults point at the *production* image repo. The workflow pins only the image *tags*, so omitting these applies your `-dev` tags to the production repo name and the pods land in `ImagePullBackOff`. |

Strongly recommended, or the run is not a controlled measurement:
`api.hpa.enable: false`, `triton.servers.<server>.hpa.enable: false`,
`api.rateLimits.enable: false`, `api.pdbEnabled: false`, `triton.pdbEnabled: false`, and
zone/instance-family `nodeSelector`s so every pod lands in one AZ on a known instance type.

### Keys the workflow always pins

These go through `--helm-set`, which in ArgoCD outranks `helm.values`. Setting them in your
values file has **no effect** — leave them out:

| Key | Source |
| --- | --- |
| `nameOverride` | `ml-load-test-<run_id>` |
| `api.image.tag` | `api_image_tag` input |
| `triton.image.tag` | `gpu_image_tag` input |
| `loadtest.image.tag` | `test_image_tag` input |
| `loadtest.runId` | resolved run ID |
| `loadtest.timeoutSeconds` | `loadtest_timeout_seconds` input |
| `loadtest.analyzeExtraArgs`, `loadtest.locustExtraArgs` | built by the workflow |
| `api.cpuResource`, `api.enableCpuLimit` | **only when the `cpu_resource` input is non-empty** |

That last row is a trap worth spelling out: if your caller workflow passes a non-empty
`cpu_resource` on every run (e.g. `${{ inputs.cpu_resource || '1500m' }}`), it will silently
override `api.cpuResource` from your values file. Default that input to `''` so it stays an
opt-in, dispatch-time override. Note that `cpu_resource` also passes `--cpu-resource` to the
analyzer so the report can recommend an HPA CPU trigger — a values file cannot do that part.

### Starting template

Copy this and adjust; it mirrors the K8S-main `values-loadtest.yaml` it replaces.

```yaml
api:
  pdbEnabled: false
  cpuResource: 1500m
  enableCpuLimit: true
  replicas: 1
  image:
    # Tag is pinned by the workflow — only the name matters here.
    name: docker.io/incodetech/my-service-dev
  scheduling:
    nodeSelector:
      karpenter.k8s.aws/instance-family: "m7a"
      topology.kubernetes.io/zone: us-east-1a
  rateLimits:
    enable: false
  hpa:
    enable: false

triton:
  pdbEnabled: false
  image:
    name: docker.io/incodetech/my-service-dev
  servers:
    modelverse:
      replicas: 1
      scheduling:
        nodeSelector:
          gpu_needs: "gpu_g6"
          topology.kubernetes.io/zone: us-east-1a
      hpa:
        enable: false

loadtest:
  enabled: true
  image:
    name: docker.io/incodetech/my-service-dev
  scheduling:
    nodeSelector:
      topology.kubernetes.io/zone: us-east-1a
  serviceAccount: ml-load-test
  # Must match the bucket process_results downloads from.
  bucket: ml-load-test-artifacts-incode-staging
  # Set by the workflow — listed only so the shape is visible.
  runId: ''
  timeoutSeconds: 7200
  analyzeExtraArgs: ''
  locustExtraArgs: ''

generalConfig:
  envName: loadtest
```

Because the file now lives in your repo, it also has to track K8S-main chart changes that
`env/values-loadtest.yaml` used to absorb centrally. That is the trade-off for owning it.

## Extending with Custom TaskSets

### Creating Custom TaskSets

Each TaskSet must implement this interface:

```python
from locust import TaskSet

class MyCustomTaskSet(TaskSet):
    # Required: endpoint identifier
    endpoint = "/my-endpoint"

    # Required: test implementation
    def test_endpoint(self) -> None:
        with self.client.post(
            self.endpoint,
            json={"data": "example"},
            name=self.endpoint,
            catch_response=True,
        ) as response:
            if response.status_code == 200:
                response.success()
            else:
                response.failure(f"Failed with {response.status_code}")
```

### Using Custom TaskSets

Create a weights module (e.g., `distribution_weights.py`):

```python
from my_tasks import TaskSet1, TaskSet2, TaskSet3

production_weights = {
    TaskSet1: 50,  # 50% of requests
    TaskSet2: 30,  # 30% of requests
    TaskSet3: 20,  # 20% of requests
}
```

Run with custom weights:

```bash
locust -f $(ml-loadtest-file) \
  --host http://api:8000 \
  --weights-module distribution_weights
```

## Architecture

### Core Components

1. **locustfile.py** - Test orchestration with adaptive scaling
   - `EndpointCapacityExplorer`: Manages test modes and adaptive scaling
   - `LoadTestHttpUser`: Executes weighted endpoint tasks
   - Daemon thread monitors P99 and adjusts user count dynamically

2. **analyze.py** - Post-test analysis
   - `LoadTestAnalyzer`: Regression detection and rate limit calculation
   - Compares against baselines (10% regression threshold)
   - Recommends safe limits (70% of measured capacity by default)

3. **distribution_weights.py** - Production traffic weights
   - Example weight configuration for bundled TaskSets
   - Template for custom weight modules

4. **notion_sync.py** - Notion integration
   - Syncs test results to Notion databases
   - Tracks performance metrics over time

### Data Flow

1. Locust users send requests to target endpoints
2. Load-mode response times are captured in circular buffers; isolated latency uses its fixed 20-sample set
3. Daemon thread checks P99 every `--check-interval` seconds
4. User count adjusted by `--step` using the deadband controller (increase below the lower bound, decrease above the upper bound, hold inside the band — see [Adaptive Scaling](#adaptive-scaling))
5. After test completion, JSON/TXT reports saved
6. Analyzer loads reports for regression detection and recommendations

## Development

### Running Tests

```bash
make test
```

### Linting and Formatting

```bash
make type-check
make format
make lint
```

