Skip to content

fix(registry): lock ModeRegistry lazy-load resolution against concurrent first access#2439

Open
ErenAta16 wants to merge 2 commits into
567-labs:mainfrom
ErenAta16:fix/mode-registry-lazy-load-race
Open

fix(registry): lock ModeRegistry lazy-load resolution against concurrent first access#2439
ErenAta16 wants to merge 2 commits into
567-labs:mainfrom
ErenAta16:fix/mode-registry-lazy-load-race

Conversation

@ErenAta16

Copy link
Copy Markdown

Summary

Fixes #2422. ModeRegistry.get_handlers() has an unguarded check -> pop -> import -> set sequence over self._handlers and self._lazy_loaders (plain dicts shared across threads on the module-level mode_registry singleton). On the first concurrent access to a given (provider, mode) key in a process, one thread pops the lazy loader and starts the module import; every other thread that already passed the mode_key not in self._handlers check then finds the loader also already popped from self._lazy_loaders, and raises KeyError (surfaced elsewhere as RegistryError) for a mode that genuinely is registered, just not finished resolving yet. Root cause and regression bisection (introduced between 1.15.2 and 1.15.3 when the v2 lazy-loading registry landed) were already thoroughly diagnosed in the issue.

Fix

instructor/v2/core/registry.py: wrapped the lazy-load resolution in a lock, with a re-check inside the lock so a thread that waited for another's in-flight load reuses the result instead of raising. Held for the full resolution rather than a per-key lock, since lazy-loading only ever runs once per key ever (subsequent calls hit the fast mode_key in self._handlers check before acquiring anything), this is a one-time cost per key at first use, not an ongoing bottleneck.

Testing

Added test_get_handlers_concurrent_first_access_does_not_race to tests/v2/test_registry.py, using a fresh ModeRegistry() instance (not the global singleton) with a loader that sleeps briefly to simulate a slow import, and a Barrier so all 8 worker threads call get_handlers() as close to simultaneously as possible.

Verified this is a meaningful regression test, not just a test that happens to pass either way:

  • Against unfixed registry.py: 7 of 8 threads raise KeyError (reproduces the issue's exact failure mode).
  • Against the fix: all 8 threads get the same ModeHandlers instance back, and the loader runs exactly once (asserted via a counter).

Ran the full tests/v2/ suite locally: 1578 passed, 1 pre-existing unrelated failure (test_vertexai_process_helpers_build_tools_and_config, missing optional jsonref dependency in my environment, nothing to do with this change), 170 skipped.

…ent first access

ModeRegistry.get_handlers() had an unguarded check -> pop -> import ->
set sequence over two plain dicts (self._handlers, self._lazy_loaders)
shared across threads. On the first concurrent access to a given
(provider, mode) key in a process, one thread pops the lazy loader and
starts (possibly slow) module import; every other thread that already
passed the "not yet in self._handlers" check then finds the loader
also already popped from self._lazy_loaders, and raises KeyError
(surfaced as RegistryError) for a mode that genuinely is registered,
just not finished resolving yet.

Wrapped the lazy-load resolution in a lock, with a re-check inside
the lock so a thread that waited for another's in-flight load reuses
the result instead of raising. Held for the full resolution rather
than per-key, since lazy-loading only ever runs once per key, this is
a one-time cost per key at first use, not an ongoing bottleneck.

Added a regression test that forces the race deterministically (all
threads start together via a Barrier, loader sleeps briefly to
simulate a slow import) rather than relying on timing. Confirmed it
fails on unfixed registry.py (7 of 8 threads raise KeyError,
reproducing the issue exactly) and passes with the fix (all 8 threads
get the same ModeHandlers instance, loader runs exactly once).

Fixes 567-labs#2422

Signed-off-by: ErenAta16 <erena6466@gmail.com>
…rrupted)

The prior commit on this branch had both changed files replaced with
a handful of bytes of binary garbage instead of real source, a gh CLI
file-encoding issue on my end when uploading the base64 blob content
(same root cause as litellm/litellm#33139, unrelated to this repo).
Recreated both blobs via a Python-built JSON request body, verified
sizes match the local files exactly (14268 and 6721 bytes) before
pushing this time.

No functional change from the intended fix, this just repairs the
upload.

Signed-off-by: ErenAta16 <erena6466@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Thread-unsafe lazy loading in ModeRegistry.get_handlers() causes RegistryError under concurrency

1 participant