Skip to content

Add lazy model registry - #256

Merged
timkpaine merged 1 commit into
mainfrom
tkp/lazy-registry
Aug 19, 2026
Merged

Add lazy model registry#256
timkpaine merged 1 commit into
mainfrom
tkp/lazy-registry

Conversation

@timkpaine

Copy link
Copy Markdown
Member

No description provided.

@github-actions

github-actions Bot commented Aug 19, 2026

Copy link
Copy Markdown
Contributor

Test Results

    1 files  ± 0      1 suites  ±0   3m 7s ⏱️ -2s
1 327 tests +18  1 325 ✅ +18  2 💤 ±0  0 ❌ ±0 
1 333 runs  +18  1 331 ✅ +18  2 💤 ±0  0 ❌ ±0 

Results for commit 03eac9c. ± Comparison against base commit 3cb9493.

♻️ This comment has been updated with latest results.

@codecov

codecov Bot commented Aug 19, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 85.86066% with 69 lines in your changes missing coverage. Please review.
✅ Project coverage is 93.47%. Comparing base (3cb9493) to head (03eac9c).

Files with missing lines Patch % Lines
ccflow/base.py 78.21% 36 Missing and 30 partials ⚠️
ccflow/tests/test_base_registry.py 98.37% 1 Missing and 2 partials ⚠️
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.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@timkpaine
timkpaine marked this pull request as ready for review August 19, 2026 19:25

@ptomecek ptomecek left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.md has a "The Model Registry" section that walks through ModelRegistry, sub-registries, and name-based references. That's the natural home for lazy loading. Worth covering: the _recursive_: false requirement 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.md has a table of the notable classes (LazyEvaluator, etc.); LazyRegistry should 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.

Comment thread ccflow/base.py
def _materialize(self, name: str) -> BaseModel:
from hydra.utils import instantiate

with self._lock:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread ccflow/base.py Outdated
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}")

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

_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.

Comment thread ccflow/base.py
search_registry = search_registries[idx]
try:
return search_registry[v]
except RegistryKeyError:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread ccflow/base.py
if not isinstance(item, str):
return False
if REGISTRY_SEPARATOR in item:
if "." in item:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread ccflow/base.py
return self

def __getitem__(self, item) -> ModelType:
if REGISTRY_SEPARATOR in item:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Signed-off-by: Tim Paine <3105306+timkpaine@users.noreply.github.com>

@ptomecek ptomecek left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@timkpaine
timkpaine merged commit 1a186fd into main Aug 19, 2026
18 of 20 checks passed
@timkpaine
timkpaine deleted the tkp/lazy-registry branch August 19, 2026 23:24
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.

2 participants