diff --git a/ccflow/base.py b/ccflow/base.py index 5df0a3c..7a831eb 100644 --- a/ccflow/base.py +++ b/ccflow/base.py @@ -1,14 +1,16 @@ """This module defines the base model and registry for flow.""" import collections.abc +import contextvars import copy import inspect import logging import pathlib import platform import sys +import threading import warnings -from collections.abc import Callable +from collections.abc import Callable, Iterator, Mapping from types import GenericAlias, MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Generic, Optional, TypeVar, Union, get_args, get_origin @@ -43,6 +45,7 @@ "BaseModel", "ContextBase", "ContextType", + "LazyRegistry", "ModelRegistry", "ModelType", "RegistryLookupContext", @@ -61,6 +64,47 @@ class RegistryKeyError(KeyError): """Subclass for KeyError specific to Registry lookup errors.""" +class _LazyMaterialization: + def __init__(self, owner: int, registry: "LazyRegistry", name: str): + self.owner = owner + self.registry = registry + self.name = name + self.event = threading.Event() + self.error: BaseException | None = None + + +_LAZY_LOADING_STACK: contextvars.ContextVar[tuple[tuple["LazyRegistry", str], ...]] = contextvars.ContextVar("ccflow_lazy_loading_stack", default=()) +_LAZY_COORDINATION_LOCK = threading.RLock() +_LAZY_MATERIALIZATIONS: dict[tuple[int, str], _LazyMaterialization] = {} +_LAZY_THREAD_WAITS: dict[int, tuple[int, str]] = {} + + +def _lazy_entry_path(registry: "LazyRegistry", name: str) -> str: + prefix = registry._debug_name.rstrip(REGISTRY_SEPARATOR) + return f"{prefix}{REGISTRY_SEPARATOR}{name}" if prefix else f"{REGISTRY_SEPARATOR}{name}" + + +def _registry_key_error_from_exception(error: BaseException) -> RegistryKeyError | None: + pending = [error] + seen = set() + while pending: + current = pending.pop() + if id(current) in seen: + continue + seen.add(id(current)) + if isinstance(current, RegistryKeyError): + return current + if isinstance(current, ValidationError): + for detail in current.errors(): + context_error = detail.get("ctx", {}).get("error") + if isinstance(context_error, BaseException): + pending.append(context_error) + for nested in (current.__cause__, current.__context__): + if nested is not None: + pending.append(nested) + return None + + class _RegistryMixin: def get_registrations(self) -> list[tuple["ModelRegistry", str]]: """Return the set of registrations that has happened for this model""" @@ -584,6 +628,20 @@ def __getitem__(self, item) -> ModelType: else: raise KeyError(f"No registered model found by the name '{item}' in registry '{self._debug_name}'") + def __contains__(self, item: object) -> bool: + if not isinstance(item, str): + return False + if REGISTRY_SEPARATOR in item: + if "." in item: + raise ValueError("Path references to registry objects do not support '.' or '..'") + registry_name, name = item.split(REGISTRY_SEPARATOR, 1) + if registry_name == "": + registry = ModelRegistry.root() + else: + registry = self._models.get(registry_name) + return isinstance(registry, ModelRegistry) and name in registry + return item in self._models + def __iter__(self): for key, model in self._models.items(): yield key @@ -673,6 +731,317 @@ def load_config_from_path( return self.load_config(cfg, overwrite=overwrite) +class _LazyRegistryModels(Mapping[str, BaseModel]): + """Read-only direct-child view of a lazy registry.""" + + def __init__(self, registry: "LazyRegistry"): + self._registry = registry + + def __getitem__(self, name: str) -> BaseModel: + if REGISTRY_SEPARATOR in name: + raise KeyError(name) + return self._registry[name] + + def __iter__(self) -> Iterator[str]: + return iter(tuple(self._registry._entry_order)) + + def __len__(self) -> int: + return len(self._registry._entry_order) + + def __contains__(self, name: object) -> bool: + return isinstance(name, str) and REGISTRY_SEPARATOR not in name and name in self._registry._entry_order + + +class LazyRegistry(ModelRegistry): + """Registry that instantiates configured models when they are first accessed.""" + + name: str = "" + _pending: dict[str, Any] = PrivateAttr(default_factory=dict) + _pending_lookup_registries: dict[str, list[ModelRegistry]] = PrivateAttr(default_factory=dict) + _entry_order: list[str] = PrivateAttr(default_factory=list) + _lookup_registries: list[ModelRegistry] = PrivateAttr(default_factory=list) + _lock: Any = PrivateAttr(default_factory=threading.RLock) + + @field_validator("name") + @classmethod + def _validate_name(cls, value): + return value + + def __init__(self, name: str = "", **config): + super().__init__(name=name) + lookup_registries = list(RegistryLookupContext.registry_search_paths() or []) + self._lookup_registries = [*lookup_registries, self] if lookup_registries else [self] + self._load_config(config, overwrite=False, lookup_registries=self._lookup_registries) + + def __eq__(self, other: object) -> bool: + if not isinstance(other, LazyRegistry) or self.name != other.name or self._entry_order != other._entry_order: + return False + self.materialize_all() + other.materialize_all() + return self._models == other._models + + @model_serializer(mode="wrap") + def _lazy_registry_serializer(self, handler): + values = handler(self) + values["_recursive_"] = False + values.update(self._config_tree()) + return values + + def _config_tree(self) -> dict[str, Any]: + values = {} + for name in self._entry_order: + if name in self._models: + model = self._models[name] + values[name] = model._config_tree() if isinstance(model, LazyRegistry) else model.model_dump(by_alias=True) + else: + values[name] = copy.deepcopy(self._pending[name]) + return values + + @property + def models(self) -> Mapping[str, BaseModel]: + """Return a read-only view that materializes values only when accessed.""" + return _LazyRegistryModels(self) + + def is_loaded(self, name: str) -> bool: + """Return whether a direct child has been instantiated.""" + return name in self._models + + def get_loaded(self, name: str) -> BaseModel | None: + """Return an instantiated direct child without loading a pending entry.""" + return self._models.get(name) + + def get_pending_config(self, name: str) -> Mapping[str, Any] | None: + """Return a pending direct child's configuration without instantiating it.""" + if name in self._models or name not in self._pending: + return None + return MappingProxyType(copy.deepcopy(self._pending[name])) + + def clone(self, name: str | None = None) -> Self: + """Shallow-clone this registry without instantiating pending entries.""" + cloned = type(self)(name=name or self.name) + cloned._lookup_registries = [cloned if registry is self else registry for registry in self._lookup_registries] + for entry_name in self._entry_order: + if entry_name in self._models: + cloned.add(entry_name, self._models[entry_name]) + else: + lookup_registries = [cloned if registry is self else registry for registry in self._pending_lookup_registries[entry_name]] + cloned._add_pending(entry_name, copy.deepcopy(self._pending[entry_name]), False, lookup_registries) + return cloned + + def __getstate__(self): + state = super().__getstate__() + private = dict(state.get("__pydantic_private__") or {}) + private.pop("_lock", None) + state["__pydantic_private__"] = private + return state + + def __setstate__(self, state): + super().__setstate__(state) + self._lock = threading.RLock() + + def _load_config(self, cfg, overwrite: bool, lookup_registries: list[ModelRegistry]) -> None: + from omegaconf import DictConfig, OmegaConf + + if isinstance(cfg, DictConfig): + cfg = OmegaConf.to_container(cfg, resolve=True) + for key, value in cfg.items(): + if isinstance(value, BaseModel): + raise TypeError( + f"Lazy registry entry '{key}' was instantiated before reaching LazyRegistry. " + "Set '_recursive_: false' on the LazyRegistry Hydra config." + ) + if not isinstance(value, (dict, DictConfig)): + continue + if isinstance(value, DictConfig): + value = OmegaConf.to_container(value, resolve=True) + if _is_config_model(value): + self._add_pending(key, value, overwrite, lookup_registries) + elif _is_config_subregistry(value): + child = LazyRegistry(name=key) + child._lookup_registries = [*lookup_registries, child] + child._load_config(value, overwrite=False, lookup_registries=child._lookup_registries) + self.add(key, child, overwrite=overwrite) + + def _add_pending(self, name: str, cfg: Any, overwrite: bool, lookup_registries: list[ModelRegistry]) -> None: + if REGISTRY_SEPARATOR in name: + raise ValueError(f"Cannot add '{name}' to '{self._debug_name}' because it contains '{REGISTRY_SEPARATOR}'") + if name in self._entry_order and not overwrite: + raise ValueError(f"Cannot add '{name}' to '{self._debug_name}' as it already exists!") + with _LAZY_COORDINATION_LOCK: + if (id(self), name) in _LAZY_MATERIALIZATIONS: + raise RuntimeError(f"Cannot replace '{name}' in '{self._debug_name}' while it is being materialized") + if name in self._models: + ModelRegistry.remove(self, name) + if name not in self._entry_order: + self._entry_order.append(name) + self._pending[name] = cfg + self._pending_lookup_registries[name] = list(lookup_registries) + + def _materialize(self, name: str) -> BaseModel: + from hydra.utils import instantiate + + key = (id(self), name) + stack = _LAZY_LOADING_STACK.get() + for index, entry in enumerate(stack): + if (id(entry[0]), entry[1]) == key: + cycle = [*stack[index:], (self, name)] + paths = " -> ".join(_lazy_entry_path(registry, entry_name) for registry, entry_name in cycle) + raise RegistryKeyError(f"Circular lazy registry dependency detected: {paths}") + + with self._lock: + if name in self._models: + return self._models[name] + if name not in self._pending: + raise KeyError(f"No registered model found by the name '{name}' in registry '{self._debug_name}'") + cfg = self._pending[name] + lookup_registries = self._pending_lookup_registries[name] + + owner = threading.get_ident() + with _LAZY_COORDINATION_LOCK: + materialization = _LAZY_MATERIALIZATIONS.get(key) + if materialization is None: + materialization = _LazyMaterialization(owner, self, name) + _LAZY_MATERIALIZATIONS[key] = materialization + owns_materialization = True + else: + owns_materialization = False + _LAZY_THREAD_WAITS[owner] = key + waiting_owner = materialization.owner + wait_chain = [key] + while waiting_owner in _LAZY_THREAD_WAITS: + waiting_key = _LAZY_THREAD_WAITS[waiting_owner] + wait_chain.append(waiting_key) + waiting_materialization = _LAZY_MATERIALIZATIONS.get(waiting_key) + if waiting_materialization is None: + break + waiting_owner = waiting_materialization.owner + if waiting_owner == owner: + _LAZY_THREAD_WAITS.pop(owner, None) + cycle = [*stack, *((entry.registry, entry.name) for entry in (materialization,))] + cycle.extend( + (_LAZY_MATERIALIZATIONS[waiting_key].registry, _LAZY_MATERIALIZATIONS[waiting_key].name) + for waiting_key in wait_chain[1:] + if waiting_key in _LAZY_MATERIALIZATIONS + ) + paths = " -> ".join(_lazy_entry_path(registry, entry_name) for registry, entry_name in cycle) + raise RegistryKeyError(f"Circular lazy registry dependency detected: {paths}") + + if not owns_materialization: + materialization.event.wait() + with _LAZY_COORDINATION_LOCK: + _LAZY_THREAD_WAITS.pop(owner, None) + if materialization.error is not None: + raise materialization.error + return self._materialize(name) + + token = _LAZY_LOADING_STACK.set((*stack, (self, name))) + try: + with RegistryLookupContext(registries=lookup_registries): + model = instantiate(cfg, _convert_="all") + if not isinstance(model, BaseModel): + raise TypeError(f"Lazy registry entry '{name}' produced '{type(model)}', not a child class of {BaseModel}.") + if hasattr(model, "meta") and hasattr(model.meta, "name") and model.meta.name == "": + model.meta.name = name + with self._lock: + ModelRegistry.add(self, name, model) + del self._pending[name] + del self._pending_lookup_registries[name] + return model + except BaseException as error: + registry_error = _registry_key_error_from_exception(error) + materialization.error = registry_error or error + if registry_error is not None: + raise registry_error from error + raise + finally: + _LAZY_LOADING_STACK.reset(token) + with _LAZY_COORDINATION_LOCK: + _LAZY_MATERIALIZATIONS.pop(key, None) + materialization.event.set() + + def add(self, name: str, model: ModelType, overwrite: bool = False) -> ModelType: + with self._lock: + if name in self._pending and not overwrite: + raise ValueError(f"Cannot add '{name}' to '{self._debug_name}' as it already exists!") + result = super().add(name, model, overwrite=overwrite) + if name not in self._entry_order: + self._entry_order.append(name) + if name in self._pending: + del self._pending[name] + del self._pending_lookup_registries[name] + return result + + def remove(self, name: str) -> None: + with self._lock: + if name in self._models: + super().remove(name) + self._pending.pop(name, None) + self._pending_lookup_registries.pop(name, None) + elif name in self._pending: + del self._pending[name] + del self._pending_lookup_registries[name] + else: + raise ValueError(f"Cannot remove '{name}' from '{self._debug_name}' as it does not exist there!") + self._entry_order.remove(name) + + def clear(self) -> Self: + for name in list(self._entry_order): + self.remove(name) + return self + + 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) + + def __contains__(self, item: object) -> bool: + if not isinstance(item, str): + return False + if REGISTRY_SEPARATOR in item: + return super().__contains__(item) + return item in self._entry_order + + def __iter__(self): + for key in tuple(self._entry_order): + yield key + model = self._models.get(key) + if isinstance(model, ModelRegistry): + for subkey in model: + yield REGISTRY_SEPARATOR.join((key, subkey)) + + def __len__(self) -> int: + count = len(self._entry_order) + for model in self._models.values(): + if isinstance(model, ModelRegistry): + count += len(model) + return count + + def load_config( + self, + cfg: "DictConfig", + overwrite: bool = False, + skip_exceptions: bool = False, + resolve_from: ModelRegistry | None = None, + ) -> Self: + if skip_exceptions: + raise ValueError("skip_exceptions is not supported by LazyRegistry") + lookup_registries = [resolve_from, self] if resolve_from is not None and resolve_from is not self else self._lookup_registries + with self._lock: + self._load_config(cfg, overwrite=overwrite, lookup_registries=lookup_registries) + return self + + def materialize_all(self) -> Self: + """Instantiate every pending model in this registry tree.""" + for name in tuple(self._entry_order): + model = self[name] + if isinstance(model, LazyRegistry): + model.materialize_all() + return self + + class RootModelRegistry(ModelRegistry): """ Class to represent the singleton, i.e. "root" ModelRegistry, @@ -771,6 +1140,8 @@ def load_config( if hasattr(model, "meta") and hasattr(model.meta, "name") and model.meta.name == "": model.meta.name = k + if isinstance(model, LazyRegistry) and not model.name: + model.name = k registries[-1].add(k, model, overwrite=self._overwrite) if not unresolved_models: @@ -817,7 +1188,7 @@ class RegistryLookupContext: Do not confuse the name with "Context" from callable.py. """ - _REGISTRIES: ClassVar[list[ModelRegistry]] = [] + _REGISTRIES: ClassVar[contextvars.ContextVar[tuple[ModelRegistry, ...]]] = contextvars.ContextVar("ccflow_registry_lookup_context", default=()) def __init__(self, registries: list[ModelRegistry] | None = None): """Constructor. @@ -827,19 +1198,19 @@ def __init__(self, registries: list[ModelRegistry] | None = None): this context. """ self.registries = registries - self._previous_registries = [] + self._token = None @classmethod def registry_search_paths(cls) -> list[ModelRegistry]: """Return the active list of additional registry search paths.""" - return cls._REGISTRIES + return list(cls._REGISTRIES.get()) def __enter__(self): - self._previous_registries = self._REGISTRIES - RegistryLookupContext._REGISTRIES = self.registries + self._token = self._REGISTRIES.set(tuple(self.registries or ())) def __exit__(self, exc_type, exc_value, exc_tb): - RegistryLookupContext._REGISTRIES = self._previous_registries + if self._token is not None: + self._REGISTRIES.reset(self._token) def resolve_str(v: str) -> ModelType: @@ -864,6 +1235,8 @@ def resolve_str(v: str) -> ModelType: search_registry = search_registries[idx] try: return search_registry[v] + except RegistryKeyError: + raise except KeyError: # A common mistake is to forget to start an absolute lookup with a forward slash. Return a better error message in that case. if not original_v.startswith("/"): diff --git a/ccflow/tests/test_base_registry.py b/ccflow/tests/test_base_registry.py index 36e8288..c92059d 100644 --- a/ccflow/tests/test_base_registry.py +++ b/ccflow/tests/test_base_registry.py @@ -1,9 +1,12 @@ import collections.abc +import concurrent.futures import json import os import pickle import sys -from unittest import TestCase +import threading +from typing import ClassVar +from unittest import TestCase, mock import pytest from hydra.errors import InstantiationException @@ -11,7 +14,7 @@ from omegaconf.errors import InterpolationKeyError from pydantic import ConfigDict, Field -from ccflow import BaseModel, ModelRegistry, RegistryLookupContext, RootModelRegistry, model_alias +from ccflow import BaseModel, LazyRegistry, ModelRegistry, RegistryLookupContext, RootModelRegistry, model_alias from ccflow.base import RegistryKeyError, resolve_str @@ -30,6 +33,32 @@ class MyTestModelSubclass(MyTestModel): pass +class LazyTestModel(MyTestModel): + constructions: ClassVar[int] = 0 + + def __init__(self, **data): + type(self).constructions += 1 + super().__init__(**data) + + +class ConcurrentLazyTestModel(MyTestModel): + barrier: ClassVar[threading.Barrier | None] = None + + def __init__(self, **data): + if self.barrier is not None: + self.barrier.wait(timeout=2) + super().__init__(**data) + + +class IndexableLazyTestModel(BaseModel): + child: MyTestModel + + def __getitem__(self, name: str) -> MyTestModel: + if name != "child": + raise KeyError(name) + return self.child + + class MyClass: def __init__(self, p="p", q=10.0): self.p = p @@ -663,6 +692,363 @@ def test_load_config_resolve_from_self_noop(self): self.assertEqual(r["foo"], MyTestModel(a="test", b=0.0)) +class TestLazyRegistry(TestCase): + def setUp(self) -> None: + ModelRegistry.root().clear() + LazyTestModel.constructions = 0 + + def tearDown(self) -> None: + ModelRegistry.root().clear() + + @staticmethod + def _load_registry(): + cfg = OmegaConf.create( + { + "lazy": { + "_target_": "ccflow.base.LazyRegistry", + "_recursive_": False, + "group": { + "source": { + "_target_": "ccflow.tests.test_base_registry.LazyTestModel", + "a": "source", + "b": 1.0, + }, + "dependent": { + "_target_": "ccflow.tests.test_base_registry.MyNestedModel", + "x": "/lazy/group/source", + "y": "source", + }, + }, + } + } + ) + root = ModelRegistry.root() + root.load_config(cfg) + return root + + def test_defers_models_until_access(self): + root = self._load_registry() + lazy = root["lazy"] + + self.assertIsInstance(lazy, LazyRegistry) + self.assertEqual(lazy.name, "lazy") + self.assertEqual(LazyTestModel.constructions, 0) + self.assertIn("lazy/group/source", root) + self.assertIn("group/source", lazy) + self.assertIn("source", lazy["group"].models) + self.assertIn("lazy/group/source", list(root.keys())) + self.assertEqual(LazyTestModel.constructions, 0) + + source = root["/lazy/group/source"] + self.assertIsInstance(source, LazyTestModel) + self.assertEqual(LazyTestModel.constructions, 1) + self.assertIs(root["/lazy/group/source"], source) + self.assertEqual(LazyTestModel.constructions, 1) + + def test_materializes_dependency_closure(self): + root = self._load_registry() + + dependent = root["/lazy/group/dependent"] + + self.assertEqual(LazyTestModel.constructions, 1) + self.assertIs(dependent.x, root["/lazy/group/source"]) + self.assertIs(dependent.y, root["/lazy/group/source"]) + + def test_nested_path_materializes_pending_head(self): + root = ModelRegistry.root() + root.load_config( + OmegaConf.create( + { + "lazy": { + "_target_": "ccflow.base.LazyRegistry", + "_recursive_": False, + "parent": { + "_target_": "ccflow.tests.test_base_registry.IndexableLazyTestModel", + "child": { + "_target_": "ccflow.tests.test_base_registry.MyTestModel", + "a": "child", + "b": 1.0, + }, + }, + } + } + ) + ) + lazy = root["lazy"] + + child = root["/lazy/parent/child"] + + self.assertEqual(child.a, "child") + self.assertTrue(lazy.is_loaded("parent")) + + def test_pending_entries_support_mutation(self): + registry = LazyRegistry( + name="lazy", + pending={ + "_target_": "ccflow.tests.test_base_registry.LazyTestModel", + "a": "pending", + "b": 1.0, + }, + ) + replacement = MyTestModel(a="replacement", b=2.0) + + with self.assertRaises(ValueError): + registry.add("pending", replacement) + registry.add("pending", replacement, overwrite=True) + self.assertIs(registry["pending"], replacement) + self.assertEqual(LazyTestModel.constructions, 0) + + registry.load_config( + OmegaConf.create( + { + "other": { + "_target_": "ccflow.tests.test_base_registry.LazyTestModel", + "a": "other", + "b": 3.0, + } + } + ) + ) + registry.remove("other") + self.assertNotIn("other", registry) + self.assertEqual(LazyTestModel.constructions, 0) + + def test_cycle_reports_dependency_chain(self): + root = ModelRegistry.root() + root.load_config( + OmegaConf.create( + { + "lazy": { + "_target_": "ccflow.base.LazyRegistry", + "_recursive_": False, + "a": { + "_target_": "ccflow.tests.test_base_registry.MyNestedModel", + "x": "b", + "y": "b", + }, + "b": { + "_target_": "ccflow.tests.test_base_registry.MyNestedModel", + "x": "a", + "y": "a", + }, + } + } + ) + ) + + with self.assertRaisesRegex(RegistryKeyError, "Circular lazy registry dependency detected: /lazy/a -> /lazy/b -> /lazy/a"): + root["/lazy/a"] + + def test_cross_registry_cycle_reports_full_path(self): + root = ModelRegistry.root() + left = LazyRegistry(name="left") + right = LazyRegistry(name="right") + root.add("left", left) + root.add("right", right) + left.load_config( + OmegaConf.create({"model": {"_target_": "ccflow.tests.test_base_registry.MyNestedModel", "x": "/right/model", "y": "/right/model"}}), + resolve_from=root, + ) + right.load_config( + OmegaConf.create({"model": {"_target_": "ccflow.tests.test_base_registry.MyNestedModel", "x": "/left/model", "y": "/left/model"}}), + resolve_from=root, + ) + + with self.assertRaisesRegex( + RegistryKeyError, + "Circular lazy registry dependency detected: /left/model -> /right/model -> /left/model", + ): + left["model"] + + def test_concurrent_cross_registry_cycle_does_not_deadlock(self): + from hydra.utils import instantiate as hydra_instantiate + + root = ModelRegistry.root() + left = LazyRegistry(name="left") + right = LazyRegistry(name="right") + root.add("left", left) + root.add("right", right) + for registry, dependency in ((left, "/right/model"), (right, "/left/model")): + registry.load_config( + OmegaConf.create({"model": {"_target_": "ccflow.tests.test_base_registry.MyNestedModel", "x": dependency, "y": dependency}}), + resolve_from=root, + ) + + barrier = threading.Barrier(2) + local = threading.local() + + def synchronized_instantiate(*args, **kwargs): + if not getattr(local, "started", False): + local.started = True + barrier.wait(timeout=2) + return hydra_instantiate(*args, **kwargs) + + with ( + mock.patch("hydra.utils.instantiate", side_effect=synchronized_instantiate), + concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor, + ): + futures = [executor.submit(registry.__getitem__, "model") for registry in (left, right)] + messages = [] + for future in futures: + with self.assertRaises(RegistryKeyError) as context: + future.result(timeout=3) + messages.append(str(context.exception)) + + for message in messages: + self.assertIn("/left/model", message) + self.assertIn("/right/model", message) + + def test_missing_reference_raises_registry_key_error_directly(self): + registry = LazyRegistry( + name="lazy", + model={"_target_": "ccflow.tests.test_base_registry.MyNestedModel", "x": "missing", "y": "missing"}, + ) + + with self.assertRaisesRegex(RegistryKeyError, "Could not resolve model 'missing'"): + registry["model"] + + def test_contains_rejects_dotted_paths_consistently(self): + registry = LazyRegistry(name="lazy") + + with self.assertRaisesRegex(ValueError, "do not support"): + _ = "foo.bar/baz" in registry + + def test_serialization_round_trip_preserves_pending_entries(self): + registry = self._load_registry()["lazy"] + + dumped = registry.model_dump(by_alias=True) + + self.assertEqual(LazyTestModel.constructions, 0) + self.assertNotIn("models", dumped) + self.assertFalse(registry["group"].is_loaded("source")) + restored = LazyRegistry.model_validate(dumped) + self.assertEqual(list(restored), list(registry)) + self.assertFalse(restored["group"].is_loaded("source")) + + def test_clone_preserves_pending_entries(self): + registry = self._load_registry()["lazy"] + + cloned = registry.clone(name="clone") + + self.assertIsInstance(cloned, LazyRegistry) + self.assertEqual(cloned.name, "clone") + self.assertEqual(list(cloned), list(registry)) + self.assertFalse(registry["group"].is_loaded("source")) + self.assertFalse(cloned["group"].is_loaded("source")) + + def test_failed_overwrite_preserves_pending_entry(self): + registry = LazyRegistry( + name="lazy", + pending={ + "_target_": "ccflow.tests.test_base_registry.LazyTestModel", + "a": "pending", + "b": 1.0, + }, + ) + + with self.assertRaises(TypeError): + registry.add("pending", object(), overwrite=True) + + self.assertIn("pending", registry) + self.assertIsNotNone(registry.get_pending_config("pending")) + self.assertEqual(registry["pending"].a, "pending") + + def test_pending_config_is_read_only_copy(self): + registry = LazyRegistry( + name="lazy", + pending={ + "_target_": "ccflow.tests.test_base_registry.LazyTestModel", + "a": "pending", + "b": 1.0, + }, + ) + config = registry.get_pending_config("pending") + + with self.assertRaises(TypeError): + config["a"] = "changed" + + self.assertEqual(registry["pending"].a, "pending") + + def test_recursive_hydra_instantiation_is_rejected(self): + from hydra.utils import instantiate + + cfg = OmegaConf.create( + { + "_target_": "ccflow.base.LazyRegistry", + "name": "lazy", + "model": { + "_target_": "ccflow.tests.test_base_registry.LazyTestModel", + "a": "pending", + "b": 1.0, + }, + } + ) + + with self.assertRaisesRegex(InstantiationException, "Set '_recursive_: false'"): + instantiate(cfg, _convert_="all") + + def test_resolve_from_applies_to_direct_pending_entry(self): + root = ModelRegistry.root() + source = MyTestModel(a="source", b=1.0) + root.add("source", source) + registry = LazyRegistry(name="lazy") + registry.load_config( + OmegaConf.create( + { + "dependent": { + "_target_": "ccflow.tests.test_base_registry.MyNestedModel", + "x": "/source", + "y": "/source", + } + } + ), + resolve_from=root, + ) + + self.assertIs(registry["dependent"].x, source) + + def test_equality_is_independent_of_materialization_state(self): + left = LazyRegistry( + name="lazy", + model={"_target_": "ccflow.tests.test_base_registry.LazyTestModel", "a": "same", "b": 1.0}, + ) + right = LazyRegistry( + name="lazy", + model={"_target_": "ccflow.tests.test_base_registry.LazyTestModel", "a": "same", "b": 1.0}, + ) + + left["model"] + + self.assertEqual(left, right) + + def test_independent_registries_materialize_concurrently(self): + ConcurrentLazyTestModel.barrier = threading.Barrier(2) + left = LazyRegistry( + name="left", + model={"_target_": "ccflow.tests.test_base_registry.ConcurrentLazyTestModel", "a": "left", "b": 1.0}, + ) + right = LazyRegistry( + name="right", + model={"_target_": "ccflow.tests.test_base_registry.ConcurrentLazyTestModel", "a": "right", "b": 1.0}, + ) + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor: + futures = [executor.submit(registry.__getitem__, "model") for registry in (left, right)] + for future in futures: + future.result(timeout=3) + ConcurrentLazyTestModel.barrier = None + + def test_pickle_preserves_pending_entries_and_recreates_lock(self): + registry = LazyRegistry( + name="lazy", + model={"_target_": "ccflow.tests.test_base_registry.LazyTestModel", "a": "pending", "b": 1.0}, + ) + + restored = pickle.loads(pickle.dumps(registry)) + + self.assertFalse(restored.is_loaded("model")) + self.assertEqual(restored["model"].a, "pending") + + class TestRegistryLoadingErrors(TestCase): def setUp(self) -> None: ModelRegistry.root().clear() diff --git a/docs/wiki/explanation/Configuration-and-Hydra.md b/docs/wiki/explanation/Configuration-and-Hydra.md index a187f04..718007b 100644 --- a/docs/wiki/explanation/Configuration-and-Hydra.md +++ b/docs/wiki/explanation/Configuration-and-Hydra.md @@ -76,6 +76,23 @@ Hydra is excellent at composing files and instantiating objects, but on its own **The same power without files.** The registry provides the interactive half of the story. `ModelRegistry.load_config(...)` loads a dictionary of configs (resolving name references as it goes), and `load_config_from_path(...)` loads the same Hydra files a CLI would use — so a researcher can pull the production configuration into a notebook, inspect it, tweak an object, and re-run, all as live Python objects. +### Deferred registry construction + +`LazyRegistry` keeps composed model configuration without importing or constructing each model immediately. Hydra must pass nested targets through unchanged, so a lazy registry target includes `_recursive_: false`: + +```yaml +models: + _target_: ccflow.LazyRegistry + _recursive_: false + + source: + _target_: my_package.SourceModel +``` + +Accessing `registry["source"]` constructs and caches that model. References encountered during validation materialize their dependencies in the same way, so forward references do not require the eager loader's retry loop. `materialize_all()` constructs the complete tree when eager validation is useful, such as in configuration tests. + +Deferral changes when errors appear: missing references, dependency cycles, import failures, and validation errors surface when the affected entry is accessed rather than when the registry configuration is loaded. + ## The payoff Put together, the configuration story gives you: diff --git a/docs/wiki/reference/Built-in-Models.md b/docs/wiki/reference/Built-in-Models.md index a80a0f1..cc7902c 100644 --- a/docs/wiki/reference/Built-in-Models.md +++ b/docs/wiki/reference/Built-in-Models.md @@ -15,6 +15,13 @@ Beyond the `CallableModel` and `BaseModel` subclasses you write yourself, `ccflo Additional data-reading models are being open-sourced over time. +## Registries + +| Name | Path | Description | +| :-------------- | :------- | :----------------------------------------------------------------- | +| `ModelRegistry` | `ccflow` | Stores named models and nested registries. | +| `LazyRegistry` | `ccflow` | Stores model configuration and constructs each model on first use. | + ## Publishers Publishers (`ccflow.publishers`) are models that write or send data. A common interface lets one be substituted for another purely through configuration. See [Bind Logic to Configs](Bind-Logic-to-Configs#write-a-custom-publisher) to write your own.