feat: add auto3d models test <engine> health check

Loads an engine and runs one tiny (methane) forward pass to verify it works in
the current environment, surfacing the common setup problems up front: a missing
torchani for the ANI engines, a failed/blocked aimnet registry download, or a
broken custom model file. Exits 0 on success, 3 if a dependency is missing, 5 if
the model loads but produces non-finite output. Also fixes the `models info`
unknown-engine hint to list the real engine names (ModelFactory.available_models)
instead of the upper-cased internal dict keys. Docs + tests included.

Author: isayev

CHANGELOG.md

@@ -14,6 +14,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   optimization, thermochemistry, and tautomer ranking from the CLI (previously
   Python-only). Each supports `--engine`, `--gpu/--no-gpu`, `--gpu-idx`,
   `-o/--output`, and `--json`.
+- **`auto3d models test <engine>`** - loads an engine and runs a tiny forward
+  pass as a health check (catches a missing torchani, a failed aimnet registry
+  download, or a broken custom model file before a full run).
 - **CLI ergonomics** - `auto3d run` gains `--job-name` and `--save-intermediate`;
   `config init` gains `--force`; choice flags use enums with shell completion;
   input paths are validated up front; and commands return differentiated exit

README.md

@@ -91,6 +91,7 @@ for mol in mols:
 | `auto3d config validate <file>` | Validate a configuration file |
 | `auto3d models list` | List available NNP models |
 | `auto3d models info <engine>` | Show model details |
+| `auto3d models test <engine>` | Load an engine and run a forward pass to verify it works |
 | `auto3d validate <input>` | Validate input file |
 
 ### Common Options

docs/source/cli.rst

@@ -356,6 +356,30 @@ Show detailed information about a specific engine.
    # Get ANI2x details
    auto3d models info ANI2x
 
+auto3d models test
+^^^^^^^^^^^^^^^^^^
+
+Load an engine and run a single tiny forward pass to confirm it works in the
+current environment -- catching a missing ``torchani``, a failed aimnet registry
+download, or a broken custom model file up front rather than mid-run.
+
+**Usage:**
+
+.. code:: console
+
+   auto3d models test ENGINE [--gpu/--no-gpu] [--gpu-idx N]
+
+**Examples:**
+
+.. code:: console
+
+   auto3d models test AIMNET            # verify the default engine on GPU
+   auto3d models test ANI2x --no-gpu    # verify ANI2x loads/runs on CPU
+   auto3d models test ./my_model.pt     # verify a custom NNP file
+
+Exits 0 on success; 3 if a dependency is missing, 5 if the model loads but
+produces non-finite output.
+
 Shell Completion
 ----------------
 

src/Auto3D/cli/app.py

