TL;DR Improved the p95 latency of a machine learning model endpoint from ~112ms to ~8ms by switching from Flask to FastAPI, using Kubernetes for deployment, restructuring the codebase to OOP, and adjusting scaling rules. Load testing and profiling helped identify bottlenecks, and monitoring was done via DataDog, enhancing both performance and reliability.
Excerpt from analysis
When I took over this project, I was faced with a significant performance issue. The machine learning model endpoint wasn’t meeting its Service Level Agreement (SLA) of a 95th percentile (p95) latency of 100ms. In fact, the p95 latency was hovering around 112ms, which was unacceptable for our client-facing applications consuming the endpoint. I set what I thought was an ambitious yet achievable goal to reduce this latency to below 50ms, which would comfortably meet our SLA requirements and provide a buffer for future growth.
Overhauling the technology stack
My first step was to modernize our technology stack. I migrated the endpoint from Flask to FastAPI, which immediately yielded latency improvements. However, I knew this was just the beginning. I implemented Redis caching to further reduce response times, especially for frequently requested data.
One of the most impactful changes I made was to the data retrieval process. Instead of querying our Redshift data warehouse directly for each request, I set up a system to create daily SQLite dumps of the relevant data. This dramatically cut down on query times. Had DuckDB been available at the time, I would have opted for that instead of SQLite due to its superior performance for analytical queries.
In collaboration with our DevOps team, we transitioned the deployment from a single EC2 instance to a Kubernetes cluster. This move allowed for more flexible horizontal scaling, which was crucial for handling varying loads efficiently.
Fine-tuning performance and monitoring
With the new infrastructure in place, I turned my attention to fine-tuning the application itself:
- I restructured the codebase using object-oriented programming principles to improve maintainability and performance.
- We adjusted our scaling rules to focus on memory usage rather than the typical CPU or I/O metrics, as this was more relevant for our ML model’s performance.
- I implemented custom tracing via DataDog, allowing us to monitor performance in real-time and identify bottlenecks.
Through this monitoring, I discovered two significant bottlenecks:
- JSON serialization was slowing us down, so I created a custom JSON response class using the
ujsonmodule, which was substantially faster than the defaultjson. - Pydantic validation, while useful for external APIs, was unnecessary overhead for our internal service calls. I removed this validation, further reducing latency.
Results and reflections
| Latency percentile | Before | After |
|---|---|---|
| p50 | 41.4ms | 2.2ms |
| p80 | 68.3ms | 4.3ms |
| p90 | 90.0ms | 6.2ms |
| p95 | 111.8ms | 8.4ms |
| p99 | 171.6ms | 13.8ms |
The culmination of these efforts led to results that far exceeded my initial expectations. Not only did we meet our goal of reducing p95 latency to below 50ms, we dramatically reduced it from 111.8ms to a mere 8.4ms. This improvement wasn’t just at the p95 level; we saw significant gains across all percentiles, from p50 to p99.
Looking back, I realize the importance of a holistic approach to performance optimization. It wasn’t just about changing one thing; it was about examining and improving every aspect of our system, from the database queries to the deployment infrastructure to the application code itself.
This project taught me the value of continuous monitoring and iterative improvement. By staying attuned to our system’s performance and being willing to question every aspect of its architecture, we were able to achieve results that far surpassed our initial goals. What started as an attempt to meet our SLA turned into a showcase of how far we could push our system’s performance, setting a new standard for our ML model endpoints.