Load test AI endpoints with Locust

You've spent days perfecting your AI pipeline, but the moment real traffic hits your LLM endpoint, latency spikes, responses time out, and your users are looking at a loading spinner that never ends. You need to know how your endpoint behaves under pressure before it happens in production. That's where load testing with Locust comes in — this tutorial will walk you through creating a Locust load test for your AI endpoints, from a simple request to realistic, streaming-capable simulations, so you can measure, understand, and fix performance issues with confidence.

The problem this lesson solves

AI endpoints are unlike typical REST APIs. They're stateful, latency-heavy, and often asynchronous — a single request to a large language model can take seconds, and the response might be streamed or delivered via WebSocket. Traditional load testing tools like ab or siege assume quick, synchronous HTTP calls, so they give you misleading results: your endpoint could handle 100 requests per second when the GPU is idle, but collapse under realistic concurrent sessions where each user waits for a multi-second response.

Another pain point is background inference: your AI endpoint might immediately return a 202 Accepted, then process the request in a job queue. In that case, you're testing the web framework's ability to accept requests, not the AI inference itself. Without a proper load test, you're blind to the real bottleneck — GPU saturation, token generation speed, or request queue buildup.

The consequence? You ship an AI feature that works in QA but falls apart in production, and you only find out when your incident alert fires at 2 AM. Locust provides a Python-native, scriptable way to simulate realistic user behavior against AI endpoints, so you can predict and prevent these failures.

Core concept / mental model

Think of load testing an AI endpoint like rehearsing a busy restaurant kitchen.

• Your AI model is the head chef — fast when working alone, but slow when every station is slammed.
• Locust's users are customers placing orders — each one waits for their meal (the response) before placing another.
• The system under test is the whole restaurant — web server, GPU, database, and network.

If you flood the restaurant with customers who don't wait for their food (fire-and-forget requests), you won't see the real choke point. Locust lets you simulate realistic customer behavior: a user sends a request, waits for the response (or streams it), and then sends another. This is what you actually want to measure: how many concurrent users can sit and eat without the kitchen collapsing?

Key definitions before we dive in:

• Users — simulated clients that each run a task (your test logic) in a loop.
• Wait time — how long a simulated user pauses between tasks, mimicking real think time.
• Requests per second (RPS) — the rate of requests your endpoint receives; not the most important metric for AI, but still useful.
• Response time percentiles — e.g., p50, p95, p99. For AI endpoints, p95 is your user experience; p99 is your worst case.
• Failures — any non-200 response or timeout — these should be near zero in a healthy test.
• Spawn rate — how quickly users are added to the test, letting you see how the system degrades gradually.

In the Locust mental model, each user is a tiny Python script that knows how to interact with your endpoint. The framework coordinates hundreds of these scripts, collects metrics, and presents them in a real-time web UI.

How it works step by step

Let's break down the lifecycle of a Locust load test against an AI endpoint:

1. Define your user behavior — a Python class inheriting from HttpUser, with tasks that define what a user does. For AI, this is typically: send a prompt, wait for the response, optionally parse and validate it.
2. Set wait times — how long a user "thinks" between tasks. For AI, this should match your real usage pattern (e.g., 2–5 seconds). Don't set it to zero unless you want a synthetic maximum stress test.
3. Run Locust — from the command line, specifying the number of users to simulate and how quickly to spawn them.
4. Monitor real-time results — Locust serves a web UI at http://localhost:8089 showing RPS, response times, and failures as the test progresses.
5. Analyze output — either from the UI, or exported to CSV for deeper analysis.

The crucial part for AI endpoints is how you write the requests. A simple GET to /health is fine, but to test real inference, you need to:

• Send a realistic JSON payload (e.g., with prompt, max_tokens, temperature).
• Handle streaming responses — catch a StreamingResponse from FastAPI or SSE events.
• Treat 202 Accepted responses properly — if the endpoint is async, you may need to poll a status endpoint as part of the user task.

Pro tip: Always start with a low number of users (like 10) and a slow spawn rate (e.g., 1 user/sec). Watch the response time percentiles. If p95 starts climbing linearly with user count, you've found your breaking point — the maximum concurrency your endpoint can handle.

Hands-on walkthrough

Let's build a practical Locust test for a typical AI endpoint. We'll assume you have an HTTP endpoint like /generate that accepts a JSON body and returns a JSON response (not streaming) for simplicity.

Setup

First, install Locust in your Python environment:

pip install locust

Now create a file named locustfile.py (this is the default filename Locust looks for). Here's a simple test:

from locust import HttpUser, task, between

class AIUser(HttpUser):
    # Wait between 1 and 5 seconds between tasks to mimic real users
    wait_time = between(1, 5)

    @task
    def generate_text(self):
        # Realistic prompt, similar to what your app would send
        payload = {
            "prompt": "Write a short paragraph about the benefits of load testing",
            "max_tokens": 100,
        }
        # The endpoint returns a JSON like {"text": "...", "latency_ms": 123}
        response = self.client.post("/generate", json=payload)
        # Validate the response - a 200 with expected structure
        if response.status_code == 200:
            data = response.json()
            if "text" not in data:
                response.failure("Response missing 'text' field")
        # else: Locust automatically records 4xx/5xx as failures

