Skip to content

Commit ffa9e8e

Browse files
committed
fix: best-effort HNSW thread-pin retrofit + drop dead attempt-cap constant
Addresses remaining PR MemPalace#976 review items after rebase on develop. ## Copilot Review MemPalace#3 (chroma.py:134) — retrofit legacy collections `get_collection(create=False)` previously returned existing collections without re-applying `hnsw:num_threads=1`, so palaces created before the fix kept the unsafe parallel-insert path. Add `_pin_hnsw_threads()` helper that calls `collection.modify(configuration=UpdateCollectionConfiguration( hnsw=UpdateHNSWConfiguration(num_threads=1)))` best-effort on every `get_collection` call (including the MCP server's `_get_collection`). In chromadb 1.5.x the runtime config does not persist to disk across `PersistentClient` reopens, so the retrofit is re-applied each process start rather than being a one-shot migration. Fresh palaces keep the metadata-based pin as primary defense; legacy palaces now also get per-session protection without requiring `mempalace nuke` + re-mine. ## mvalentsev feedback — drop dead `MAX_PRECOMPACT_BLOCK_ATTEMPTS` After the rebase on develop, `hook_precompact` delegates to `_mine_sync` and no longer emits `decision: block`, so the attempt-cap constant was orphaned. Grep confirms 0 usages in the repo — remove it. ## Tests - `_pin_hnsw_threads` retrofits legacy collection (num_threads None -> 1) - `_pin_hnsw_threads` swallows all errors (never raises) - `ChromaBackend.get_collection(create=False)` applies retrofit on legacy palace - 62 tests pass (10 backends + 6 palace locks + 46 hooks_cli)
1 parent 404d73f commit ffa9e8e

4 files changed

Lines changed: 103 additions & 15 deletions

File tree

mempalace/backends/chroma.py

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,38 @@
1111
logger = logging.getLogger(__name__)
1212

1313

