The Background - An Agentic UI for Postgres AI
These days everybody who is building a SaaS product wants to add an Agentic UI to it.
So rather than bothering to learn the product's GUI new users can just tell it to do stuff.
At EDB we are no different. A customer may want to create a new Postgres cluster or set up an existing one with Analytics or create and populate a vector store for RAG. New users should just be able to ask the Chatbot UI and it will come back with clarifying questions and do it for you.
I was involved in building this Chatbot for Postgres AI and was responsible for the core help agent.
The help agent holds all the knowledge about these components and the full docs for the last few versions of Postgres. The control plane and data plane documentation etc. It acts as the knowledge base for appropriate context for any Chatbot interaction. The bot will first consult the documentation and then the related skills that wrap the tools to execute a task, whether that is migrating a database, or getting performance telemetry from a Postgres cluster in Australia.
OOM Killed
It was a week or so before GA release and everything was working fine. I had created a Python framework with a full set of eval test suites for all the agents including the help agent, see this post on testing Agentic systems. This has helped iron out a few performance bugs and ensured final tweaks had not caused any regressions in. Meanwhile the docs team had done a full rewrite of all its documentation corpus. Their changes finally came in. The documentation for the help agent had doubled in size.
I ran the skill that crawls and aggregates all the documents. It packs them up as markdown and text indexes them so that they can be locally stored and searched without any overhead of web retrieval and LLM calls. The core knowledge base needs to be local and fast so that the whole trajectory of LLM calls and tool steps that are required to complete a task have an acceptable total execution time.
From perfectly functional the agent proxy with knowledge base Golang micro-service became a completely dead Gopher 😢 OOMKill, OOMkill whenever the Chatbot UI was queried.
I did the standard thing and doubled the pod memory, then doubled it again. Success the Gopher lives again. But throwing 4 times the memory resource at it was hardly an ideal solution. Not only that latency was 49% worse when the doc eval suite was run on an AKS deployment with before and after corpus increase runs. Surely it shouldn't require such a huge increase in memory, what if the docs double again in future? The challenge was on, could I improve the code's Go memory usage significantly to reduce the memory cost, ideally removing the latency increase too?
Stuffing double the docs into the same memory footprint
Try the obvious
First I tried aggressively trimming the LLM context to save memory (and tokens). It looked like an obvious memory optimization, but it made the evaluation 7.5x slower (84 mins vs 11 mins) and dropped TaskCompletion from 0.86 to 0.74 because the agent lost disambiguating signals and looped. I rejected the change based on complete-task evidence.
Roll out the AI (parallel Claude, Codex or OpenCode sessions)
Clearly guessing what might save memory was no way to fix this problem.
I needed empirical data, an experimental process and a list of potential fixes that I could try out with full Golang pprof profiling to determine how effective they were.
Instrumenting the code on a remote deployment and returning it via http using net/http/pprof is the option that caters for profiling real kubernetes pod deployments. However I did not need to do that.
All I needed to do was to compare the memory usage for running the code via a test harness locally.
Since I was trying to reduce the relative memory use against the control sample data provided by running the same test against the current code with the doubled doc corpus.
Once I had likely candidate fixes that showed significant memory savings I could confirm they were good by deploying a fix image for real and running the eval tests to check latency was good and OOMkills were eliminated. I already had skills written to create a test deployment for various cloud providers with a PR version of the code and so I could add to them the wrapper instructions to run the doc test eval suite against them.
The main specific temporary skill to sort out was the context file that ran the local pprof and analysed results then scored each fix accordingly. Having come up with a set of 12 possible memory fixes.
I could then run 3 or 4 parallel AI sessions to test and validate the good fixes. So the process became:
Stage 1 - find fix candidates
create code fix > local pprof microbench > pass if it cuts memory over 20%
Stage 2 - candidate validation
create fix PR > autobuild image > rollout to full test deploy > eval doc suite A/B test > push to Langfuse > analyse telemetry
The process was rather tiring dealing with switching between the various multi agent runs and keeping everything on track, with some iterations of the agentic process and reruns required. The bottleneck was that developers are only allowed a few personal ephemeral test clusters by default. The need to restrict development budgets means rolling out 12 clusters is way beyond standard quota. Similarly the amount of Claude use burned through over a weeks worth of my token budget. Productivity is expensive!
But it enabled doing the equivalent of a weeks work in a day or so. Easily completing the discovery of 4 fixes, 1 for latency and 3 for memory usage reducton.
The rewrite stacked up all 4 fixes and after full test validation, could be merged for meeting the GA release deadline.
The Winning Fixes
Latency: Profiling showed phrase boosting was re-tokenizing document bodies on every query. I precomputed normalized text at index-build time, reducing multi-term BM25 p95 latency from 3.9 seconds to 507 ms.
Memory: pprof identified three heap-heavy structures. I replaced map[string]bool token sets with sorted []uint64 arrays and binary search, and I stopped retaining full body strings after JSON unmarshaling, applying expensive phrase reranking only to the top-K results.
Result:
1. HeapInuse dropped from 276 MiB to 60 MiB - a 78% reduction of the heap slice!
The other half of memory is fixed based on baseline Go use rather than the application.
2. The doubled corpus stabilised perfectly at the original 512 MiB limit.
3. Top-3 retrieval remained at 96.8%, and the agent's end-to-end task completion matched the baseline.
Takeaway
The biggest memory saving was back to basics - the Golang data structures used.
Using an array of integers via converting search to binary the index of all the documents can occupy a contiguous slice of memory with each item always only 8 bytes. Where as using a map to store the index keys as text will take 10 times as much memory, since a map is actually a complex hash table with metadata and pointers. Maps also auto grow but do not auto shrink with item deletion, so they suck up memory and keep hold of it.
The lesson learned was don't assume any code change will retain performance. Also don't assume what should fix a performance degradation. Empirical evidence is an absolute requirement.
Anything that may effect memory or latency should be thoroughly profiled and tested.
Performance telemetry should always inform code rewrites and retaining it and enabling A/B evaluation of it is required to avoid changes that can produce performance regressions - sometimes so bad they are critical OOM killers.
No comments:
Post a Comment