Cut idle embedding memory without changing nomic quality

[TorstenDittmann]

Fixes #7

The embedding process kept full-precision Nomic ONNX sessions resident from boot. Appwrite self-hosted does not set EMBEDDING_POOL_SIZE. Idle RSS on this host was ~870 MiB for a single unused nomic session.

Changes

• Lazy-load ONNX sessions on first /embed; /health no longer pins weights. Load runs in spawn_blocking with a std::sync::Mutex held inside that blocking task, so a cancelled cold /embed cannot start a second pool, and a current-thread runtime gets an EmbedError instead of panicking.
• Idle unload after EMBEDDING_IDLE_UNLOAD_SECS (default 300; 0 disables), then reload the same NomicEmbedTextV15 checkpoint from cache. In-flight requests block unload so a slow embed cannot create a second pool. Unload runs on a blocking thread.
• Pool default is 1. Operators opt into EMBEDDING_POOL_SIZE>1 for parallel HTTP throughput. Extra sessions are RAM-capped on first load and that cap is reused on later reloads.
• Intra-op threads default to CPU count on the single default session. When the pool is larger than one, every session uses nproc / pool_size threads.
• GET /health is liveness (200 while unused). After a failed model load it returns 503 for 30 seconds, then 200 again so a probe can recover; the next /embed retries. Failed loads retry at most every 30 seconds.
• jemalloc is the crate allocator (tests and the binary) so ONNX Runtime's C++ heap is released after unload
• Image/compose defaults cache nomic only. Breaking: {"model":"bge-small"} is 400 unless you add it to EMBEDDING_MODELS. With lazy load that no longer costs idle RSS, only image size.

Quality

nomic / nomic-embed-text still map to full-precision NomicEmbedTextV15 (768-d). No quantized default. Ignored e2e covers MiniLM embed → idle unload → embed (including concurrent unload) and Nomic 768-d cosine ranking.

RSS on this host (4 CPUs, ~16 GiB, full-precision nomic)

State	master	this PR
/health (model never used)	~870 MiB	13 MiB
After one /embed	~885 MiB	587 MiB (one nomic session)
After idle unload	n/a (always resident)	83 MiB

MiniLM on master was 138 MiB at pool=1 vs 435 MiB at nproc=4 while loaded. That multiplier is now opt-in via EMBEDDING_POOL_SIZE.

  

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge; the previously reported cancellation race is fixed and no new actionable failure remains.

Summary

This PR changes embedding models from eagerly resident sessions to lazily loaded, idle-unloaded pools while preserving model quality and container-aware memory controls.

• Loads model sessions on first use and serializes initialization independently of request cancellation.
• Unloads inactive sessions while protecting in-flight inference from duplicate pool creation.
• Defaults to one session and divides intra-op threads across explicitly larger pools.
• Makes health report recent load failures without permanently failing liveness.
• Adds observable lifecycle and Nomic quality coverage and updates container defaults to cache Nomic only.

_{Reviews (10) · Last reviewed commit: "fix: keep model-load ownership if /embed..."}

[cursor]

Review

The memory work is sound and the headline numbers reproduce on my box: boot RSS 13 MiB with /health not touching weights, 585 MiB with one nomic session loaded, and 88–130 MiB after idle unload. I ran ten consecutive load→unload cycles and saw no monotonic RSS growth, so there's no leak across reload cycles. cargo fmt, cargo clippy -D warnings, and all six e2e tests (including the #[ignore]d MiniLM unload/reload and nomic cosine tests) pass on stable 1.98.

The problem is that moving model loading onto the request path put a large amount of blocking work inside async code, and it took the container's only liveness signal with it. I'd fix those two things before merging.

1. Model loading blocks tokio worker threads

embed is async, and it calls acquire_instance → ensure_loaded → load_model synchronously. That path does file IO, an HTTP download, ORT session construction, and a warmup inference, and it serializes concurrent callers on a std::sync::Mutex:

