fix(registry): lock ModeRegistry lazy-load resolution against concurrent first access#2439
Open
ErenAta16 wants to merge 2 commits into
Open
fix(registry): lock ModeRegistry lazy-load resolution against concurrent first access#2439ErenAta16 wants to merge 2 commits into
ErenAta16 wants to merge 2 commits into
Conversation
…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>
8 tasks
…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>
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
Fixes #2422.
ModeRegistry.get_handlers()has an unguarded check -> pop -> import -> set sequence overself._handlersandself._lazy_loaders(plain dicts shared across threads on the module-levelmode_registrysingleton). 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 themode_key not in self._handlerscheck then finds the loader also already popped fromself._lazy_loaders, and raisesKeyError(surfaced elsewhere asRegistryError) 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 fastmode_key in self._handlerscheck 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_racetotests/v2/test_registry.py, using a freshModeRegistry()instance (not the global singleton) with a loader that sleeps briefly to simulate a slow import, and aBarrierso all 8 worker threads callget_handlers()as close to simultaneously as possible.Verified this is a meaningful regression test, not just a test that happens to pass either way:
registry.py: 7 of 8 threads raiseKeyError(reproduces the issue's exact failure mode).ModeHandlersinstance 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 optionaljsonrefdependency in my environment, nothing to do with this change), 170 skipped.