fix: cap ONNX intra_op threads via MEMPAL_MAX_THREADS#1071
Open
sha2fiddy wants to merge 1 commit intoMemPalace:developfrom
Open
fix: cap ONNX intra_op threads via MEMPAL_MAX_THREADS#1071sha2fiddy wants to merge 1 commit intoMemPalace:developfrom
sha2fiddy wants to merge 1 commit intoMemPalace:developfrom
Conversation
jphein
added a commit
to jphein/mempalace
that referenced
this pull request
Apr 21, 2026
…emPalace#1088, revise PR ordering
This was referenced Apr 27, 2026
5f98cea to
fcd52aa
Compare
MEMPAL_MAX_THREADSWithout this, ORT spawns ~physical-core-count workers in its intra_op pool and a background mine pegs 400-500% CPU. OMP_NUM_THREADS does not control the ORT pool — ORT has its own. Cap is applied in mempalace.embedding by overriding the model cached_property on the EF subclass so the InferenceSession is built with explicit SessionOptions (intra_op_num_threads=N, inter_op=1). Default cap is 2; "0"/"off"/"default"/"none" disables it. HNSW thread pinning is already handled by _pin_hnsw_threads on develop, so this PR only addresses the ONNX side. The auto_save and precompact hooks export MEMPAL_MAX_THREADS=2 so background mines stay throttled. Closes MemPalace#1068.
fcd52aa to
a1b076b
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Caps the ONNX Runtime intra_op pool so background mines don't pin every core. Controlled by
MEMPAL_MAX_THREADS(default2;0/off/default/nonedisables). ORT has its own intra-op pool, soOMP_NUM_THREADSdoesn't reach it.HNSW thread pinning was independently fixed by #1191, so after the rebase this PR is ONNX-only.
Closes #1068.
Changes
mempalace/embedding.py:_read_thread_cap()parses the env var._build_ef_class(thread_cap)overrides@cached_property modelon the EF subclass so the ORT session is built withSessionOptions(intra_op_num_threads=N, inter_op=1, log_severity_level=3).get_embedding_function()cache key extended to(providers, thread_cap).hooks/mempal_save_hook.sh,hooks/mempal_precompact_hook.sh: exportMEMPAL_MAX_THREADS=2 TOKENIZERS_PARALLELISM=falseso the background mine inherits the cap.tests/test_embedding.py: 7 new tests covering env-var parsing, the capped subclass, session options, and cache keying.How to test
```bash
ruff check . && python -m pytest tests/ -v --ignore=tests/benchmarks
```
1478 passed locally.
CPU sanity check:
```bash
MEMPAL_MAX_THREADS=2 mempalace mine
another terminal:
while true; do ps -o pid,%cpu -p "$(pgrep -f 'mempalace mine')" 2>/dev/null; sleep 1; done
```
Benchmarks (M-series Mac, 10 cores)
Design
chromadb 1.5 doesn't expose `SessionOptions` on `init`, so the cap goes on the EF subclass. Overriding `cached_property model` is the smallest stable hook into ORT session construction. When `thread_cap=0` the override is skipped and chromadb's defaults apply.