Skip to content

Optimize embedding processor performance - reuse LlamaContext #776

Description

@malibio

Overview

The embedding processor currently creates a new LlamaContext for each embedding call, causing significant CPU overhead and underutilizing the GPU. During bulk imports (172 docs), CPU usage spikes to 1200%+ while GPU utilization remains low (~10%).

Problem Statement

Current Behavior:

  • Each embedding call creates a new LlamaContext via model.new_context()
  • Metal kernel compilation occurs for every context creation
  • Context is immediately freed after one embedding
  • Results in 1200% CPU during bulk operations with ~10% GPU utilization

Evidence:

  • Log shows repeated ggml_metal_init: loaded kernel_* messages (context recreation)
  • Log shows ggml_metal_free: deallocating after each embedding
  • Activity Monitor shows high CPU, low GPU during embedding processing
  • 172 document import took ~54 minutes with heavy CPU load

Code Location: packages/nlp-engine/src/embedding.rs:207-253

// Current: New context created for EACH embedding
fn generate_embedding_llama(&self, text: &str) -> Result<Vec<f32>> {
    // ... 
    let mut ctx = model.new_context(backend, ctx_params)...; // Created every call
    // ... use ctx ...
    // ctx dropped here, Metal resources freed
}

Proposed Solution

Reuse the LlamaContext across multiple embedding calls instead of recreating it each time:

  1. Store context as Option<LlamaContext> in EmbeddingService
  2. Initialize context lazily on first embedding call
  3. Reuse context for subsequent calls
  4. Only recreate if context parameters change

Alternative Approach:
Implement batch embedding that processes multiple texts in a single context session.

Acceptance Criteria

  • LlamaContext is reused across embedding calls (not recreated each time)
  • GPU utilization increases during embedding processing
  • CPU usage decreases during bulk embedding operations
  • Embedding throughput improves (measure docs/second)
  • No regression in embedding quality or correctness
  • Tests pass (bun run test:all)
  • Code passes bun run quality:fix

Technical Specifications

Reference Files

  • embedding.rs: packages/nlp-engine/src/embedding.rs - Main embedding service
  • embedding_processor.rs: packages/core/src/services/embedding_processor.rs - Background processor
  • embedding_service.rs: packages/core/src/services/embedding_service.rs - Node embedding service

Key Changes

Option A: Persistent Context

pub struct EmbeddingService {
    // ... existing fields ...
    #[cfg(feature = "embedding-service")]
    ctx: Option<LlamaContext<'static>>, // Reusable context
}

Option B: Batch Processing

pub fn generate_batch(&self, texts: Vec<&str>) -> Result<Vec<Vec<f32>>> {
    // Create single context
    let ctx = self.get_or_create_context()?;
    // Process all texts in one session
    texts.iter().map(|t| self.embed_with_context(&ctx, t)).collect()
}

Considerations

  • Context lifetime management with llama.cpp backend
  • Thread safety if context is shared
  • Memory usage of persistent context vs recreation overhead
  • Graceful handling if context becomes invalid

Dependencies

  • None

Related Issues

Non-Goals

  • Changing the embedding model itself
  • Modifying the debounce/queue system
  • Changing the root-aggregate embedding model

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions