Cut idle embedding memory without changing nomic quality

.env

@@ -1,3 +1,2 @@
-EMBEDDING_MODELS=nomic,bge-small
+EMBEDDING_MODELS=nomic
 EMBEDDING_CACHE_DIR=./models
-EMBEDDING_POOL_SIZE=1

Cargo.toml

@@ -5,13 +5,16 @@ edition = "2024"
 
 [dependencies]
 axum = "0.8.7"
-fastembed = "5.13.4"
+fastembed = "5.15.0"
 futures = "0.3.32"
 hf-hub = "0.5.0"
 serde = { version = "1.0.219", features = ["derive"] }
 serde_json = "1.0.145"
 sysinfo = "0.39.2"
 tokenizers = "0.22.2"
-tokio = { version = "1.52.3", features = ["macros", "rt-multi-thread", "net", "signal"] }
+tokio = { version = "1.52.3", features = ["macros", "rt-multi-thread", "net", "signal", "time"] }
 tracing = "0.1.44"
 tracing-subscriber = { version = "0.3.20", features = ["env-filter"] }
+
+[target.'cfg(not(target_env = "msvc"))'.dependencies]
+tikv-jemallocator = { version = "0.6", features = ["override_allocator_on_supported_platforms"] }

Dockerfile

@@ -5,6 +5,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
         pkg-config \
         libssl-dev \
         g++ \
+        make \
     && rm -rf /var/lib/apt/lists/*
 WORKDIR /app
 COPY . .
@@ -34,7 +35,7 @@ COPY --from=builder /usr/local/bin/warmup /usr/local/bin/warmup
 # the model set with `--build-arg EMBEDDING_MODELS=...` (docker compose passes
 # this from .env). Pool size is forced to 1 to keep the build's memory low —
 # it only affects the warmup, not the runtime pool.
-ARG EMBEDDING_MODELS=nomic,bge-small
+ARG EMBEDDING_MODELS=nomic
 RUN EMBEDDING_MODELS="${EMBEDDING_MODELS}" EMBEDDING_POOL_SIZE=1 /usr/local/bin/warmup
 
 # EXPOSE is build-time metadata only; the actual port is controlled by the

README.md

@@ -10,6 +10,8 @@ docker compose up --build
 
 First request triggers the model download into `./models` (bind-mounted into the container); subsequent restarts reuse it.
 
+`GET /health` is liveness: it stays `200` while unused so idle RSS is not pinned. After a failed model load it returns `503` with the error for 30 seconds, then `200` again so a probe can recover; the next `/embed` retries the load.
+
 ```bash
 curl -X POST http://localhost:3000/embed \
   -H 'content-type: application/json' \
@@ -23,9 +25,11 @@ Configured via environment variables (set them in `.env`):
 | Variable | Default | Description |
 | --- | --- | --- |
 | `EMBEDDING_PORT` | `3000` | Port the service listens on. |
-| `EMBEDDING_MODELS` | `nomic` | Comma-separated list of models to load. |
+| `EMBEDDING_MODELS` | `nomic` | Comma-separated list of models allowed to load. ONNX sessions are created on first `/embed`, not at process start. Image/compose defaults are nomic only; `bge-small` and other aliases still work if you add them here. |
 | `EMBEDDING_CACHE_DIR` | _(default cache)_ | Directory for downloaded model files. |
-| `EMBEDDING_POOL_SIZE` | _(memory-derived)_ | Number of model instances per pool. |
+| `EMBEDDING_POOL_SIZE` | `1` | Number of ONNX sessions per model while it is loaded, then capped by available RAM. Concurrent `/embed` calls round-robin across sessions. Raise this for parallel HTTP throughput; each extra session keeps another copy of the weights resident until idle unload. |
+| `EMBEDDING_INTRA_THREADS` | CPU count | ONNX Runtime intra-op threads per session. The default uses the whole machine on the single default session. When `EMBEDDING_POOL_SIZE` is greater than one, threads are split across sessions (`nproc / pool_size`, still capped by this value) so concurrent embeds do not oversubscribe the host. |
+| `EMBEDDING_IDLE_UNLOAD_SECS` | `300` | Drop a model's sessions this many seconds after last use (`0` disables). The next `/embed` reloads the same checkpoint from `EMBEDDING_CACHE_DIR`. |
 
 ## API
 
@@ -49,5 +53,6 @@ Response:
 
 Errors:
 
-- `400 Bad Request` — `texts` is empty.
+- `400 Bad Request` — `texts` is empty or the model alias is not in `EMBEDDING_MODELS`.
 - `500 Internal Server Error` — embedding or tokenizer failure (message in `error` field).
+- `503 Service Unavailable` — `GET /health` for 30 seconds after a model load has failed (message in `error` field).

docker-compose.yml

@@ -4,7 +4,7 @@ services:
       context: .
       dockerfile: Dockerfile
       args:
-        EMBEDDING_MODELS: ${EMBEDDING_MODELS:-nomic,bge-small}
+        EMBEDDING_MODELS: ${EMBEDDING_MODELS:-nomic}
     image: embedding:latest
     container_name: embedding
     ports:

src/bin/warmup.rs

@@ -13,7 +13,8 @@ fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
         "warmup: downloading and initializing {} model(s)",
         config.models.len()
     );
-    let _ = EmbeddingClient::new(config)?;
+    let client = EmbeddingClient::new(config)?;
+    client.preload()?;
     tracing::info!("warmup: models cached and ready");
 
     Ok(())