Performance Testing ML Serving APIs With Locust

ByJordan Brown

Machine learning is becoming increasingly integral to modern software applications, with many systems now relying on ML inference APIs to power intelligent features and automate key processes. However, deploying ML models efficiently at scale introduces significant engineering challenges. Chief among these is ensuring the performance and reliability of the APIs that serve model predictions.

Performance issues with ML inference APIs can lead to slow response times, timeouts, and service interruptions that degrade the user experience and even render dependent systems unusable. Rigorous performance testing is essential to avoid these problems. By accurately simulating production workloads, measuring system behavior, and identifying bottlenecks, we can gain confidence that our ML APIs will meet performance requirements before releasing to end users.

In this in-depth guide, we‘ll explore how to effectively load test machine learning APIs using Locust, an open-source performance testing tool. We‘ll cover the key concepts, walk through a real-world example, and share best practices to help you ensure your ML serving APIs are production-ready.

Understanding ML Serving APIs

Before diving into performance testing, let‘s clarify what we mean by "ML serving APIs." In essence, these are web services that host trained ML models and expose endpoints for requesting predictions on new data. The API acts as an abstraction layer, allowing other systems to integrate ML capabilities without concerning themselves with the underlying model implementations.

There are many common use cases for serving models via APIs, such as:

  • Fraud detection systems that need to score transactions in real-time
  • Recommendation engines that surface personalized content suggestions
  • Computer vision APIs that process user-uploaded images and videos
  • Chatbots and virtual assistants that engage in conversations
  • Predictive maintenance solutions that estimate equipment failure risk

The benefit of this approach is it makes ML services more accessible, reusable, and scalable. However, it also means the APIs must be treated as critical infrastructure. Poor performance can have a compounding negative impact, since the effects are felt not just by the API itself but by all its dependent systems.

Key considerations when building ML serving APIs include:

  • Choosing an appropriate web framework and server to host the API
  • Packaging trained models for efficient distribution and deployment
  • Implementing request validation, authentication, and security safeguards
  • Instrumenting monitoring for performance and model drift
  • Architecting infrastructure for high availability and scalability

While the training and accuracy of the underlying models is important, the APIs themselves deserve equal attention to ensure reliable operation. This makes techniques like load testing invaluable.

Locust: A Powerful Tool for API Load Testing

Locust is a popular open-source load testing tool designed to help developers assess the performance of their APIs and web applications. Its intuitive, code-based approach makes it easy to define complex testing scenarios and user behaviors. Let‘s take a closer look at how it works and what makes it effective.

Locust test scripts are written in plain Python. The basic building blocks are:

  • TaskSets – Classes that define a collection of related tasks (requests) that a simulated user will perform. This allows modeling realistic usage patterns.

  • Tasks – The actual HTTP requests sent to your API, such as GET, POST, etc. Each is defined as a method and decorated with @task.

  • Users – Locust‘s representation of concurrent users interacting with your system. They are subclasses of HttpUser, assigned TaskSets that they execute repeatedly.

Here‘s a simple example of a Locust test script:

from locust import HttpUser, TaskSet, task, between

class UserBehavior(TaskSet):
    @task(1)
    def index(self):
        self.client.get("/")

    @task(2)
    def predict(self):
        payload = {...}
        headers = {‘Content-Type‘: ‘application/json‘}
        self.client.post("/predict", json=payload, headers=headers)

class WebsiteUser(HttpUser):
    tasks = {UserBehavior}  
    wait_time = between(1, 5)

This script defines a TaskSet called UserBehavior with two tasks – one sending a GET request to the root / endpoint, and the other a POST to /predict with a JSON payload. We also define a WebsiteUser that runs the UserBehavior tasks repeatedly, waiting between 1 and 5 seconds after each task.

When you execute a Locust script, you specify a target number of concurrent users and a spawn rate at which they are added. Locust then runs the script, collects metrics, and provides a web UI to monitor the results in real-time. It records key statistics like the number of requests, average response times, and failure rates.

The architecture is fully distributed, so you can scale out tests by running worker instances across multiple machines. Locust consolidates metrics from all workers, giving you an aggregated view of performance.

Some of the advantages of Locust include:

  • Defining tests in code, rather than clunky UIs, allows versioning and reproducibility
  • Extensibility and flexibility to model complex scenarios
  • Ability to gradually ramp up load to assess scalability limits
  • Choice of running locally or distributed for large scale tests
  • Being open-source and well-documented enables customization

With an understanding of Locust‘s capabilities, let‘s walk through applying it to load test an ML serving API.

Testing an ML Serving API with Locust

To demonstrate Locust in action, we‘ll develop a test plan for a sample ML inference API. Imagine we have a deep learning model trained to estimate the value of a house based on property details like the location, size, number of bedrooms, etc. We‘ve packaged this model into an API using a framework like FastAPI and now need to verify it performs acceptably under load.

Here are the key steps:

1. Define performance requirements

First clarify the API‘s expected usage patterns and load:

  • How many inference requests per second must be supported?
  • What is the maximum acceptable response latency?
  • How many concurrent users are expected?

