Optimizing memory usage with AI

1. Optimizing memory usage with AI

Welcome back! So far, we've focused on time complexity: finding CPU bottlenecks, understanding dominant operations, and optimizing performance-critical code paths.

2. From time to space complexity

In real-world systems, however, performance problems are not always about speed. They're often about memory usage, scalability, and how data grows over time. In this video, we'll extend our workflow to space complexity and see how AI can help us detect, analyze, and fix memory bottlenecks in production code.

3. Testing at scale

Up to now, we've been analyzing Atlas using a relatively small sample dataset. That was useful for understanding the application, but memory issues often only appear at scale. To test Atlas at scale, we could run it in one of our real tourism databases. That is closer to what we'd expect in a real deployment.

4. Memory growth warning

After doing this, we notice something concerning: the application's memory usage grows much faster than expected. This suggests that the memory cost may grow faster than linearly with input size, indicating a space-complexity issue rather than a simple configuration problem.

5. AI suggests tracemalloc

This is where AI-assisted profiling becomes useful. Instead of guessing which tools to try, we can ask AI for guidance. From this prompt, the model suggests tracemalloc, Python's built-in memory allocation profiler, and explains that it is well-suited for identifying where memory is allocated and retained.

6. Running tracemalloc

Following the AI-generated instructions, we run Atlas under tracemalloc and capture memory usage snapshots during execution.

7. AI interprets memory traces

Now we have raw profiler output, but interpreting memory traces can be tricky. Once again, we can bring AI into the loop! By reasoning over the profiler output, the model confirms that a large fraction of memory is retained by repeated materialization of intermediate data structures during aggregation.

8. Root cause identified

In practice, issues like this are often rooted in data structure choices and how intermediate results are represented and reused. This gives us confidence that we're dealing with a real space complexity bottleneck, not just noisy measurements.

9. AI proposes fixes

Once the bottleneck is identified, we can ask the AI to reason about possible fixes. The model proposes several options: reducing intermediate allocations, reusing buffers, and introducing caching where recomputation is frequent but results are reusable.

10. Caching trade-offs

At this point, we want to be careful. Caching can improve performance, but it can also increase memory usage if applied incorrectly.

11. Caching trade-offs

So we ask a follow-up question. The AI model reasons about access frequency and confirms that a bounded cache can make peak memory worse for this workload. In this case, the AI suggestion is to reduce intermediate allocations first.

12. AI implementation

Based on this reasoning, we ask the model to propose a concrete implementation. As always, we review the suggested code carefully before applying it. Then we validate with benchmarks before accepting changes.

13. Validation results

Then, we rerun the application under tracemalloc to validate the effectiveness of the code changes. As a result, we observe that peak memory usage drops significantly, while runtime remains stable. This confirms that the proposed optimization is both correct and effective.

14. Benchmarking prompts

Finally, we can make this analysis reproducible by constructing a benchmarking prompt. A benchmarking prompt is a structured way of asking AI to reason from measurements. Instead of vague questions, we could write a more complete prompt that provides the profiler output from before and after a change, specifies which metrics matter, and asks the model to justify its conclusions using numbers. This turns performance optimization into a repeatable, evidence-based process. Benchmarking prompts can help us to systematically test and validate optimization decisions.

15. Let's practice!

Now it is your turn to practice data structures optimization!

Create Your Free Account

By continuing, you accept our Terms of Use, our Privacy Policy and that your data is stored in the USA.