14+
def _pin_hnsw_threads(collection) -> None:
15+
"""Best-effort retrofit: pin ``hnsw:num_threads=1`` on an existing collection.
16+
17+
Fresh collections set this via ``metadata=`` at creation. Legacy palaces
18+
built before that change keep the default (parallel insert) and can hit
19+
the HNSW race described in #974/#965. ChromaDB's
20+
``collection.modify(configuration=...)`` lets us re-apply ``num_threads=1``
21+
in memory at load time so every new process is protected.
22+
23+
Note: in chromadb 1.5.x the modified ``configuration_json["hnsw"]`` does
24+
not persist to disk across ``PersistentClient`` reopens, so this must
25+
run on every ``get_collection`` call, not just once.
26+
"""
27+
try:
28+
from chromadb.api.collection_configuration import (
29+
UpdateCollectionConfiguration,
30+
UpdateHNSWConfiguration,
31+
)
32+
except ImportError:
33+
# Older chromadb (pre-1.5) doesn't expose UpdateCollectionConfiguration.
34+
logger.debug("_pin_hnsw_threads skipped: chromadb too old", exc_info=True)
35+
return
36+
try:
37+
collection.modify(
38+
configuration=UpdateCollectionConfiguration(
39+
hnsw=UpdateHNSWConfiguration(num_threads=1)
40+
)
41+
)
42+
except Exception:
43+
logger.debug("_pin_hnsw_threads modify failed", exc_info=True)
44+
45+
1446
def _fix_blob_seq_ids(palace_path: str):
1547
"""Fix ChromaDB 0.6.x -> 1.5.x migration bug: BLOB seq_ids -> INTEGER.
1648
@@ -131,6 +163,7 @@ def get_collection(self, palace_path: str, collection_name: str, create: bool =
131163
)
132164
else:
133165
collection = client.get_collection(collection_name)
166+
_pin_hnsw_threads(collection)
134167
return ChromaCollection(collection)
135168

136169
def get_or_create_collection(

mempalace/hooks_cli.py

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -262,9 +262,6 @@ def hook_session_start(data: dict, harness: str):
262262
_output({})
263263

264264

265-
MAX_PRECOMPACT_BLOCK_ATTEMPTS = 2
266-
267-
268265
def hook_precompact(data: dict, harness: str):
269266
"""Precompact hook: mine transcript synchronously, then allow compaction."""
270267
parsed = _parse_harness_input(data, harness)

mempalace/mcp_server.py

Lines changed: 14 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@
5757
sanitize_content,
5858
)
5959
from .version import __version__ # noqa: E402
60-
from .backends.chroma import ChromaBackend, ChromaCollection # noqa: E402
60+
from .backends.chroma import ChromaBackend, ChromaCollection, _pin_hnsw_threads # noqa: E402
6161
from .query_sanitizer import sanitize_query # noqa: E402
6262
from .searcher import search_memories # noqa: E402
6363
from .palace_graph import ( # noqa: E402
@@ -219,20 +219,23 @@ def _get_collection(create=False):
219219
if create:
220220
# hnsw:num_threads=1 disables ChromaDB's multi-threaded ParallelFor
221221
# HNSW insert path, which has a race in repairConnectionsForUpdate /
222-
# addPoint (see issues #974, #965). The setting is only honored at
223-
# collection creation time — pre-existing palaces created before
224-
# this fix keep the unsafe default; users must `mempalace nuke` +
225-
# re-mine to get the protection on legacy palaces.
226-
_collection_cache = ChromaCollection(
227-
client.get_or_create_collection(
228-
_config.collection_name,
229-
metadata={"hnsw:space": "cosine", "hnsw:num_threads": 1},
230-
)
222+
# addPoint (see issues #974, #965). Set via metadata on fresh
223+
# collections and re-applied via _pin_hnsw_threads() for legacy
224+
# palaces whose collections were created before this fix (the
225+
# runtime config does not persist cross-process in chromadb 1.5.x,
226+
# so the retrofit runs every time _get_collection opens a cache).
227+
raw = client.get_or_create_collection(
228+
_config.collection_name,
229+
metadata={"hnsw:space": "cosine", "hnsw:num_threads": 1},
231230
)
231+
_pin_hnsw_threads(raw)
232+
_collection_cache = ChromaCollection(raw)
232233
_metadata_cache = None
233234
_metadata_cache_time = 0
234235
elif _collection_cache is None:
235-
_collection_cache = ChromaCollection(client.get_collection(_config.collection_name))
236+
raw = client.get_collection(_config.collection_name)
237+
_pin_hnsw_threads(raw)
238+
_collection_cache = ChromaCollection(raw)
236239
_metadata_cache = None
237240
_metadata_cache_time = 0
238241
return _collection_cache

tests/test_backends.py

Lines changed: 56 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,12 @@
33
import chromadb
44
import pytest
55

6-
from mempalace.backends.chroma import ChromaBackend, ChromaCollection, _fix_blob_seq_ids
6+
from mempalace.backends.chroma import (
7+
ChromaBackend,
8+
ChromaCollection,
9+
_fix_blob_seq_ids,
10+
_pin_hnsw_threads,
11+
)
712

813

914
class _FakeCollection:
@@ -140,3 +145,53 @@ def test_fix_blob_seq_ids_noop_without_blobs(tmp_path):
140145
def test_fix_blob_seq_ids_noop_without_database(tmp_path):
141146
"""No error when palace has no chroma.sqlite3."""
142147
_fix_blob_seq_ids(str(tmp_path)) # should not raise
148+
149+
150+
def test_pin_hnsw_threads_retrofits_legacy_collection(tmp_path):
151+
"""Legacy collections (created without num_threads) get the retrofit applied."""
152+
palace_path = tmp_path / "legacy-palace"
153+
palace_path.mkdir()
154+
155+
client = chromadb.PersistentClient(path=str(palace_path))
156+
col = client.create_collection(
157+
"mempalace_drawers",
158+
metadata={"hnsw:space": "cosine"}, # no num_threads — legacy
159+
)
160+
assert col.configuration_json.get("hnsw", {}).get("num_threads") is None
161+
162+
_pin_hnsw_threads(col)
163+
164+
assert col.configuration_json["hnsw"]["num_threads"] == 1
165+
166+
167+
def test_pin_hnsw_threads_swallows_all_errors():
168+
"""Retrofit never raises even when collection.modify explodes."""
169+
170+
class _ExplodingCollection:
171+
def modify(self, *args, **kwargs):
172+
raise RuntimeError("boom")
173+
174+
_pin_hnsw_threads(_ExplodingCollection()) # must not raise
175+
176+
177+
def test_get_collection_applies_retrofit_on_existing_palace(tmp_path):
178+
"""ChromaBackend.get_collection(create=False) applies the retrofit."""
179+
palace_path = tmp_path / "palace"
180+
palace_path.mkdir()
181+
182+
# Simulate a legacy palace: create collection without num_threads
183+
bootstrap_client = chromadb.PersistentClient(path=str(palace_path))
184+
bootstrap_client.create_collection(
185+
"mempalace_drawers", metadata={"hnsw:space": "cosine"}
186+
)
187+
del bootstrap_client # drop reference so a fresh client reopens cleanly
188+
189+
wrapper = ChromaBackend().get_collection(
190+
str(palace_path),
191+
collection_name="mempalace_drawers",
192+
create=False,
193+
)
194+
195+
assert (
196+
wrapper._collection.configuration_json["hnsw"]["num_threads"] == 1
197+
)

0 commit comments

Comments
 (0)