Skip to content

feat: expose ONNX Runtime session config entries through InitOptions - #292

Merged
anush008 merged 2 commits into
anush008:mainfrom
paulnsorensen:feat/session-config-entries
Sep 12, 2026
Merged

feat: expose ONNX Runtime session config entries through InitOptions#292
anush008 merged 2 commits into
anush008:mainfrom
paulnsorensen:feat/session-config-entries

Conversation

@paulnsorensen

@paulnsorensen paulnsorensen commented Sep 12, 2026

Copy link
Copy Markdown
Contributor

Closes #291.

What changed

  • init_session_builder takes a session_config: Vec<(String, String)> parameter. It applies each entry with SessionBuilder::with_config_entry after with_intra_threads. The DirectML block is unchanged.
  • InitOptions<M> and InitOptionsWithLength<M> gain a session_config field and a with_session_config(key, value) builder method. The default is empty.
  • InitOptionsUserDefined, RerankInitOptionsUserDefined, and ImageInitOptionsUserDefined gain the same field and method. The From impls copy the field.
  • All ten init_session_builder call sites pass the field through.
  • README: one line next to the with_intra_threads example.

Why

Some ONNX Runtime behaviors have a per-session switch only. mlas.disable_kleidiai is one example. On Apple Silicon with SME, ONNX Runtime 1.28.0 keeps packed GEMM buffers in thread-local storage. The buffers never shrink. A long-lived process retains gigabytes after one embedding pass. With with_session_config("mlas.disable_kleidiai", "1") the same process retains 216 MB instead of 2189 MB. See #291 and microsoft/onnxruntime#29538.

The shape follows the existing execution_providers and intra_threads passthroughs.

Usage

let model = TextEmbedding::try_new(
    TextInitOptions::new(EmbeddingModel::SnowflakeArcticEmbedS)
        .with_session_config("mlas.disable_kleidiai", "1"),
)?;

Tests

  • init_session_builder_applies_config_entry calls the real SessionBuilder with one entry and asserts success. It loads no model.
  • Unit tests assert that with_session_config appends in order and that the From impls carry the field.
  • cargo fmt --all -- --check, cargo clippy ... -- -D warnings, and cargo test --lib pass with ort-download-binaries-native-tls,hf-hub-native-tls.

Compatibility

The option structs are #[non_exhaustive], so the new field is not a breaking change for callers that use the builder methods. Callers that construct the structs by field name inside this crate are updated.

🤖 Generated with Claude Code

@paulnsorensen
feat: expose ONNX Runtime session config entries through InitOptions
e6494ad 
fastembed builds the ONNX Runtime SessionBuilder internally. Callers
had no way to set session-only configuration entries such as
mlas.disable_kleidiai, which controls a per-session GEMM buffer switch
on Apple SME hardware.

Add a session_config field to every init-options struct. Add a
with_session_config(key, value) builder method that appends one entry.
init_session_builder applies each entry with
SessionBuilder::with_config_entry after setting intra-op threads and
before the DirectML block.

Update all 10 call sites that build a session from init options. All
From impls between default and user-defined option structs now carry
session_config across.
anush008
anush008 reviewed Sep 12, 2026
View reviewed changes
Comment thread README.md Outdated
@anush008

anush008 commented Sep 12, 2026

Copy link
Copy Markdown
Owner

Thanks for the contribution @paulnsorensen.

@anush008
anush008 merged commit adfefb9 into anush008:main Sep 12, 2026
github-actions Bot pushed a commit that referenced this pull request Sep 12, 2026
@KShivendu @github-actions
chore(release): 6.1.0 [skip ci]
3c5c9f4 
## [6.1.0](v6.0.3...v6.1.0) (2026-09-12)

### 🍕 Features