// src/embedding.rs
fn acquire_instance(&self, loaded: &LoadedModel) -> Result<Arc<Mutex<TextEmbedding>>, String> {
    for _ in 0..2 {
        self.ensure_loaded(loaded)?;   // <- synchronous download + ORT init + warmup

Nothing goes through spawn_blocking or block_in_place. Measured against the release binary on a 4-CPU host (so 4 tokio workers), firing 4 concurrent cold /embed requests:

Scenario	/health latency
Idle baseline	0.9 ms
4 concurrent cold embeds, warm model cache	683 ms
4 concurrent cold embeds, empty EMBEDDING_CACHE_DIR	5.2 s

One worker is inside the load; the other three park on the load mutex, so there is no worker left to answer /health. My VM has a fast link, so the 550 MB nomic download finished in ~6 s — on a typical self-hosted connection the whole service, /health included, would be unresponsive for the entire download. Idle unload makes this recurring rather than a one-time startup cost: every 300 s idle window is followed by another blocking reload (~1 s warm for nomic here).

Wrapping load_model in tokio::task::block_in_place works with &self and is the smallest fix; making the load-serialization lock a tokio::sync::Mutex would keep the waiters off the workers too.

unload_idle has the same shape at lower stakes — it drops the session (freeing ~500 MiB and joining ORT thread pools) on a worker thread from the background ticker.

2. Fail-fast is gone, and /health no longer reflects it

EmbeddingClient::new can't fail anymore, so load errors surface only per-request. With a read-only EMBEDDING_CACHE_DIR the process stays up and:

health: 200
embed 1 -> 500  {"error":"Failed to initialize embedding model: Failed to retrieve onnx/model.onnx"}
embed 2 -> 500  (same)
embed 3 -> 500  (same)
health after failures: 200

On master this was a nonzero exit and a visible crash-loop under restart: unless-stopped. Now the container looks healthy forever while being completely non-functional, which is the worst outcome for self-hosted users who only have /health. Each failed request also retries the full download with no backoff. Either keep an eager preload on by default (EMBEDDING_PRELOAD, opt out for fast boots) or have /health report readiness once a model has loaded at least once.

Worth addressing

Pool sizing is now recomputed on every reload from a noisy measurement. The mem_before/mem_after delta uses host available_memory, and it now runs while other models may be serving traffic and right after an unload when jemalloc may not have returned pages yet. With EMBEDDING_POOL_SIZE>1 the effective pool size can differ between reload cycles for no visible reason. Caching the cap on LoadedModel, or logging at warn when it changes, would make that debuggable.

in_flight_embed_survives_idle_unload may not test what it claims. With idle_unload_secs = 1 and unix_now() truncated to whole seconds, should_unload can stay false for the entire ~100 ms unload loop, in which case the loop never reaches the in_flight recheck it exists to cover. Whether it does depends on sub-second alignment of when embed started. Injecting the clock, or asserting a counter that at least one attempt reached the recheck, would make it deterministic. Same root cause: second truncation makes the real idle window idle_unload_secs - 1 .. idle_unload_secs, which is immaterial at 300 but not at small values, and Instant would also sidestep wall-clock jumps.

The throughput side of the tradeoff isn't quantified. pool_size default 1 means concurrent /embed calls now queue on a single mutex where master gave self-hosted nproc sessions. The RSS table is thorough; a p50/p99 comparison at ~4 concurrent requests would make the tradeoff reviewable on the same footing. Relatedly, the per-chunk spawn_blocking fan-out buys nothing at pool_size=1 — every chunk contends on the same session — so it's now just task overhead, worth a comment or a single-session fast path.

Dropping bge-small from .env and compose is user-visible breaking behavior: {"model":"bge-small"} starts returning 400 after a rebuild. That's fine as a decision, but with lazy loading it no longer costs idle RSS (only image size), and it should be called out as breaking in the PR body.

Nits

The #[global_allocator] lives in src/main.rs, so the library and all integration tests run under the system allocator. I hit this directly: the same ten embed→unload cycles leave RSS at 609–807 MiB inside an integration test versus 88–130 MiB against the binary. Anyone reproducing your numbers via cargo test will conclude idle unload doesn't work — worth moving to src/lib.rs or noting in the README. Enabling jemalloc's background_threads would also stop purging from depending on later allocation activity.

Smaller things:

• The tikv-jemalloc-sys direct dependency and use tikv_jemalloc_sys as _; are redundant — override_allocator_on_supported_platforms on tikv-jemallocator already enables tikv-jemalloc-sys/override_allocator_on_supported_platforms, and the #[global_allocator] static keeps the crate linked.
• cap_pool_from_memory's None => 1 arm is unreachable because per_instance_bytes == 0 is handled above, so checked_div can't return None.
• idle_unload_secs() is never called (main.rs reads config.idle_unload_secs before moving the config).
• load_model clones the whole EmbeddingConfig twice just to override intra_threads; passing the thread count into init_model would read better.
• When the RAM cap shrinks the pool, session_intra != probe_intra so the probe session is discarded and rebuilt/rewarmed. That's correct, but a comment noting the deliberate double build would help the next reader — and note the probe measures a session with probe_intra threads while serving sessions use more, so the cap can underestimate.

Two things I checked that are fine

Tokenizer reload after unload does not need network — hf_hub::ApiRepo::get consults the local cache and only falls back to download, so offline reload from EMBEDDING_CACHE_DIR works.

fastembed's intra_threads default is None → available_parallelism(), so the new CPU-count default at pool_size=1 matches what the code was already getting. The real change is the split across sessions when the pool is larger, which is a genuine improvement over master's nproc sessions × nproc threads.

[cursor]

Thanks — both merge blockers are in 05a38d7.

1. Cold load no longer occupies tokio workers: ensure_loaded takes a tokio::sync::Mutex so waiters .await, and load_model runs in block_in_place. Idle unload is spawn_blocking.
2. /health stays 200 while unused (so idle RSS stays low) and returns 503 with the error after a failed load. Failed loads retry at most every 30s instead of re-downloading on every /embed.

Also done from the rest of the review: monotonic last-access (ms), reuse the RAM-capped pool size across reloads, jemalloc on the library crate so tests match the binary, drop the redundant tikv-jemalloc-sys link, comment the probe rebuild, and note that dropping bge-small from compose is a breaking alias change unless you set EMBEDDING_MODELS.

Not in this commit: concurrent p50/p99 vs master (needs a loaded nomic bench on the same box) and jemalloc background_threads.