Add lazy model registry - #256
Conversation
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #256 +/- ##
==========================================
- Coverage 93.66% 93.47% -0.20%
==========================================
Files 176 176
Lines 19840 20319 +479
Branches 1279 1350 +71
==========================================
+ Hits 18584 18994 +410
- Misses 1015 1052 +37
- Partials 241 273 +32 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
9a5bcec to
de7f4da
Compare
ptomecek
left a comment
There was a problem hiding this comment.
I went through this fairly closely, mostly focused on what happens when a config references something that hasn't been instantiated yet. The lazy-resolution design holds up well there: because access is what drives materialization, forward references and out-of-order definitions just work, and you don't need the eager loader's fixpoint-retry loop at all. I confirmed that with a few throwaway configs (entry defined before its dependency, relative and absolute refs) and it behaves as you'd hope.
A few inline notes below. The concurrency one is the only thing I'd consider blocking; the rest are about error-message quality and a minor inconsistency.
The bigger gap for me is documentation. LazyRegistry is public now (exported from ccflow, listed in __all__), but nothing under docs/ mentions it:
docs/wiki/Configuration.mdhas a "The Model Registry" section that walks throughModelRegistry, sub-registries, and name-based references. That's the natural home for lazy loading. Worth covering: the_recursive_: falserequirement on the Hydra config, that entries instantiate on first access,materialize_all()for validating a whole tree at once (handy in tests), and the behavioral change that missing/cyclic references now fail at access time instead of load time.docs/wiki/Key-Features.mdhas a table of the notable classes (LazyEvaluator, etc.);LazyRegistryshould be in it for discoverability.
I don't think a dedicated notebook is necessary here — a short subsection with a snippet in Configuration.md would be enough. Happy to draft that if it helps.
| def _materialize(self, name: str) -> BaseModel: | ||
| from hydra.utils import instantiate | ||
|
|
||
| with self._lock: |
There was a problem hiding this comment.
Holding self._lock across the instantiate call means the lock stays held while materialization recurses into other registries to resolve references. Two lazy registries that reference each other can deadlock under concurrent access: thread A locks A and waits on B, thread B locks B and waits on A.
I reproduced this with A.x -> /B/p and B.q -> /A/z materialized on two threads — both hang indefinitely. test_independent_registries_materialize_concurrently doesn't catch it because those registries don't reference each other.
A few ways out: a single process-wide materialization lock, releasing the lock before recursing (resolve dependencies outside the critical section), or explicitly documenting that concurrent materialization across mutually-referencing registries isn't supported. Whichever you pick, a test with cross-references would be good to have.
| raise KeyError(f"No registered model found by the name '{name}' in registry '{self._debug_name}'") | ||
| if name in self._loading: | ||
| cycle = " -> ".join([*self._loading, name]) | ||
| raise RegistryKeyError(f"Circular lazy registry dependency detected: {cycle}") |
There was a problem hiding this comment.
_loading is per-registry, so a cycle that spans two lazy registries only shows the part tracked by whichever registry happens to raise. A cycle like /A/x -> /B/y -> /A/x reports x -> x instead of the full path. Detection itself is fine (no infinite recursion), but the chain is the reason we're tracking this at all, and a partial one is more confusing than helpful. A shared loading stack (thread-local) would let you report the complete cycle.
| search_registry = search_registries[idx] | ||
| try: | ||
| return search_registry[v] | ||
| except RegistryKeyError: |
There was a problem hiding this comment.
This re-raise only helps for top-level lookups. When the cycle is hit while validating a nested model field, pydantic wraps the RegistryKeyError in a ValidationError before it ever gets back here, and hydra then wraps that in an InstantiationException — so the "Circular lazy registry dependency detected" text ends up buried two layers down. test_cycle_reports_dependency_chain only passes because it regex-searches the full stringified exception. Same story for a plain missing reference: you get a raw InstantiationException at access time rather than the eager loader's grouped, readable RegistryKeyError. Not blocking, but the failure UX for cycles and missing refs is a step down from the eager path, and that's exactly where people will be debugging.
| if not isinstance(item, str): | ||
| return False | ||
| if REGISTRY_SEPARATOR in item: | ||
| if "." in item: |
There was a problem hiding this comment.
Minor inconsistency: in returns False for a path containing ., while __getitem__ raises ValueError for the same input. Someone guarding with if key in reg: then indexing will get different behavior depending on which method notices the .. Not worth much, just noting it.
| return self | ||
|
|
||
| def __getitem__(self, item) -> ModelType: | ||
| if REGISTRY_SEPARATOR in item: |
There was a problem hiding this comment.
This only materializes for separator-free keys. When item contains a /, it's handed straight to the base __getitem__, which reads self._models and never materializes a pending head segment — so a nested path fails cold if the parent hasn't been touched yet.
This bites a fairly common shape: a model that exposes derived children through its own __getitem__, so registry["/grp/parent/child"] resolves child via parent.__getitem__. If parent is still pending you get KeyError: No sub-registry found by the name 'parent', even though registry["/grp/parent"] alone resolves fine (and after that, .../parent/child works too). I hit this with a small repro: cold registry["/grp/parent/derived"] raises, but touching /grp/parent first makes it pass — so it's specifically the pending-head path.
Splitting off the leading segment and materializing it before delegating fixes it:
def __getitem__(self, item) -> ModelType:
if REGISTRY_SEPARATOR in item:
head, _, _ = item.partition(REGISTRY_SEPARATOR)
if head != "" and head in self._pending:
self._materialize(head)
return super().__getitem__(item)
return self._materialize(item) if item in self._pending else super().__getitem__(item)The head != "" guard keeps the absolute /...-from-root behavior intact. Worth a regression test that does the nested lookup on a freshly loaded registry.
de7f4da to
06586d8
Compare
Signed-off-by: Tim Paine <3105306+timkpaine@users.noreply.github.com>
06586d8 to
03eac9c
Compare
ptomecek
left a comment
There was a problem hiding this comment.
All the issues from my earlier review have been addressed:
- Cross-registry deadlock is resolved (the lock is no longer held across instantiate; coordination lock + cross-thread wait-graph). Verified my original no-cycle concurrent repro now completes.
- Cycle errors report the full cross-registry path and surface as a clean top-level RegistryKeyError; missing refs do too.
- Nested-path lookups now materialize a pending head segment, so derived sub-paths (models exposing children via getitem) resolve cold. Verified with a repro; regression test added.
- contains/getitem handle dotted paths consistently.
- LazyRegistry is now documented in the explanation and reference pages.
Full TestLazyRegistry suite passes (18 tests). Nice work.
No description provided.