fix: cap ONNX intra_op threads via MEMPAL_MAX_THREADS - #1071
Closed
sha2fiddy wants to merge 2 commits into
Closed
Conversation
sha2fiddy
requested review from
bensig,
igorls and
milla-jovovich
as code owners
April 21, 2026 13:21
jphein
added a commit
to techempower-org/mempalace
that referenced
this pull request
Apr 21, 2026
…1088, revise PR ordering
This was referenced Apr 27, 2026
sha2fiddy
force-pushed
the
fix/ort-intra-op-thread-cap
branch
from
April 28, 2026 16:09
5f98cea to
fcd52aa
Compare
sha2fiddy
changed the title
MEMPAL_MAX_THREADS
sha2fiddy
force-pushed
the
fix/ort-intra-op-thread-cap
branch
from
April 29, 2026 23:01
fcd52aa to
a1b076b
Compare
sha2fiddy
force-pushed
the
fix/ort-intra-op-thread-cap
branch
from
May 15, 2026 19:27
a1b076b to
5e34637
Compare
sha2fiddy
force-pushed
the
fix/ort-intra-op-thread-cap
branch
2 times, most recently
from
May 27, 2026 00:18
62f6fca to
c5fe9f7
Compare
sha2fiddy
added a commit
to sha2fiddy/mempalace
that referenced
this pull request
May 27, 2026
PR MemPalace#1071 changes _build_ef_class to accept thread_cap=N. The test in tests/test_embeddinggemma.py (added on develop) stubs _build_ef_class with a zero-arg lambda, which now fails after the rebase. Update the lambda and clear MEMPAL_MAX_THREADS so the cap branch is exercised deterministically.
sha2fiddy
added a commit
to sha2fiddy/mempalace
that referenced
this pull request
May 27, 2026
PR MemPalace#1071 changes _build_ef_class to accept thread_cap=N. The test in tests/test_embeddinggemma.py (added on develop) stubs _build_ef_class with a zero-arg lambda, which now fails after the rebase. Update the lambda and clear MEMPAL_MAX_THREADS so the cap branch is exercised deterministically.
sha2fiddy
force-pushed
the
fix/ort-intra-op-thread-cap
branch
from
May 27, 2026 18:16
bbf709b to
a8a25d6
Compare
sha2fiddy
added a commit
to sha2fiddy/mempalace
that referenced
this pull request
May 30, 2026
PR MemPalace#1071 changes _build_ef_class to accept thread_cap=N. The test in tests/test_embeddinggemma.py (added on develop) stubs _build_ef_class with a zero-arg lambda, which now fails after the rebase. Update the lambda and clear MEMPAL_MAX_THREADS so the cap branch is exercised deterministically.
sha2fiddy
force-pushed
the
fix/ort-intra-op-thread-cap
branch
from
May 30, 2026 15:51
a8a25d6 to
44918f0
Compare
Without 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.
PR MemPalace#1071 changes _build_ef_class to accept thread_cap=N. The test in tests/test_embeddinggemma.py (added on develop) stubs _build_ef_class with a zero-arg lambda, which now fails after the rebase. Update the lambda and clear MEMPAL_MAX_THREADS so the cap branch is exercised deterministically.
sha2fiddy
force-pushed
the
fix/ort-intra-op-thread-cap
branch
from
June 3, 2026 13:09
44918f0 to
e780749
Compare
jphein
added a commit
to techempower-org/mempalace
that referenced
this pull request
Jun 14, 2026
#349) * feat(miner): restore concurrent mining via parallel-prepare/serial-write The `--workers` CLI flag had been declared (cli.py) but consumed nowhere — `args.workers` had no reader, so mining always ran serially. Concurrent mining originally shipped in 5cd14bd ("feat: concurrent mining with ThreadPoolExecutor"), but a later upstream sync dropped the miner.py wiring while leaving the flag stranded. This restores concurrency on today's architecture (pgvector/RFC-001 backends, id_recipe, collision_scan, novelty, closets) rather than cherry-picking the pre-architecture code. Design — parallel-prepare / serial-write split: - `_prepare_file()` does the embedding-free, read-only work (read, detect_room, room_resolver, chunk_text, chunk-cap check, source_mtime + content_date) and returns a `_PreparedFile`. It touches no backend state, so it is safe to fan out across a ThreadPoolExecutor. - `_write_prepared()` performs the backend write for one prepared file (already-mined check, mine_lock, post-lock re-check, stale-drawer purge, novelty tag, batched assert_no_collisions + collection.upsert, closet build). Byte-identical to the write half of the unchanged `process_file`. - `_mine_files_concurrently()` submits prep across the pool and, as each future completes, runs `_write_prepared` on the MAIN thread. Why writes stay serial: embedding happens inside `collection.upsert(documents=...)` (the backend embeds on write). Keeping every upsert on the main thread keeps the ONNX/encoder path single-threaded no matter how many workers prepare files in parallel — the data-safety property behind the concurrent-ingest thread-safety gate (#330). This composes with upstream MemPalace#1071 (encoder thread-safety) but does not depend on it: even without MemPalace#1071, the encoder is only ever entered from one thread here. Behaviour: - Default `workers=1` everywhere (library + CLI). `workers <= 1` (and every dry-run) keeps the EXACT original sequential path — zero behaviour change for existing callers. `workers > 1` is strictly opt-in. - Progress logging, counts, dry-run, limit, and per-file error handling are preserved. A file that errors during prep is logged and skipped (does not abort the mine); a serial-write failure propagates with the existing partial-progress summary, same as today's `process_file` exception path. - Honest `--workers` help text: prep parallelizes, writes stay serialized so embedding is single-threaded. Tests (no real ONNX — injected stub collections): - concurrent mine writes byte-identical drawers (ids + documents) to sequential - writes are serialized: a thread-recording collection asserts no two upserts overlap and all upserts land on a single thread under workers=4 - a file that raises during prep is skipped while every other file still files Closes the parallel-corruption concern by construction (writes never run concurrently). The "#357" reference in the task framing does not resolve in this repo; the live data-safety gate is #330. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * docs: sync test count (4282→4285) + regen api docs for concurrent-mining helpers The new _prepare_file / _write_prepared / _mine_files_concurrently functions and 3 new tests tripped check-docs (README count + auto-rendered python-api). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * docs(changelog): record concurrent-mining restore in fork-changes.yaml Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * docs: render README fork-change-queue table for concurrent-mining entry Completes the previous changelog commit — the README table render output was inadvertently left unstaged. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Member
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
2 participants
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.