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)