Let‘s assume our API needs to handle up to 100 requests per second with an average latency under 200ms. We expect around 1000 concurrent users at peak times.

2. Create a Locust test plan

Next outline the structure of the load test in a Locust script:

from locust import HttpUser, TaskSet, task, between

class UserBehavior(TaskSet):
    @task
    def get_root(self):
        self.client.get("/")

    @task
    def predict_home_value(self):
        payload = {
            ‘address‘: ‘123 Main St‘,
            ‘bedrooms‘: 3,
            ‘bathrooms‘: 2, 
            ‘square_feet‘: 2000
        }
        headers = {‘Content-Type‘: ‘application/json‘}
        self.client.post("/predict", json=payload, headers=headers)

class MLApiUser(HttpUser):
    tasks = [UserBehavior]
    wait_time = between(1, 5)

Here we define two tasks – one hitting the API root endpoint, and the other sending a request to the /predict endpoint with a mock payload representing a house‘s details. The MLApiUser is configured to send requests repeatedly with a random wait between 1 and 5 seconds.

3. Configure test execution

With the test plan defined, configure the execution settings:

  • Target host address (e.g. http://api.example.com)
  • Number of simulated users (e.g. 1000)
  • Spawn rate (e.g. 10 users/second)

Specify these via command-line arguments to the Locust CLI:

locust -f ai_api_test.py --host=http://api.example.com -u 1000 -r 10

This configures Locust to use the ai_api_test.py script, target http://api.example.com, simulate 1000 users, and add 10 new users every second.

4. Run the test and analyze results

Launch the test and monitor its progress through the Locust web UI. Keep an eye on key metrics:

  • Requests per second: Are we reaching the target throughput of 100 rps?
  • Response times: Is the average latency staying under 200ms? Watch the 95th percentile to spot outliers.
  • Failure rate: Are any requests generating errors? Investigate 4xx and 5xx responses.

Locust‘s charts help visualize behavior over time. Look out for any anomalies like latency increasing as concurrency ramps up, which could indicate a resource bottleneck.

5. Optimize based on results

Analyze the test results to identify areas for improvement. Some common optimization targets:

  • Increasing hardware resources (CPU, RAM, etc.) on the API servers
  • Adjusting auto-scaling policies to accommodate traffic bursts
  • Enabling caching of frequently accessed data
  • Optimizing database queries and indexes
  • Tweaking model inference batch sizes and parallelization

After making changes, re-run the load test to verify their impact. Consider automating this into a continuous testing pipeline.

Best Practices for Testing ML Serving APIs

We‘ve walked through a basic example of load testing an ML API with Locust, but there are many additional factors to consider as you develop a comprehensive testing strategy. Some best practices include:

  • Test early and often. Performance testing should be an integral part of the ML lifecycle, not an afterthought. Continuously test as you build to catch issues early.

  • Simulate realistic workloads. Craft test scenarios that closely mimic production usage patterns. Consider factors like the mix of inference requests, authentication patterns, and data payload shapes.

  • Verify API semantics. Load tests are a good opportunity to sanity check API contracts. Ensure you‘re sending valid request payloads and receiving expected response schemas.

  • Test edge cases. Inject outliers and boundary values into your test data. Verify the API gracefully handles large payloads, malformed inputs, and unexpected data distributions.

  • Monitor resource utilization. During load tests, keep an eye on infrastructure metrics like CPU usage, memory consumption, and network I/O. Bottlenecks here can help explain performance issues.

  • Perform distributed testing. To accurately gauge limits, run large-scale tests using distributed worker instances. This better simulates real-world conditions and identifies scaling bottlenecks.

  • Integrate with CI/CD. Automate load tests as part of your build and deployment pipelines. Configure performance benchmarks as pass/fail criteria to enforce regressions.

  • Collaborate with stakeholders. Engage with product owners, data scientists, and system architects to ensure your test scenarios align with business requirements and infrastructure realities.

Conclusion

Performance is a critical concern when deploying machine learning models at scale. Users expect responsive, reliable experiences and even brief slowdowns or outages can significantly impact engagement and revenue. Load testing is an invaluable tool to proactively identify and address performance issues before they impact end users.

Locust‘s versatile and code-driven approach makes it a compelling choice for load testing ML inference APIs. By simulating complex scenarios and modeling realistic traffic patterns, it allows you to verify your APIs can handle production workloads. Following testing best practices and continuously assessing performance helps build more robust, scalable ML-powered applications.

While we covered the basics of using Locust here, it has many more advanced capabilities worth exploring. Techniques like distributed testing, custom client behavior, and integration with monitoring platforms enable even more sophisticated and informative load tests. As ML infrastructures evolve, Locust is a powerful addition to any engineering team‘s toolkit.

Further Reading:

How useful was this post?

Click on a star to rate it!

Average rating 0 / 5. Vote count: 0

No votes so far! Be the first to rate this post.

You May Like to Read,

Similar Posts