From bdde437441d2e070d62e1a713ab43a12c539d1af Mon Sep 17 00:00:00 2001 From: isayev Date: Thu, 13 Aug 2026 17:27:37 -0400 Subject: [PATCH] refactor: resolve model engine names through the shared registry Item 5, step 2 -- the two collapsible lists of the three. Non-breaking: `available_models()` returns the same six names in the same order, and `_adapters` was private. Three parallel lists of engine names existed, holding *different* sets: ModelFactory._adapters 2 -- ANI2X, ANI2XT only ModelFactory.available_models() 6 -- a hand-written literal, derived from nothing ENGINE_INFO (cli/commands) 6 -- AIMNET, three aimnet variants, two ANI `_adapters` held only the engines built from a local adapter class; the four aimnet entries are names resolved through the aimnet registry, not adapters. So one list was keyed by "has a local adapter" and two by "user-facing name", and nothing connected them. `_engines` is now the one registry, keyed by user-facing name, with the adapter class as the value where there is one and `None` where the name is resolved through the aimnet registry. `available_models()` derives from it rather than restating it -- adding an engine and forgetting that list used to leave it working but undiscoverable. The import-time `assert set(_adapters) == set(BUILTIN_ANI_MODELS)` survives, now derived from the registry so the two cannot drift. It moved to module scope because a comprehension inside a class body cannot see that class's own attributes -- a scoping rule, not a style choice. `ENGINE_INFO` is deliberately still separate. Moving it into `info=` belongs with the CLI's changes rather than this factory's, and it is the remaining piece of item 5. Naming it here so the third list is not mistaken for handled. Suite: 1756 passed, 1 skipped, 70 deselected. --- src/Auto3D/model_factory.py | 63 +++++++++++++++++++++++++----------- tests/test_cli_exit_codes.py | 6 +++- tests/test_model_factory.py | 35 ++++++++++++++------ 3 files changed, 76 insertions(+), 28 deletions(-) diff --git a/src/Auto3D/model_factory.py b/src/Auto3D/model_factory.py index 39f133b..d3baafa 100644 --- a/src/Auto3D/model_factory.py +++ b/src/Auto3D/model_factory.py @@ -22,6 +22,7 @@ ANI2xtAdapter, CustomModelAdapter, ) +from Auto3D.registry import Registry if TYPE_CHECKING: # Annotation-only. Every signature here promises the CONTRACT @@ -57,13 +58,27 @@ class ModelFactory: >>> ModelFactory.clear_cache() """ - # Built-in (non-aimnet) engines kept for back-compat; keys are exactly the - # members of BUILTIN_ANI_MODELS. - _adapters: dict[str, type[BaseModelAdapter]] = { - MODEL_ANI2XT.upper(): ANI2xtAdapter, - MODEL_ANI2X.upper(): ANI2xAdapter, - } - assert set(_adapters) == set(BUILTIN_ANI_MODELS) + #: Every engine name a user may pass, in the order `auto3d models list` + #: shows them. The value is the adapter class for engines built from one + #: locally, and ``None`` for names resolved through the aimnet registry -- + #: which is the distinction that made this three separate lists before. + #: + #: `_adapters` held only the two ANI engines, `available_models()` restated + #: all six as a hand-written literal derived from nothing, and + #: `cli/commands/models.py`'s `ENGINE_INFO` held six display entries keyed by + #: the same names. Nothing connected them; they agreed by hand. Two of the + #: three are collapsed here. `ENGINE_INFO` is the remaining one, and moving + #: it into `info=` is the follow-up -- it belongs with the CLI changes rather + #: than with this factory's. + _engines: Registry[type[BaseModelAdapter] | None] = Registry( + "optimizing engine", case_insensitive=True + ) + _engines.register(MODEL_AIMNET, None) + _engines.register("aimnet2-2025", None) + _engines.register("aimnet2-nse", None) + _engines.register("aimnet2-pd", None) + _engines.register(MODEL_ANI2X, ANI2xAdapter) + _engines.register(MODEL_ANI2XT, ANI2xtAdapter) # Model instance cache: key = (name, device_str, compile_model) _cache: dict[tuple[str, str, bool], ModelAdapter] = {} @@ -141,12 +156,14 @@ def create( # name first; that helper is gone -- the Hessian path takes a # ModelAdapter now -- so this factory is the only name resolver # left, and the reason above stands on its own.) - if name_upper in cls._adapters: + if name_upper in cls._engines and cls._engines.resolve(name_upper) is not None: cache_key = (name_upper, str(device), compile_model) if use_cache and cache_key in cls._cache: return cls._cache[cache_key] try: - adapter = cls._adapters[name_upper](device, compile_model=compile_model) + adapter_cls = cls._engines.resolve(name_upper) + assert adapter_cls is not None # guarded by the branch condition + adapter = adapter_cls(device, compile_model=compile_model) except ImportError as exc: # Only translate the *absence of torchani itself*. A broken # torchani whose own transitive import fails names a different @@ -185,15 +202,25 @@ def create( @classmethod def available_models(cls) -> list[str]: - """Return list of registered model names.""" - return [ - MODEL_AIMNET, - "aimnet2-2025", - "aimnet2-nse", - "aimnet2-pd", - MODEL_ANI2X, - MODEL_ANI2XT, - ] + """Registered engine names, in declaration order. + + Was a hand-written literal that restated the registry's contents and was + derived from nothing -- so adding an engine and forgetting this list left + it working but undiscoverable. + """ + return cls._engines.available() + + +# The engines built from a local adapter class must be exactly +# BUILTIN_ANI_MODELS. This was an assert over a hand-maintained dict; it is now +# derived from the registry, so the two cannot drift apart. At module scope +# because a comprehension inside a class body cannot see that class's own +# attributes. +assert { + name.upper() + for name in ModelFactory._engines.available() + if ModelFactory._engines.resolve(name) is not None +} == set(BUILTIN_ANI_MODELS) def create_model( diff --git a/tests/test_cli_exit_codes.py b/tests/test_cli_exit_codes.py index 3022e44..b7f62bc 100644 --- a/tests/test_cli_exit_codes.py +++ b/tests/test_cli_exit_codes.py @@ -345,7 +345,11 @@ def _broken_adapter(device, compile_model=False): "No module named 'some_transitive_dep'", name="some_transitive_dep" ) - monkeypatch.setitem(mf.ModelFactory._adapters, "ANI2X", _broken_adapter) + monkeypatch.setitem( + mf.ModelFactory._engines._entries, + "ANI2x", + mf.ModelFactory._engines.entry("ANI2x").__class__(name="ANI2x", value=_broken_adapter), + ) monkeypatch.setattr(mf.ModelFactory, "_cache", {}) result = runner.invoke(app, ["models", "test", "ANI2x", "--no-gpu"]) diff --git a/tests/test_model_factory.py b/tests/test_model_factory.py index 119bae8..5ece158 100644 --- a/tests/test_model_factory.py +++ b/tests/test_model_factory.py @@ -35,7 +35,10 @@ def test_the_factory_promises_the_contract_not_the_base_class(): assert inspect.get_annotations(ModelFactory.create.__func__)["return"] == "ModelAdapter" # BaseModelAdapter survives in exactly one type position: a registry of # Auto3D's OWN adapter classes, which really is the concrete base. - assert inspect.get_annotations(ModelFactory)["_adapters"] == "dict[str, type[BaseModelAdapter]]" + assert ( + inspect.get_annotations(ModelFactory)["_engines"] + == "Registry[type[BaseModelAdapter] | None]" + ) # ...and it is no longer an incidental runtime re-export of this module. assert not hasattr(model_factory, "BaseModelAdapter") @@ -51,8 +54,17 @@ def test_registry_is_populated(self): assert "ANI2x" in models assert "ANI2xt" in models # AIMNET is NOT a hard-coded adapter key anymore: only ANI engines are. - assert set(ModelFactory._adapters) == {"ANI2X", "ANI2XT"} - assert "AIMNET" not in ModelFactory._adapters + # The registry now holds every user-facing engine name; the ones built + # from a local adapter class are those whose entry carries one. AIMNET + # is registered (it is a name a user may pass) but resolves to None, + # because it is looked up through the aimnet registry instead. + built_locally = { + n.upper() + for n in ModelFactory._engines.available() + if ModelFactory._engines.resolve(n) is not None + } + assert built_locally == {"ANI2X", "ANI2XT"} + assert ModelFactory._engines.resolve("AIMNET") is None def test_create_unknown_model_raises_error(self, monkeypatch): """Unknown non-path names no longer raise a ValueError up front; they @@ -93,11 +105,10 @@ def __init__(self, model_name, device, **kw): assert captured["aimnet"] == "aimnet2" # ANI engines resolve case-insensitively to their adapter class. - assert ( - model_factory.ModelFactory._adapters["ani2x".upper()] - is model_factory.ModelFactory._adapters["ANI2X"] - ) - assert "ANI2XT" in model_factory.ModelFactory._adapters + assert model_factory.ModelFactory._engines.resolve( + "ani2x" + ) is model_factory.ModelFactory._engines.resolve("ANI2X") + assert "ANI2XT" in model_factory.ModelFactory._engines @pytest.mark.slow def test_create_aimnet_returns_aimnet2_adapter(self, aimnet_model): @@ -364,7 +375,13 @@ class _FakeANI2xtAdapter: def __init__(self, device, **kw): captured["built_in"] = True - monkeypatch.setitem(model_factory.ModelFactory._adapters, "ANI2XT", _FakeANI2xtAdapter) + monkeypatch.setitem( + model_factory.ModelFactory._engines._entries, + "ANI2xt", + model_factory.ModelFactory._engines.entry("ANI2xt").__class__( + name="ANI2xt", value=_FakeANI2xtAdapter + ), + ) model_factory.ModelFactory.clear_cache() result = model_factory.create_model("ANI2xt", torch.device("cpu"), use_cache=False)