* expose ONNX Runtime session config entries through InitOptions ([#292](#292)) ([adfefb9](adfefb9))
* inference-free SPLADE (opensearch-neural-sparse-encoding-doc-v3-gte) ([#289](#289)) ([6b9e45d](6b9e45d))

### 📝 Documentation

* Update LICENSE ([#290](#290)) ([a535a2f](a535a2f))
@github-actions

github-actions Bot commented Sep 12, 2026

Copy link
Copy Markdown

🎉 This PR is included in version 6.1.0 🎉

The release is available on:

Your semantic-release bot 📦🚀

anush008 pushed a commit that referenced this pull request Sep 12, 2026
@KShivendu @github-actions
chore(release): 6.1.0 [skip ci]
c79d8dc 
## [6.1.0](v6.0.3...v6.1.0) (2026-09-12)

### 🍕 Features

* expose ONNX Runtime session config entries through InitOptions ([#292](#292)) ([adfefb9](adfefb9))
* inference-free SPLADE (opensearch-neural-sparse-encoding-doc-v3-gte) ([#289](#289)) ([6b9e45d](6b9e45d))

### 📝 Documentation

* Update LICENSE ([#290](#290)) ([a535a2f](a535a2f))
anush008 pushed a commit that referenced this pull request Sep 12, 2026
@KShivendu @github-actions
chore(release): 6.1.0 [skip ci]
7dfa3ff 
## [6.1.0](v6.0.3...v6.1.0) (2026-09-12)

### 🍕 Features

* expose ONNX Runtime session config entries through InitOptions ([#292](#292)) ([adfefb9](adfefb9))
* inference-free SPLADE (opensearch-neural-sparse-encoding-doc-v3-gte) ([#289](#289)) ([6b9e45d](6b9e45d))

### 📝 Documentation

* Update LICENSE ([#290](#290)) ([a535a2f](a535a2f))
anush008 pushed a commit that referenced this pull request Sep 12, 2026
@KShivendu @github-actions
chore(release): 6.1.0 [skip ci]
8caec12 
## [6.1.0](v6.0.3...v6.1.0) (2026-09-12)

### 🍕 Features

* expose ONNX Runtime session config entries through InitOptions ([#292](#292)) ([adfefb9](adfefb9))
* inference-free SPLADE (opensearch-neural-sparse-encoding-doc-v3-gte) ([#289](#289)) ([6b9e45d](6b9e45d))

### 📝 Documentation

* Update LICENSE ([#290](#290)) ([a535a2f](a535a2f))
@paulnsorensen paulnsorensen mentioned this pull request Sep 14, 2026
Merged
4 tasks
paulnsorensen added a commit to paulnsorensen/hallouminate that referenced this pull request Sep 14, 2026
@paulnsorensen @claude
fix(adapters): stop ONNX Runtime from retaining gigabytes after embed…
e393d95 
…ding (#505)

## What changed

Semantics-altering. ONNX Runtime sessions in `hallouminate-adapters` now
start with the CPU memory arena off, a 416-token embed `max_length`, and
the `mlas.disable_kleidiai=1` session config entry. fastembed moves to
the released 6.1.0, which exposes `with_session_config`.

## Why

The resident daemon retained 4.8 GB after one boot catch-up embed of 240
markdown files. `malloc_history` traced every retained block to
`MatMul<float> -> MlasGemmBatch -> ArmKleidiAI::MlasGemmBatch`. ORT 1.28
caches packed GEMM operands in `static thread_local KaiTlsBuffers` that
never shrink (onnxruntime#29538). The ORT BFCArena added a second
retention layer on top. Neither is released when the session drops.

`mlas.disable_kleidiai` is a per-session config entry. fastembed 6.0.3
exposed no hook, so this branch first pinned a fork through
`[patch.crates-io]`. Upstream merged anush008/fastembed-rs#292 (issue
#291) and released 6.1.0 on 2026-09-12. The second commit removes the
patch and requires `fastembed = "6.1"`.

Measured on the final binary, 240 files, boot catch-up embed,
arctic-embed-s fp32:

| Build | Settled footprint | Peak |
|---|---|---|
| Before | 4843 MB | 5041 MB |
| After | 217 MB | 2067 MB |

Embed pass runs about 40 % slower on SME-capable Apple Silicon (M4+).
KleidiAI never ran elsewhere, so the flag is a no-op on other hardware.

## How to verify

```sh
just verify
cargo test -p hallouminate-adapters --lib -- session_config cpu_no_arena
```

Memory check on an M4+ Mac: install the build, run `hallouminate daemon
restart`, wait for boot catch-up, then compare `footprint <pid>` against
the previous binary.

## Non-obvious changes

- `EMBED_MAX_LENGTH = 416`: `CHUNK_BUDGET_TOKENS` (384) plus headroom
for `[CLS]`/`[SEP]` and the e5 `query:`/`passage:` prefix. Attention
scratch scales with `seq_len²`.
- `EMBED_BATCH_SIZE = 32` stays. With the arena off it bounds transient
peak, not arena growth. The doc comment says so.
- `Cargo.lock`: only `fastembed` moves (`6.0.3` git fork -> `6.1.0`
crates.io).

## Risk / rollout notes

- ~40 % slower embed pass on M4+ Macs. Acceptable for a resident daemon.
- `.hallouminate/wiki/ort-arena-retention.md` rewritten: KleidiAI TLS is
the dominant retention on SME Macs; the arena was secondary.
- Upstream evidence posted:
microsoft/onnxruntime#29538 (comment)
- Composes with #472 (reranker admission). No shared hunks.
- Parked follow-ups: reconcile hot loop on a missing corpus path; 8040
orphan `_indices` dirs.

## Checklist

- [x] Conventional Commits PR title (`feat:`, `fix:`, `chore:`, `docs:`,
...)
- [x] Tests added or updated (or N/A)
- [x] Documentation updated (or N/A)
- [x] No secrets or credentials added

🤖 Generated with [Claude Code](https://claude.com/claude-code)


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

- **Bug Fixes**
- Reduced retained CPU memory during embedding and reranking workloads.
- Improved memory bounds by limiting processed text length and
batch-related allocations.
- Applied more predictable CPU execution behavior across supported
platforms.

- **Documentation**
- Updated technical documentation with current memory-retention
mechanisms, verified measurements, fixes, and guidance for measuring
memory usage.
  - Removed outdated idle-eviction information.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@anush008 anush008 anush008 left review comments

Assignees

No one assigned

Labels

released

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

Expose ONNX Runtime session configuration entries through InitOptions

2 participants

@paulnsorensen @anush008