2015887 - llama.cpp leaks ~500mb for each link preview performed until threads are exhausted

Status: RESOLVED FIXED

(Core :: Machine Learning: On Device, defect, P2)

Description

[Greg Tatum [:gregtatum]]

I got the inference process up to 1.5G of memory usage before link preview refused to generate any more requests. From the performance profile it looks like it's creating threads, and not re-using them. Eventually it seems to exhaust the threads and breaks.

https://share.firefox.dev/3OeYJqz

I'm tentatively marking as S2 since it gets up to a pretty problematic amount of memory. It is quite easy to accidentally activate the link preview feature since it's a long press on a link. Eventually, the process is terminated after some timeouts, so the memory is eventually freed, but during this time a user could have some negative impacts from this.

Comment 1

[Greg Tatum [:gregtatum]]

Here's what claude suggests. No idea if it's accurate, but is a good place to start the investigation:

Root Cause

Every time LinkPreviewModel.generateTextAI() is called, the following happens:

1. New Engine Creation (LinkPreviewModel.sys.mjs:415-480):
   - Line 417-445: Creates a new engine via createEngine()
   - Line 479: Terminates the engine after use
2. Engine Caching Issue (EngineProcess.sys.mjs:133-137 and MLEngineParent.sys.mjs:287-295):
   - Engines are supposed to be reused if pipeline options match
   - The engineId for link-preview is "wllamapreview" (line 174 in EngineProcess.sys.mjs)
   - However, the engine caching only works if called through the same MLEngine instance
3. Native Thread Creation (LlamaBackend.cpp:221-254):
   - In ReinitializeContext(), new threadpools are created:
   • Line 247: mThreadpool.reset(mLib->ggml_threadpool_new(&tpp))
   • Line 235: mThreadpoolBatch.reset(mLib->ggml_threadpool_new(&tppBatch))
     - These create llama.cpp worker threads with callbacks (lines 223-224, 228-231)
4. Worker Thread (LlamaRunner.cpp:437-439):
   - Each generation also creates a new "LlamaWorker" thread via NewNamedThread
   - This thread is properly shut down via AsyncShutdown() (line 391)

The Problem

The issue is that engines are NOT being reused properly. Looking at LinkPreviewModel.sys.mjs, each call to generateTextAI():

• Creates a fresh engine (line 417-445)
• Terminates it immediately after use (line 479)

This means:

• New llama.cpp threadpools are created each time (multiple threads per threadpool)
• The threadpools should be freed when the LlamaBackend is destroyed (via custom deleters in LlamaBackend.h:93-94)
• However, if there are any lingering references to the backend, the threadpools won't be freed

Potential Solutions

1. Reuse the engine instead of creating a new one each time:
   - Cache the engine at the LinkPreviewModel level
   - Only recreate if pipeline options change
2. Ensure proper cleanup of threadpools:
   - Verify that all RefPtr references to LlamaBackend are dropped when engine terminates
   - Check if LlamaStreamSource might be keeping the backend alive

Let me check if there's an engine caching issue:

⏺ Searched for 1 pattern (ctrl+o to expand)

⏺ Based on my analysis, the issue is that LinkPreviewModel.generateTextAI() creates a new engine for every preview, which creates new llama.cpp threadpools each time, and these threads are not being properly
reused.

The Fix

The solution is to cache and reuse the engine at the LinkPreviewModel level. Here's what needs to be changed in browser/components/genai/LinkPreviewModel.sys.mjs:

The problem is in the generateTextAI function (lines 366-481). Currently:

• Line 415-445: Creates a new engine every time
• Line 479: Terminates it after use

Instead, the engine should be:

1. Created once and cached in the LinkPreviewModel object
2. Reused across multiple generateTextAI() calls
3. Only recreated if pipeline options change
4. Properly terminated only when LinkPreviewModel is torn down

Would you like me to implement this fix? The solution would involve:

1. Adding a cached engine field to LinkPreviewModel
2. Checking if an engine exists and its options match before creating a new one
3. Reusing the existing engine when possible
4. Only calling engine.terminate() when the options change or on module teardown

This would eliminate the thread leak by ensuring threadpools are created once and reused, rather than being created and destroyed on every preview.

Comment 2

[Greg Tatum [:gregtatum]]

<s>Bug 2015887 may be related.</s> I apparently link to the same bug here accidentally.

Comment 3

[Valentin Pollet]

To investigate this issue a bit more, I started by adding some logging in native code to observe construction/destruction of:
LlamaRunner, LlamaGenerateTask, LlamaStreamSource and LlamaBackend. The lifetime of native objects and flow of operations during a Link Preview is roughly:

1. Pipeline initialization: LlamaRunner::ctor -> LlamaBackend::ctor
2. Generation stream creation: from LlamaRunner, LlamaStreamSource::ctor and LlamaGenerateTask::ctor.
3. Actual generation: Link Preview JS calls runWithGenerator which consumes LlamaStreamSource from JS.
4. (Optional) Early termination when enough has been generated.
5. Teardown

We'd expect in a healthy Link Preview that each constructor is called once, and each destructor is called once. When I ran Link Previews with my logging enabled, I observed that two destructors were not called: LlamaBackend and LlamaStreamSource. The latter keeps the former alive by strongly referencing it.

The leak is caused by LlamaStreamSource being kept alive after teardown. LlamaStreamSource is the native half of the streaming setup, the other half living in JS.

LlamaStreamSource is likely forming a strong cycle through its stream controller. Nulling it in LlamaStreamSource::ShutdownWorkerThread fixes the leak, but I'm not satisfied because the strong cycle shouldn't be an issue as it is traversed for cycle collection .

Edit: Here's a profile with LlamaStreamSource::mStreamController nulled in ShutdownWorkerThread

Comment 4

[Valentin Pollet]

On a hunch, I tried reproducing the leak with the changes from this revision. It can't reproduce the leak with it applied. The change adds a link to be traversed during CC for LlamaRunner->mInitPromise. I have no clue why this change prevents the leak from happening.

Comment 5

[Joe Bowser]

I can reproduce what looks like a relatively small leak here:

https://share.firefox.dev/3RZqCEX

The total available remains the same but the allocated goes up a about 500K or so each time I run LinkPreview and it doesn't get fully cleaned up. I used bikepacking.com because unlike Wikipedia, it has a lot more graphics and it also doesn't have annoying hover popups.

Comment 6

[Valentin Pollet]

Attachment: Bug 2015887 - Added missing fclose on exit of LlamaRunner::OnMetadataReceived r=jbowser

Comment 7

[Valentin Pollet]

The relatively small leak that the profiler now reports is a bit elusive. I dumped memory before and after 10 link previews or so and I could not find where those 2MB went. When looking at about:memory or about:processes, the 2MB increase per call we see in profiler is not visible. Furthermore, hitting minimize memory in about memory will return the inference process to a healthy baseline which makes me confident in saying this is retention rather than leaking.

I did find a tiny leak by using DMD: LlamaRunner::OnMetadataReceived opens a file descriptor and never closes it.

Comment 8

[Pulsebot]

Pushed by [e-mail]:
https://github.com/mozilla-firefox/firefox/commit/f3c135230eec
https://hg.mozilla.org/integration/autoland/rev/d2a7b6f09794
Added missing fclose on exit of LlamaRunner::OnMetadataReceived r=jbowser,ai-platform-reviewers

Comment 9

[Sandor Molnar[:smolnar]]

https://hg.mozilla.org/mozilla-central/rev/d2a7b6f09794