diff --git a/src/Auto3D/registry.py b/src/Auto3D/registry.py new file mode 100644 index 0000000..459e7ce --- /dev/null +++ b/src/Auto3D/registry.py @@ -0,0 +1,124 @@ +"""A named-backend registry, with no knowledge of what is registered. + +Auto3D has two families of swappable backend -- neural network potentials and +isomer engines -- and until now each carried its own bespoke lookup: a dict plus +an if-chain for models, a tuple plus an if/elif ladder for isomer engines. Adding +a model backend meant editing five places and an isomer backend six, and one of +those places was the *presentation* layer, where `ENGINE_INFO` kept a parallel +table of display metadata that nothing checked against the set of real backends. +A backend registered without an entry there simply stopped appearing in +`auto3d models info`. + +This module is the shared half of both. It deliberately does **not** own +construction: a model adapter is built with ``(device, compile_model)`` and an +isomer engine with eight-odd keyword arguments, so a signature both satisfy would +be a bag of keywords that hides what each backend needs and stops a type checker +helping. Each factory resolves through the registry and then calls its own +constructor. + +It is also not a plugin system. Auto3D already accepts a third-party model as a +file path -- ``--engine /path/to/my_nnp.pt``, checked against the ``CustomNNP`` +contract at load -- so entry-point discovery would add only *named, installable* +backends, and would mean freezing ``ModelAdapter`` as a public interface while it +is still gaining members. A registry is what such a loader would populate, so +this does not foreclose it. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Generic, TypeVar + +from Auto3D.exceptions import ConfigurationError + +T = TypeVar("T") + + +@dataclass(frozen=True) +class Entry(Generic[T]): + """One registered backend: the thing itself, plus how to talk about it. + + ``info`` is free-form per-registry metadata -- for models it is what + ``auto3d models info`` prints. It lives here rather than in a table beside + the registry because a second table keyed by the same names is exactly the + drift this module exists to remove. + """ + + name: str + value: T + aliases: tuple[str, ...] = () + info: Any = None + + +class Registry(Generic[T]): + """Names to backends, with one place that explains an unknown name. + + Args: + kind: What is being registered, used in error messages -- "optimizing + engine", "isomer engine". Phrased as a noun so the message reads as + a sentence. + case_insensitive: Fold names before lookup. Models resolve + case-insensitively (``--engine ani2x`` works) while isomer engine + types are exact lowercase; both behaviors are current and tested, so + the difference is configuration here rather than something this + module decides. + """ + + def __init__(self, kind: str, *, case_insensitive: bool = False) -> None: + self._kind = kind + self._case_insensitive = case_insensitive + self._entries: dict[str, Entry[T]] = {} + self._lookup: dict[str, str] = {} + + def _key(self, name: str) -> str: + return name.upper() if self._case_insensitive else name + + def register( + self, + name: str, + value: T, + *, + aliases: tuple[str, ...] = (), + info: Any = None, + ) -> None: + """Add a backend. Registering a name twice is an error, not an overwrite. + + A plain dict silently replaces, which turns a duplicate -- two modules + registering the same name, or one imported twice under different paths -- + into a backend that works or does not depending on import order. + """ + for candidate in (name, *aliases): + key = self._key(candidate) + if key in self._lookup: + raise ConfigurationError( + f"{self._kind} {candidate!r} is already registered " + f"(as {self._lookup[key]!r}); names must be unique." + ) + self._entries[name] = Entry(name=name, value=value, aliases=aliases, info=info) + for candidate in (name, *aliases): + self._lookup[self._key(candidate)] = name + + def entry(self, name: str) -> Entry[T]: + """The full :class:`Entry`, for a caller that needs ``info`` too.""" + key = self._key(name) + if key not in self._lookup: + raise ConfigurationError( + f"Unknown {self._kind} {name!r}. Available: " + + ", ".join(repr(n) for n in self.available()) + + "." + ) + return self._entries[self._lookup[key]] + + def resolve(self, name: str) -> T: + """The registered value, or a ``ConfigurationError`` naming the alternatives.""" + return self.entry(name).value + + def available(self) -> list[str]: + """Registered names, canonical spelling only, in registration order.""" + return list(self._entries) + + def __contains__(self, name: str) -> bool: + return self._key(name) in self._lookup + + def __len__(self) -> int: + return len(self._entries) diff --git a/tests/test_layer_boundaries.py b/tests/test_layer_boundaries.py index 71eea47..e830488 100644 --- a/tests/test_layer_boundaries.py +++ b/tests/test_layer_boundaries.py @@ -16,8 +16,8 @@ L2 engines models/**, model_factory, isomers/**, isomer_engine, tautomer, batch_opt/** L1 domain ranking, filtering, embedding, clash_relief, id_mapping - L0 foundation config, constants, exceptions, results, torch_config, - utils/** + L0 foundation config, constants, exceptions, registry, results, + torch_config, utils/** This map describes the package as it is today, at its current module names -- not the target layout in the plan, which renames and regroups. It is the thing that @@ -50,7 +50,7 @@ 3: ["workflow", "workflow_workers", "chunk_manager", "job_layout", "processors", "pipeline"], 2: ["models", "model_factory", "isomers", "isomer_engine", "tautomer", "batch_opt"], 1: ["ranking", "filtering", "embedding", "clash_relief", "id_mapping"], - 0: ["config", "constants", "exceptions", "results", "torch_config", "utils"], + 0: ["config", "constants", "exceptions", "registry", "results", "torch_config", "utils"], } #: Upward module-scope edges that are known, named and dated. An entry here is a diff --git a/tests/test_registry.py b/tests/test_registry.py new file mode 100644 index 0000000..9610b62 --- /dev/null +++ b/tests/test_registry.py @@ -0,0 +1,108 @@ +"""The shared backend registry.""" + +from __future__ import annotations + +import pytest + +from Auto3D.exceptions import ConfigurationError +from Auto3D.registry import Registry + + +def test_register_and_resolve(): + reg: Registry[str] = Registry("widget") + reg.register("alpha", "A") + assert reg.resolve("alpha") == "A" + assert "alpha" in reg + assert reg.available() == ["alpha"] + + +def test_aliases_resolve_to_the_same_value_but_are_not_listed(): + """An alias is a second spelling, not a second backend. + + ``available()`` is what the CLI shows and what error messages enumerate, so + listing aliases there would advertise two engines where one exists. + """ + reg: Registry[str] = Registry("widget") + reg.register("alpha", "A", aliases=("a", "first")) + assert reg.resolve("a") == reg.resolve("first") == "A" + assert reg.available() == ["alpha"] + assert "first" in reg + + +def test_unknown_name_names_the_alternatives(): + """The error has to say what *would* have worked. + + Two hand-written messages did this before, with different wording and only + one of them quoting the alternatives. One implementation means one message. + """ + reg: Registry[str] = Registry("optimizing engine") + reg.register("alpha", "A") + reg.register("beta", "B") + with pytest.raises(ConfigurationError) as exc: + reg.resolve("gamma") + message = str(exc.value) + assert "optimizing engine" in message + assert "'gamma'" in message + assert "'alpha'" in message and "'beta'" in message + + +def test_duplicate_registration_raises_instead_of_overwriting(): + """A dict would silently replace, making behavior depend on import order.""" + reg: Registry[str] = Registry("widget") + reg.register("alpha", "A") + with pytest.raises(ConfigurationError, match="already registered"): + reg.register("alpha", "B") + assert reg.resolve("alpha") == "A" + + +def test_duplicate_is_caught_across_name_and_alias(): + reg: Registry[str] = Registry("widget") + reg.register("alpha", "A", aliases=("shared",)) + with pytest.raises(ConfigurationError, match="already registered"): + reg.register("beta", "B", aliases=("shared",)) + + +def test_case_sensitivity_is_configurable_because_both_kinds_exist(): + """Models resolve case-insensitively; isomer engine types do not. + + Both are current, tested behavior -- ``--engine ani2x`` works, while + ``rdkit_sdf`` is matched exactly -- so this is configuration rather than a + policy this module picks. + """ + folding: Registry[str] = Registry("engine", case_insensitive=True) + folding.register("ANI2x", "A") + assert folding.resolve("ani2x") == folding.resolve("ANI2X") == "A" + + exact: Registry[str] = Registry("engine") + exact.register("rdkit", "R") + assert exact.resolve("rdkit") == "R" + with pytest.raises(ConfigurationError): + exact.resolve("RDKit") + + +def test_case_folding_also_applies_to_duplicate_detection(): + reg: Registry[str] = Registry("engine", case_insensitive=True) + reg.register("ANI2x", "A") + with pytest.raises(ConfigurationError, match="already registered"): + reg.register("ani2x", "B") + + +def test_entry_carries_display_metadata(): + """``info`` is why the CLI stops keeping a parallel table. + + ``ENGINE_INFO`` was keyed by engine name and checked against nothing, so a + backend registered without an entry vanished from ``auto3d models info`` + with no error anywhere. + """ + reg: Registry[str] = Registry("engine") + reg.register("alpha", "A", info={"description": "the first one"}) + assert reg.entry("alpha").info == {"description": "the first one"} + assert reg.entry("alpha").value == "A" + + +def test_available_is_registration_order(): + """Not sorted: the order backends are declared in is the order shown.""" + reg: Registry[str] = Registry("engine") + for name in ("gamma", "alpha", "beta"): + reg.register(name, name.upper()) + assert reg.available() == ["gamma", "alpha", "beta"]