Run it with:

locust --host http://localhost:8000

Then open http://localhost:8089 in your browser, set Number of users to 50, Spawn rate to 5 per second, and click Start. You'll see real-time charts.

Expected output (web UI): You'll see RPS, average response time, and a failures column. If your endpoint handles it, p95 should stay under 2 seconds; if it can't, p95 will climb and failures will appear.

A realistic test with streaming

Many AI endpoints stream responses token-by-token (e.g., using SSE). Locust can handle that using the catch_response context manager to inspect chunked responses:

from locust import HttpUser, task, between

class StreamingAIUser(HttpUser):
    wait_time = between(2, 7)

    @task
    def stream_completion(self):
        payload = {
            "prompt": "Explain quantum entanglement in simple terms",
            "stream": True,
        }
        # Use catch_response to manually control success/failure
        with self.client.post("/v1/completions", json=payload, catch_response=True) as response:
            if response.status_code == 200:
                # For streaming, response content might be a stream of SSE events.
                # You can read the text incrementally. Here we just read it all.
                content = response.text
                if "data:" in content:
                    response.success()
                else:
                    response.failure("No SSE data received")
            else:
                response.failure(f"Unexpected status: {response.status_code}")

Handling async endpoints (202 + polling)

Some AI services return a job ID and process asynchronously. Your Locust task should simulate the full lifecycle:

from locust import HttpUser, task, between
import time

class AsyncAIUser(HttpUser):
    wait_time = between(1, 3)

    @task
    def async_generation(self):
        # Trigger the job
        response = self.client.post("/jobs", json={"prompt": "Hello"})
        if response.status_code == 202:
            job_id = response.json()["job_id"]
            # Poll until complete or timeout
            for _ in range(30):  # max 30 attempts
                status = self.client.get(f"/jobs/{job_id}")
                if status.status_code == 200:
                    break
                time.sleep(1)  # pragmatic sleep, but Locust has a custom way to wait
            else:
                # If we never got 200, mark as failure
                self.client.get(f"/jobs/{job_id}").failure("Timeout waiting for async job")

Note: In real production, you wouldn't use time.sleep inside a Locust task because it blocks the greennet thread. Instead, use gevent.sleep which yields control. But for simplicity in this lesson, we use time.sleep; just be aware of this for large-scale tests.

Compare options / when to choose what

Choose Locust when: you're a Python developer, you need to simulate complex AI workflows (prompt crafting, streaming consumption, async polling), or you want to extend the test with Python logic.

Choose k6 if your team already writes JavaScript and you need high-level test orchestration scripts.

Choose JMeter if you're in a corporate environment with a non-Python QA team.

For most AI endpoint teams, Locust is the sweet spot — it keeps the tests in the same language as your app, so your AI engineers can write and maintain them without learning a new DSL.

Troubleshooting & edge cases

• Test hangs — no requests are made — Check that your host flag or --host is correct, and that the endpoint is reachable from where Locust runs. The UI shows stats; if RPS is 0, likely wrong host or a firewall.
• Failures show on every request — This is often a 409/429 response (rate limit). Your AI endpoint may have its own rate limiter, which will skew results. Consider disabling rate limits for load test or design the test to respect them.
• Response times are unrealistically high — If you're seeing memory or GPU contention due to your own test clients, use distributed mode (--master and --worker) to spread load generation across machines.
• Streaming responses appear as one blob — If you're testing an SSE endpoint, response.text may be huge; you need to parse chunks. Use catch_response and read line by line to avoid memory spikes.
• Using time.sleep slows down the test — As mentioned, use gevent.sleep(1) instead to avoid blocking the event loop, which can reduce your max RPS artificially.
• Endpoint returns 202 but never completes — Your async polling might time out. Make sure you test your polling logic separately, and increase the number of attempts in the Locust task.
• Latency is high for AI but HTTP-level tests look fine — This indicates the bottleneck is AI inference, not the web server. Monitor GPU utilization via nvidia-smi during the test; if it's maxed, you need to scale the model or use request batching.

What you learned & what's next

You now know how to load test AI endpoints with Locust — from the mental model of simulating realistic users to writing effective test scripts that handle synchronous, streaming, and async endpoints. You've also learned how to choose the right tool (hint: Locust for Python teams) and troubleshoot common pitfalls like rate limits and streaming issues.

With these skills, you can confidently predict how your AI service will perform under load and make informed capacity decisions.

Next step: In the next lesson, you'll learn how to analyze Locust test results and derive actionable insights — like identifying whether your bottleneck is GPU-bound or network-bound, and how to translate that into scaling decisions. You'll take the metrics you collect here and turn them into an optimization strategy.