@@ -235,6 +235,23 @@ def models_info(
     execute_models_info(engine=engine)
 
 
+@models_app.command("test")
+def models_test(
+    engine: Annotated[
+        str,
+        typer.Argument(
+            autocompletion=engine_autocomplete,
+            help="Engine to health-check: AIMNET, ANI2x, ANI2xt, a registry name, or a model path.",
+        ),
+    ],
+    gpu: Annotated[bool, typer.Option("--gpu/--no-gpu", help="Use GPU when available.")] = True,
+    gpu_idx: Annotated[int, typer.Option("--gpu-idx", help="CUDA device index.")] = 0,
+) -> None:
+    """Load an engine and run a tiny forward pass to verify it works."""
+    from Auto3D.cli.commands.models import execute_models_test
+    execute_models_test(engine=engine, gpu=gpu, gpu_idx=gpu_idx)
+
+
 @app.command()
 def validate(
     input_file: InputFile,

src/Auto3D/cli/commands/models.py

@@ -182,9 +182,10 @@ def execute_models_info(engine: str) -> None:
         engine_upper = "AIMNET"
 
     if engine_upper not in ENGINE_INFO:
+        from Auto3D.model_factory import ModelFactory
         print_error(
             f"Unknown engine: {engine}",
-            hint=f"Available: {', '.join(ENGINE_INFO.keys())}",
+            hint=f"Available: {', '.join(ModelFactory.available_models())}",
         )
         raise SystemExit(1)
 
@@ -207,3 +208,50 @@ def execute_models_info(engine: str) -> None:
         content += f"  - {note}\n"
 
     console.print(Panel(content, title=f"[cyan]{engine_upper}[/cyan]", border_style="blue"))
+
+
+def execute_models_test(engine: str, gpu: bool = True, gpu_idx: int = 0) -> None:
+    """Health-check an engine: load it and run one tiny forward pass.
+
+    Catches the common environment problems up front -- a missing torchani for
+    ANI, a failed/blocked aimnet registry download, or a broken custom model
+    file -- instead of having them surface deep inside a run.
+    """
+    from Auto3D.cli.console import print_success
+    from Auto3D.cli.errors import handle_error
+
+    try:
+        import time
+
+        import torch
+
+        from Auto3D.exceptions import NumericalError
+        from Auto3D.model_factory import create_model, get_device
+
+        device = get_device(gpu_idx, use_gpu=gpu)
+        with console.status(f"[bold]Loading {engine} on {device}..."):
+            t0 = time.time()
+            adapter = create_model(engine, device)
+            # A single methane molecule (H, C only -> supported by every engine).
+            coords = torch.tensor(
+                [[[0.0, 0.0, 0.0], [0.63, 0.63, 0.63], [-0.63, -0.63, 0.63],
+                  [0.63, -0.63, -0.63], [-0.63, 0.63, -0.63]]],
+                dtype=torch.float, device=device,
+            )
+            species = torch.tensor([[6, 1, 1, 1, 1]], device=device)
+            charges = torch.tensor([0.0], device=device)
+            energy, forces = adapter.forward(coords, species, charges)
+            elapsed = time.time() - t0
+
+        if not (torch.isfinite(energy).all() and torch.isfinite(forces).all()):
+            raise NumericalError(
+                f"{engine} produced non-finite energy/forces on the test molecule."
+            )
+
+        e_ev = float(energy.reshape(-1)[0])
+        print_success(
+            f"{engine} is working on {device} "
+            f"(methane E = {e_ev:.4f} eV, {elapsed:.1f}s incl. load)."
+        )
+    except Exception as e:  # noqa: BLE001 - present every failure as a clean panel
+        handle_error(e)

tests/test_cli_property_commands.py

@@ -164,6 +164,49 @@ def test_config_init_preset_enum_valid(tmp_path):
     assert target.exists()
 
 
+def test_models_test_success(monkeypatch):
+    """`models test` loads the engine and runs a forward; reports success."""
+    import torch
+
+    class _StubAdapter:
+        def forward(self, coords, species, charges):
+            return torch.zeros(1), torch.zeros(1, 5, 3)
+
+    monkeypatch.setattr("Auto3D.model_factory.get_device", lambda *a, **k: torch.device("cpu"))
+    monkeypatch.setattr("Auto3D.model_factory.create_model", lambda *a, **k: _StubAdapter())
+    res = runner.invoke(app, ["models", "test", "AIMNET", "--no-gpu"])
+    assert res.exit_code == 0, res.output
+    assert "working" in res.output
+
+
+def test_models_test_load_failure_exit_code(monkeypatch):
+    """A load failure (e.g. missing dependency) exits with the mapped code."""
+    from Auto3D.exceptions import DependencyError
+
+    def _boom(*a, **k):
+        raise DependencyError("torchani not installed")
+
+    monkeypatch.setattr("Auto3D.model_factory.get_device", lambda *a, **k: __import__("torch").device("cpu"))
+    monkeypatch.setattr("Auto3D.model_factory.create_model", _boom)
+    res = runner.invoke(app, ["models", "test", "ANI2x", "--no-gpu"])
+    assert res.exit_code == 3  # DependencyError -> 3
+    assert "Traceback" not in res.output
+
+
+def test_models_test_non_finite_exit_code(monkeypatch):
+    """Non-finite outputs are reported as a model (numerical) error -> exit 5."""
+    import torch
+
+    class _NanAdapter:
+        def forward(self, coords, species, charges):
+            return torch.tensor([float("nan")]), torch.zeros(1, 5, 3)
+
+    monkeypatch.setattr("Auto3D.model_factory.get_device", lambda *a, **k: torch.device("cpu"))
+    monkeypatch.setattr("Auto3D.model_factory.create_model", lambda *a, **k: _NanAdapter())
+    res = runner.invoke(app, ["models", "test", "AIMNET", "--no-gpu"])
+    assert res.exit_code == 5  # NumericalError (ModelError) -> 5
+
+
 def test_api_functions_expose_new_params():
     """calc_spe/opt_geometry/calc_thermo must accept out_path/use_gpu/allow_tf32
     so the CLI can drive output location, GPU choice, and TF32 uniformly."""