From 665e50836e3f71410a66d34d5c24a32a34e39fb6 Mon Sep 17 00:00:00 2001 From: njzjz-bot <48687836+njzjz-bot@users.noreply.github.com> Date: Sun, 1 Mar 2026 20:18:00 +0000 Subject: [PATCH 01/11] feat(pretrained): integrate pretrained downloader and alias backend - add dp pretrained download CLI command - move pretrained logic under deepmd/pretrained - add built-in model registry with multi-source probing and fallback - register .pretrained backend alias so DeepPot usage stays unchanged - keep deep-eval adapter lazy to avoid circular imports - add parser/backend/downloader tests Authored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.3-codex) --- deepmd/__about__.py | 4 + deepmd/backend/pretrained.py | 60 +++++ deepmd/entrypoints/main.py | 5 + deepmd/main.py | 33 +++ deepmd/pretrained/__init__.py | 2 + deepmd/pretrained/backend.py | 143 ++++++++++++ deepmd/pretrained/download.py | 206 ++++++++++++++++++ deepmd/pretrained/entrypoints.py | 29 +++ deepmd/pretrained/registry.py | 30 +++ .../tests/common/test_pretrained_backend.py | 43 ++++ .../tests/common/test_pretrained_download.py | 105 +++++++++ source/tests/common/test_pretrained_parser.py | 42 ++++ 12 files changed, 702 insertions(+) create mode 100644 deepmd/__about__.py create mode 100644 deepmd/backend/pretrained.py create mode 100644 deepmd/pretrained/__init__.py create mode 100644 deepmd/pretrained/backend.py create mode 100644 deepmd/pretrained/download.py create mode 100644 deepmd/pretrained/entrypoints.py create mode 100644 deepmd/pretrained/registry.py create mode 100644 source/tests/common/test_pretrained_backend.py create mode 100644 source/tests/common/test_pretrained_download.py create mode 100644 source/tests/common/test_pretrained_parser.py diff --git a/deepmd/__about__.py b/deepmd/__about__.py new file mode 100644 index 0000000000..d406010330 --- /dev/null +++ b/deepmd/__about__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Package metadata used when setuptools-scm version is not available.""" + +__version__ = "unknown" diff --git a/deepmd/backend/pretrained.py b/deepmd/backend/pretrained.py new file mode 100644 index 0000000000..fa7f43fe37 --- /dev/null +++ b/deepmd/backend/pretrained.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Callable, +) +from typing import ( + TYPE_CHECKING, + ClassVar, +) + +from deepmd.backend.backend import ( + Backend, +) + +if TYPE_CHECKING: + from argparse import ( + Namespace, + ) + + from deepmd.infer.deep_eval import ( + DeepEvalBackend, + ) + from deepmd.utils.neighbor_stat import ( + NeighborStat, + ) + + +@Backend.register("pretrained") +class PretrainedBackend(Backend): + """Backend for ``*.pretrained`` model aliases.""" + + name = "Pretrained" + features: ClassVar[Backend.Feature] = Backend.Feature.DEEP_EVAL + suffixes: ClassVar[list[str]] = [".pretrained"] + + def is_available(self) -> bool: + return True + + @property + def entry_point_hook(self) -> Callable[["Namespace"], None]: + raise NotImplementedError("Unsupported backend: pretrained") + + @property + def deep_eval(self) -> type["DeepEvalBackend"]: + from deepmd.pretrained.backend import ( + get_pretrained_deep_eval_backend, + ) + + return get_pretrained_deep_eval_backend() + + @property + def neighbor_stat(self) -> type["NeighborStat"]: + raise NotImplementedError("Unsupported backend: pretrained") + + @property + def serialize_hook(self) -> Callable[[str], dict]: + raise NotImplementedError("Unsupported backend: pretrained") + + @property + def deserialize_hook(self) -> Callable[[str, dict], None]: + raise NotImplementedError("Unsupported backend: pretrained") diff --git a/deepmd/entrypoints/main.py b/deepmd/entrypoints/main.py index 34ebe4d2e3..86c9687bd4 100644 --- a/deepmd/entrypoints/main.py +++ b/deepmd/entrypoints/main.py @@ -39,6 +39,9 @@ from deepmd.loggers.loggers import ( set_log_handles, ) +from deepmd.pretrained.entrypoints import ( + pretrained_entrypoint, +) def main(args: argparse.Namespace) -> None: @@ -97,5 +100,7 @@ def main(args: argparse.Namespace) -> None: convert_backend(**dict_args) elif args.command == "show": show(**dict_args) + elif args.command == "pretrained": + pretrained_entrypoint(args) else: raise ValueError(f"Unknown command: {args.command}") diff --git a/deepmd/main.py b/deepmd/main.py index 62118ae3c6..d239ec80c2 100644 --- a/deepmd/main.py +++ b/deepmd/main.py @@ -942,6 +942,38 @@ def main_parser() -> argparse.ArgumentParser: ], nargs="+", ) + + # pretrained + parser_pretrained = subparsers.add_parser( + "pretrained", + parents=[parser_log], + help="Manage builtin pretrained models", + formatter_class=RawTextArgumentDefaultsHelpFormatter, + ) + pretrained_subparsers = parser_pretrained.add_subparsers( + dest="pretrained_command", + required=True, + ) + parser_pretrained_download = pretrained_subparsers.add_parser( + "download", + help="Download one pretrained model", + ) + from deepmd.pretrained.registry import ( + available_model_names, + ) + + parser_pretrained_download.add_argument( + "MODEL", + choices=available_model_names(), + help="Pretrained model name", + ) + parser_pretrained_download.add_argument( + "--cache-dir", + default=None, + type=str, + help="Optional cache directory for pretrained model files", + ) + return parser @@ -997,6 +1029,7 @@ def main(args: list[str] | None = None) -> None: "gui", "convert-backend", "show", + "pretrained", ): # common entrypoints from deepmd.entrypoints.main import main as deepmd_main diff --git a/deepmd/pretrained/__init__.py b/deepmd/pretrained/__init__.py new file mode 100644 index 0000000000..0da3f0dbb0 --- /dev/null +++ b/deepmd/pretrained/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Pretrained model helpers for DeePMD-kit.""" diff --git a/deepmd/pretrained/backend.py b/deepmd/pretrained/backend.py new file mode 100644 index 0000000000..775a03f768 --- /dev/null +++ b/deepmd/pretrained/backend.py @@ -0,0 +1,143 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Backend helper for `*.pretrained` model aliases.""" + +from __future__ import annotations + +from functools import ( + lru_cache, +) +from pathlib import ( + Path, +) +from typing import ( + TYPE_CHECKING, + Any, +) + +from deepmd.pretrained.download import ( + resolve_model_path, +) + +if TYPE_CHECKING: + import numpy as np + + from deepmd.infer.deep_eval import ( + DeepEval, + DeepEvalBackend, + ) + + +def parse_pretrained_alias(model_file: str) -> str: + """Extract model name from ``*.pretrained`` alias string.""" + alias = Path(model_file).name + suffix = ".pretrained" + if not alias.endswith(suffix): + raise ValueError(f"Invalid pretrained alias: {model_file}") + + model_name = alias[: -len(suffix)] + if not model_name: + raise ValueError(f"Invalid pretrained alias: {model_file}") + + return model_name + + +@lru_cache(maxsize=1) +def get_pretrained_deep_eval_backend() -> type[DeepEvalBackend]: + """Build and cache the concrete DeepEval adapter lazily.""" + # Avoid circular import when deepmd backend entrypoints are loading. + from deepmd.infer.deep_eval import ( + DeepEvalBackend, + ) + + class PretrainedDeepEvalBackend(DeepEvalBackend): + """Resolve alias and delegate to backend selected by resolved model path.""" + + def __init__( + self, + model_file: str, + output_def: object, + *args: object, + auto_batch_size: object = True, + neighbor_list: object | None = None, + **kwargs: object, + ) -> None: + model_name = parse_pretrained_alias(model_file) + resolved = str(resolve_model_path(model_name)) + + # DeepEvalBackend.__new__ dispatches by resolved suffix (.pt/.pb/.dp...) + self._backend = DeepEvalBackend( + resolved, + output_def, + *args, + auto_batch_size=auto_batch_size, + neighbor_list=neighbor_list, + **kwargs, + ) + + def eval( + self, + coords: np.ndarray, + cells: np.ndarray | None, + atom_types: np.ndarray, + atomic: bool = False, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, + **kwargs: Any, + ) -> dict[str, np.ndarray]: + return self._backend.eval( + coords, + cells, + atom_types, + atomic, + fparam=fparam, + aparam=aparam, + **kwargs, + ) + + def get_rcut(self) -> float: + return self._backend.get_rcut() + + def get_ntypes(self) -> int: + return self._backend.get_ntypes() + + def get_type_map(self) -> list[str]: + return self._backend.get_type_map() + + def get_dim_fparam(self) -> int: + return self._backend.get_dim_fparam() + + def has_default_fparam(self) -> bool: + return self._backend.has_default_fparam() + + def get_dim_aparam(self) -> int: + return self._backend.get_dim_aparam() + + @property + def model_type(self) -> type[DeepEval]: + return self._backend.model_type + + def get_sel_type(self) -> list[int]: + return self._backend.get_sel_type() + + def get_numb_dos(self) -> int: + return self._backend.get_numb_dos() + + def get_has_efield(self) -> bool: + return self._backend.get_has_efield() + + def get_has_spin(self) -> bool: + return self._backend.get_has_spin() + + def get_has_hessian(self) -> bool: + return self._backend.get_has_hessian() + + def get_var_name(self) -> str: + return self._backend.get_var_name() + + def get_ntypes_spin(self) -> int: + return self._backend.get_ntypes_spin() + + def get_model(self) -> Any: + return self._backend.get_model() + + return PretrainedDeepEvalBackend diff --git a/deepmd/pretrained/download.py b/deepmd/pretrained/download.py new file mode 100644 index 0000000000..ee1775fc3d --- /dev/null +++ b/deepmd/pretrained/download.py @@ -0,0 +1,206 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Download and resolve pretrained model files.""" + +from __future__ import annotations + +import concurrent.futures +import hashlib +import logging +import shutil +import time +import urllib.error +import urllib.parse +import urllib.request +from pathlib import ( + Path, +) +from typing import ( + Any, +) + +from deepmd.pretrained.registry import ( + MODEL_REGISTRY, +) + +DEFAULT_CACHE_DIR = Path.home() / ".cache" / "deepmd" / "pretrained" / "models" +DOWNLOAD_TIMEOUT_SECONDS = 120 +SOURCE_PROBE_TIMEOUT_SECONDS = 8 + + +def _validate_download_url(url: str) -> None: + """Validate that download URL uses HTTPS scheme.""" + parsed = urllib.parse.urlparse(url) + if parsed.scheme != "https": + raise ValueError(f"Unsupported URL scheme for download: {parsed.scheme}") + + +def _sha256sum(path: Path) -> str: + """Calculate SHA256 checksum of a file.""" + hasher = hashlib.sha256() + with path.open("rb") as f: + for chunk in iter(lambda: f.read(1024 * 1024), b""): + hasher.update(chunk) + return hasher.hexdigest() + + +def _model_download_urls(model_info: dict[str, Any]) -> list[str]: + """Return candidate download URLs (deduplicated and ordered).""" + candidates: list[str] = [] + raw_urls = model_info.get("urls") + if isinstance(raw_urls, list): + candidates.extend(item for item in raw_urls if isinstance(item, str)) + + if not candidates and isinstance(model_info.get("url"), str): + # backward compatibility + candidates.append(model_info["url"]) + + seen: set[str] = set() + unique: list[str] = [] + for url in candidates: + if url not in seen: + seen.add(url) + unique.append(url) + return unique + + +def _probe_download_url(url: str) -> float | None: + """Probe one URL and return latency seconds if reachable; else None.""" + _validate_download_url(url) + request = urllib.request.Request( + url, + headers={"Range": "bytes=0-0"}, + method="GET", + ) + start = time.monotonic() + try: + with urllib.request.urlopen(request, timeout=SOURCE_PROBE_TIMEOUT_SECONDS): + pass + except (urllib.error.URLError, OSError, ValueError): + return None + + return time.monotonic() - start + + +def _rank_download_urls(urls: list[str]) -> list[str]: + """Rank candidate URLs by probe latency (fastest first).""" + if len(urls) <= 1: + return urls + + results: dict[str, float] = {} + with concurrent.futures.ThreadPoolExecutor(max_workers=min(4, len(urls))) as exe: + future_to_url = {exe.submit(_probe_download_url, url): url for url in urls} + for future in concurrent.futures.as_completed(future_to_url): + url = future_to_url[future] + latency = future.result() + if latency is not None: + results[url] = latency + + ranked_ok = sorted(results, key=lambda url: results[url]) + ranked_fail = [url for url in urls if url not in results] + return ranked_ok + ranked_fail + + +def _download_file(url: str, destination: Path) -> None: + """Download URL content to destination atomically.""" + _validate_download_url(url) + destination.parent.mkdir(parents=True, exist_ok=True) + tmp_path = destination.with_suffix(destination.suffix + ".part") + + try: + with ( + urllib.request.urlopen(url, timeout=DOWNLOAD_TIMEOUT_SECONDS) as response, + tmp_path.open("wb") as out_file, + ): + shutil.copyfileobj(response, out_file) + except Exception: + tmp_path.unlink(missing_ok=True) + raise + + tmp_path.replace(destination) + + +def download_model( + model_name: str, + *, + cache_dir: Path | None = None, + logger: logging.Logger | None = None, +) -> Path: + """Download one model and return local path. + + The function will probe all configured sources, try the fastest reachable + source first, and then fallback to others when failure happens. + """ + log = logger or logging.getLogger(__name__) + + model_info = MODEL_REGISTRY.get(model_name) + if model_info is None: + available = ", ".join(sorted(MODEL_REGISTRY)) + raise ValueError(f"Unknown model: {model_name}. Available: {available}") + + target_dir = cache_dir or DEFAULT_CACHE_DIR + output_path = target_dir / str(model_info["filename"]) + expected_sha256 = str(model_info["sha256"]) + + if output_path.exists(): + actual = _sha256sum(output_path) + if actual == expected_sha256: + log.info("Model '%s' already exists at: %s", model_name, output_path) + return output_path + log.warning( + "Cached file for '%s' failed SHA256 check, re-downloading...", + model_name, + ) + output_path.unlink(missing_ok=True) + + urls = _model_download_urls(model_info) + if not urls: + raise RuntimeError(f"No download URL configured for model '{model_name}'") + + ranked_urls = _rank_download_urls(urls) + if len(ranked_urls) > 1: + log.info( + "Selecting fastest source among %d candidates...", + len(ranked_urls), + ) + + for idx, url in enumerate(ranked_urls, start=1): + log.info("Downloading '%s' (source %d/%d): %s", model_name, idx, len(ranked_urls), url) + try: + _download_file(url, output_path) + except (urllib.error.URLError, OSError, ValueError) as exc: + log.warning("Download attempt failed from %s: %s", url, exc) + continue + + actual = _sha256sum(output_path) + if actual != expected_sha256: + output_path.unlink(missing_ok=True) + log.warning("SHA256 verification failed from source: %s", url) + log.warning("Expected: %s", expected_sha256) + log.warning("Actual: %s", actual) + continue + + log.info("Downloaded '%s' to: %s", model_name, output_path) + return output_path + + raise RuntimeError(f"Failed to download model '{model_name}' from all sources") + + +def resolve_model_path( + model_name: str, + *, + cache_dir: Path | None = None, + logger: logging.Logger | None = None, +) -> Path: + """Resolve model alias to verified local file, downloading if needed.""" + target_dir = cache_dir or DEFAULT_CACHE_DIR + model_info = MODEL_REGISTRY.get(model_name) + if model_info is None: + available = ", ".join(sorted(MODEL_REGISTRY)) + raise ValueError(f"Unknown model: {model_name}. Available: {available}") + + output_path = target_dir / str(model_info["filename"]) + expected_sha256 = str(model_info["sha256"]) + if output_path.exists() and _sha256sum(output_path) == expected_sha256: + return output_path + + return download_model(model_name, cache_dir=target_dir, logger=logger) diff --git a/deepmd/pretrained/entrypoints.py b/deepmd/pretrained/entrypoints.py new file mode 100644 index 0000000000..f463273f16 --- /dev/null +++ b/deepmd/pretrained/entrypoints.py @@ -0,0 +1,29 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""CLI entrypoint for pretrained model operations.""" + +from __future__ import annotations + +from pathlib import ( + Path, +) +from typing import ( + TYPE_CHECKING, +) + +from deepmd.pretrained.download import ( + download_model, +) + +if TYPE_CHECKING: + import argparse + + +def pretrained_entrypoint(args: argparse.Namespace) -> None: + """Handle `dp pretrained ...` subcommands.""" + if args.pretrained_command == "download": + cache_dir = Path(args.cache_dir) if args.cache_dir else None + path = download_model(args.MODEL, cache_dir=cache_dir) + print(path) # noqa: T201 + return + + raise ValueError(f"Unknown pretrained subcommand: {args.pretrained_command}") diff --git a/deepmd/pretrained/registry.py b/deepmd/pretrained/registry.py new file mode 100644 index 0000000000..b540cd3490 --- /dev/null +++ b/deepmd/pretrained/registry.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Registry of built-in pretrained model sources.""" + +from typing import ( + Any, +) + +MODEL_REGISTRY: dict[str, dict[str, Any]] = { + "DPA-3.2-5M": { + "urls": [ + "https://huggingface.co/deepmodelingcommunity/DPA-3.2-5M/resolve/main/DPA-3.2-5M.pt?download=true", + "https://hf-mirror.com/deepmodelingcommunity/DPA-3.2-5M/resolve/main/DPA-3.2-5M.pt?download=true", + ], + "filename": "DPA-3.2-5M.pt", + "sha256": "876354744aeaae17b2639a6a690514470273784f2b4836280850f50cbb799165", + }, + "DPA-3.1-3M": { + "urls": [ + "https://huggingface.co/deepmodelingcommunity/DPA-3.1-3M/resolve/main/DPA-3.1-3M.pt?download=true", + "https://hf-mirror.com/deepmodelingcommunity/DPA-3.1-3M/resolve/main/DPA-3.1-3M.pt?download=true", + ], + "filename": "DPA-3.1-3M.pt", + "sha256": "86dd3a804d78ca5d203ebf98747e8f16dff9713ba8950097ceb760b161e19907", + }, +} + + +def available_model_names() -> list[str]: + """Return available model names from built-in registry.""" + return sorted(MODEL_REGISTRY) diff --git a/source/tests/common/test_pretrained_backend.py b/source/tests/common/test_pretrained_backend.py new file mode 100644 index 0000000000..65a9472181 --- /dev/null +++ b/source/tests/common/test_pretrained_backend.py @@ -0,0 +1,43 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for pretrained backend registration and alias parsing.""" + +import unittest +from unittest.mock import ( + patch, +) + +import deepmd.backend # noqa: F401 +from deepmd.backend.backend import ( + Backend, +) +from deepmd.backend.pretrained import ( + PretrainedBackend, +) +from deepmd.pretrained.backend import ( + parse_pretrained_alias, +) + + +class TestPretrainedBackend(unittest.TestCase): + """Test pretrained backend integration points.""" + + def test_detect_backend_by_pretrained_suffix(self) -> None: + backend = Backend.detect_backend_by_model("DPA-3.2-5M.pretrained") + self.assertIs(backend, PretrainedBackend) + + def test_parse_pretrained_alias(self) -> None: + self.assertEqual( + parse_pretrained_alias("DPA-3.2-5M.pretrained"), + "DPA-3.2-5M", + ) + + def test_parse_pretrained_alias_invalid(self) -> None: + with self.assertRaises(ValueError): + parse_pretrained_alias("DPA-3.2-5M.pt") + + def test_deep_eval_property_is_lazy(self) -> None: + with patch( + "deepmd.pretrained.backend.get_pretrained_deep_eval_backend", + return_value=object, + ): + self.assertIs(PretrainedBackend().deep_eval, object) diff --git a/source/tests/common/test_pretrained_download.py b/source/tests/common/test_pretrained_download.py new file mode 100644 index 0000000000..6fa2b7af29 --- /dev/null +++ b/source/tests/common/test_pretrained_download.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for pretrained download/resolve helpers.""" + +from __future__ import annotations + +import hashlib +import tempfile +import unittest +import urllib.error +from pathlib import ( + Path, +) +from unittest.mock import ( + patch, +) + +from deepmd.pretrained import download as dl + + +class TestPretrainedDownload(unittest.TestCase): + """Test download helper behavior.""" + + def test_model_download_urls_prefers_urls(self) -> None: + info = { + "urls": ["https://a", "https://a", "https://b"], + "url": "https://legacy", + } + self.assertEqual(dl._model_download_urls(info), ["https://a", "https://b"]) + + def test_rank_download_urls(self) -> None: + with patch.object( + dl, + "_probe_download_url", + side_effect=lambda url: { + "https://a": 0.3, + "https://b": 0.1, + "https://c": None, + }[url], + ): + ranked = dl._rank_download_urls(["https://a", "https://b", "https://c"]) + + self.assertEqual(ranked, ["https://b", "https://a", "https://c"]) + + def test_download_model_fallback_on_failure(self) -> None: + payload = b"payload" + expected = hashlib.sha256(payload).hexdigest() + model_name = "DPA-3.2-5M" + + with tempfile.TemporaryDirectory() as td: + cache_dir = Path(td) + + with patch.object( + dl, + "MODEL_REGISTRY", + { + model_name: { + "filename": "model.pt", + "sha256": expected, + "urls": ["https://a", "https://b"], + } + }, + ): + with patch.object( + dl, + "_rank_download_urls", + return_value=["https://a", "https://b"], + ): + + def fake_download(url: str, destination: Path) -> None: + if url == "https://a": + raise urllib.error.URLError("timeout") + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_bytes(payload) + + with patch.object(dl, "_download_file", side_effect=fake_download): + path = dl.download_model(model_name, cache_dir=cache_dir) + + self.assertTrue(path.exists()) + self.assertEqual(path.read_bytes(), payload) + + def test_resolve_model_path_cached(self) -> None: + payload = b"payload" + expected = hashlib.sha256(payload).hexdigest() + model_name = "DPA-3.2-5M" + + with tempfile.TemporaryDirectory() as td: + cache_dir = Path(td) + target = cache_dir / "model.pt" + target.parent.mkdir(parents=True, exist_ok=True) + target.write_bytes(payload) + + with patch.object( + dl, + "MODEL_REGISTRY", + { + model_name: { + "filename": "model.pt", + "sha256": expected, + "urls": ["https://a"], + } + }, + ): + path = dl.resolve_model_path(model_name, cache_dir=cache_dir) + + self.assertEqual(path, target) diff --git a/source/tests/common/test_pretrained_parser.py b/source/tests/common/test_pretrained_parser.py new file mode 100644 index 0000000000..a31af78184 --- /dev/null +++ b/source/tests/common/test_pretrained_parser.py @@ -0,0 +1,42 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for pretrained argument parsing.""" + +import unittest + +from deepmd.main import ( + parse_args, +) +from deepmd.pretrained.registry import ( + available_model_names, +) + + +class TestPretrainedParser(unittest.TestCase): + """Test `dp pretrained` parser behavior.""" + + def test_pretrained_download_parser(self) -> None: + model = available_model_names()[0] + args = parse_args(["pretrained", "download", model]) + + self.assertEqual(args.command, "pretrained") + self.assertEqual(args.pretrained_command, "download") + self.assertEqual(args.MODEL, model) + self.assertIsNone(args.cache_dir) + + def test_pretrained_download_with_cache_dir(self) -> None: + model = available_model_names()[0] + args = parse_args( + [ + "pretrained", + "download", + model, + "--cache-dir", + "/tmp/deepmd-pretrained", + ] + ) + + self.assertEqual(args.cache_dir, "/tmp/deepmd-pretrained") + + def test_pretrained_download_rejects_unknown_model(self) -> None: + with self.assertRaises(SystemExit): + parse_args(["pretrained", "download", "NOT-EXIST"]) From 06c1ff9ad3de324b03057fe494d53eeb80b97602 Mon Sep 17 00:00:00 2001 From: njzjz-bot <48687836+njzjz-bot@users.noreply.github.com> Date: Sun, 1 Mar 2026 20:39:15 +0000 Subject: [PATCH 02/11] chore: drop temporary deepmd/__about__.py fallback The fallback file was only added for local source-tree unittest convenience.\nKeep version behavior aligned with upstream packaging flow (_version.py via build).\n\nAuthored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.3-codex) --- deepmd/__about__.py | 4 ---- 1 file changed, 4 deletions(-) delete mode 100644 deepmd/__about__.py diff --git a/deepmd/__about__.py b/deepmd/__about__.py deleted file mode 100644 index d406010330..0000000000 --- a/deepmd/__about__.py +++ /dev/null @@ -1,4 +0,0 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -"""Package metadata used when setuptools-scm version is not available.""" - -__version__ = "unknown" From 35f384bb30f6283a6a428d5ee3ac55249a36d83a Mon Sep 17 00:00:00 2001 From: njzjz-bot <48687836+njzjz-bot@users.noreply.github.com> Date: Sun, 1 Mar 2026 20:56:33 +0000 Subject: [PATCH 03/11] style: apply prek formatting fixes for pretrained integration Authored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.3-codex) --- deepmd/pretrained/backend.py | 4 +++- deepmd/pretrained/download.py | 12 ++++++++++-- deepmd/pretrained/entrypoints.py | 4 +++- source/tests/common/test_pretrained_download.py | 4 +++- 4 files changed, 19 insertions(+), 5 deletions(-) diff --git a/deepmd/pretrained/backend.py b/deepmd/pretrained/backend.py index 775a03f768..acf4a12b7e 100644 --- a/deepmd/pretrained/backend.py +++ b/deepmd/pretrained/backend.py @@ -1,7 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Backend helper for `*.pretrained` model aliases.""" -from __future__ import annotations +from __future__ import ( + annotations, +) from functools import ( lru_cache, diff --git a/deepmd/pretrained/download.py b/deepmd/pretrained/download.py index ee1775fc3d..b24508ebdc 100644 --- a/deepmd/pretrained/download.py +++ b/deepmd/pretrained/download.py @@ -1,7 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Download and resolve pretrained model files.""" -from __future__ import annotations +from __future__ import ( + annotations, +) import concurrent.futures import hashlib @@ -164,7 +166,13 @@ def download_model( ) for idx, url in enumerate(ranked_urls, start=1): - log.info("Downloading '%s' (source %d/%d): %s", model_name, idx, len(ranked_urls), url) + log.info( + "Downloading '%s' (source %d/%d): %s", + model_name, + idx, + len(ranked_urls), + url, + ) try: _download_file(url, output_path) except (urllib.error.URLError, OSError, ValueError) as exc: diff --git a/deepmd/pretrained/entrypoints.py b/deepmd/pretrained/entrypoints.py index f463273f16..ba8a23de85 100644 --- a/deepmd/pretrained/entrypoints.py +++ b/deepmd/pretrained/entrypoints.py @@ -1,7 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """CLI entrypoint for pretrained model operations.""" -from __future__ import annotations +from __future__ import ( + annotations, +) from pathlib import ( Path, diff --git a/source/tests/common/test_pretrained_download.py b/source/tests/common/test_pretrained_download.py index 6fa2b7af29..c6afe6aee5 100644 --- a/source/tests/common/test_pretrained_download.py +++ b/source/tests/common/test_pretrained_download.py @@ -1,7 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Tests for pretrained download/resolve helpers.""" -from __future__ import annotations +from __future__ import ( + annotations, +) import hashlib import tempfile From 27ac825be353df2da78b67820838ba52c36b80e1 Mon Sep 17 00:00:00 2001 From: njzjz-bot <48687836+njzjz-bot@users.noreply.github.com> Date: Sun, 1 Mar 2026 21:34:00 +0000 Subject: [PATCH 04/11] refactor(pretrained): rename backend helper and address PR review comments - rename deepmd/pretrained/backend.py -> deepmd/pretrained/deep_eval.py\n- remove unnecessary lazy import in deepmd/backend/pretrained.py\n- simplify resolve_model_path by delegating directly to download_model\n- update tests per review comments (cached resolve should not trigger _download_file)\n- keep eval_descriptor/eval_fitting_last_layer delegation in adapter\n\nAuthored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.3-codex) --- deepmd/backend/pretrained.py | 7 ++-- .../pretrained/{backend.py => deep_eval.py} | 40 +++++++++++++++++++ deepmd/pretrained/download.py | 5 --- .../tests/common/test_pretrained_backend.py | 20 +++++----- .../tests/common/test_pretrained_download.py | 4 +- 5 files changed, 55 insertions(+), 21 deletions(-) rename deepmd/pretrained/{backend.py => deep_eval.py} (76%) diff --git a/deepmd/backend/pretrained.py b/deepmd/backend/pretrained.py index fa7f43fe37..2364acf6a3 100644 --- a/deepmd/backend/pretrained.py +++ b/deepmd/backend/pretrained.py @@ -10,6 +10,9 @@ from deepmd.backend.backend import ( Backend, ) +from deepmd.pretrained.deep_eval import ( + get_pretrained_deep_eval_backend, +) if TYPE_CHECKING: from argparse import ( @@ -41,10 +44,6 @@ def entry_point_hook(self) -> Callable[["Namespace"], None]: @property def deep_eval(self) -> type["DeepEvalBackend"]: - from deepmd.pretrained.backend import ( - get_pretrained_deep_eval_backend, - ) - return get_pretrained_deep_eval_backend() @property diff --git a/deepmd/pretrained/backend.py b/deepmd/pretrained/deep_eval.py similarity index 76% rename from deepmd/pretrained/backend.py rename to deepmd/pretrained/deep_eval.py index acf4a12b7e..6eb81b625a 100644 --- a/deepmd/pretrained/backend.py +++ b/deepmd/pretrained/deep_eval.py @@ -96,6 +96,46 @@ def eval( **kwargs, ) + def eval_descriptor( + self, + coords: np.ndarray, + cells: np.ndarray | None, + atom_types: np.ndarray, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, + efield: np.ndarray | None = None, + mixed_type: bool = False, + **kwargs: Any, + ) -> np.ndarray: + return self._backend.eval_descriptor( + coords, + cells, + atom_types, + fparam=fparam, + aparam=aparam, + efield=efield, + mixed_type=mixed_type, + **kwargs, + ) + + def eval_fitting_last_layer( + self, + coords: np.ndarray, + cells: np.ndarray | None, + atom_types: np.ndarray, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, + **kwargs: Any, + ) -> np.ndarray: + return self._backend.eval_fitting_last_layer( + coords, + cells, + atom_types, + fparam=fparam, + aparam=aparam, + **kwargs, + ) + def get_rcut(self) -> float: return self._backend.get_rcut() diff --git a/deepmd/pretrained/download.py b/deepmd/pretrained/download.py index b24508ebdc..6a3fb7d8cc 100644 --- a/deepmd/pretrained/download.py +++ b/deepmd/pretrained/download.py @@ -206,9 +206,4 @@ def resolve_model_path( available = ", ".join(sorted(MODEL_REGISTRY)) raise ValueError(f"Unknown model: {model_name}. Available: {available}") - output_path = target_dir / str(model_info["filename"]) - expected_sha256 = str(model_info["sha256"]) - if output_path.exists() and _sha256sum(output_path) == expected_sha256: - return output_path - return download_model(model_name, cache_dir=target_dir, logger=logger) diff --git a/source/tests/common/test_pretrained_backend.py b/source/tests/common/test_pretrained_backend.py index 65a9472181..961296cf12 100644 --- a/source/tests/common/test_pretrained_backend.py +++ b/source/tests/common/test_pretrained_backend.py @@ -1,19 +1,16 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Tests for pretrained backend registration and alias parsing.""" +import importlib import unittest -from unittest.mock import ( - patch, -) -import deepmd.backend # noqa: F401 from deepmd.backend.backend import ( Backend, ) from deepmd.backend.pretrained import ( PretrainedBackend, ) -from deepmd.pretrained.backend import ( +from deepmd.pretrained.deep_eval import ( parse_pretrained_alias, ) @@ -21,6 +18,11 @@ class TestPretrainedBackend(unittest.TestCase): """Test pretrained backend integration points.""" + @classmethod + def setUpClass(cls) -> None: + # ensure backend registration side effects are loaded + importlib.import_module("deepmd.backend") + def test_detect_backend_by_pretrained_suffix(self) -> None: backend = Backend.detect_backend_by_model("DPA-3.2-5M.pretrained") self.assertIs(backend, PretrainedBackend) @@ -35,9 +37,5 @@ def test_parse_pretrained_alias_invalid(self) -> None: with self.assertRaises(ValueError): parse_pretrained_alias("DPA-3.2-5M.pt") - def test_deep_eval_property_is_lazy(self) -> None: - with patch( - "deepmd.pretrained.backend.get_pretrained_deep_eval_backend", - return_value=object, - ): - self.assertIs(PretrainedBackend().deep_eval, object) + def test_deep_eval_property(self) -> None: + self.assertIsNotNone(PretrainedBackend().deep_eval) diff --git a/source/tests/common/test_pretrained_download.py b/source/tests/common/test_pretrained_download.py index c6afe6aee5..b943c1f247 100644 --- a/source/tests/common/test_pretrained_download.py +++ b/source/tests/common/test_pretrained_download.py @@ -102,6 +102,8 @@ def test_resolve_model_path_cached(self) -> None: } }, ): - path = dl.resolve_model_path(model_name, cache_dir=cache_dir) + with patch.object(dl, "_download_file") as mocked_download: + path = dl.resolve_model_path(model_name, cache_dir=cache_dir) self.assertEqual(path, target) + mocked_download.assert_not_called() From 73034122ce787c5714129829fa440f2fc0046614 Mon Sep 17 00:00:00 2001 From: njzjz-bot <48687836+njzjz-bot@users.noreply.github.com> Date: Sun, 1 Mar 2026 21:58:32 +0000 Subject: [PATCH 05/11] refactor(pretrained): lazy load at backend boundary, eager deep_eval module - keep lazy import in deepmd/backend/pretrained.py\n- keep deepmd/pretrained/deep_eval.py as regular (non-lazy) module\n- preserve deep eval delegations for descriptor/fitting-last-layer\n- simplify resolve_model_path and adjust tests for cached path behavior\n\nAuthored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.3-codex) --- deepmd/backend/pretrained.py | 9 +- deepmd/pretrained/deep_eval.py | 285 +++++++++--------- .../tests/common/test_pretrained_backend.py | 6 +- 3 files changed, 145 insertions(+), 155 deletions(-) diff --git a/deepmd/backend/pretrained.py b/deepmd/backend/pretrained.py index 2364acf6a3..f6f95d8c59 100644 --- a/deepmd/backend/pretrained.py +++ b/deepmd/backend/pretrained.py @@ -10,9 +10,6 @@ from deepmd.backend.backend import ( Backend, ) -from deepmd.pretrained.deep_eval import ( - get_pretrained_deep_eval_backend, -) if TYPE_CHECKING: from argparse import ( @@ -44,7 +41,11 @@ def entry_point_hook(self) -> Callable[["Namespace"], None]: @property def deep_eval(self) -> type["DeepEvalBackend"]: - return get_pretrained_deep_eval_backend() + from deepmd.pretrained.deep_eval import ( + PretrainedDeepEvalBackend, + ) + + return PretrainedDeepEvalBackend @property def neighbor_stat(self) -> type["NeighborStat"]: diff --git a/deepmd/pretrained/deep_eval.py b/deepmd/pretrained/deep_eval.py index 6eb81b625a..c722bcab5b 100644 --- a/deepmd/pretrained/deep_eval.py +++ b/deepmd/pretrained/deep_eval.py @@ -1,13 +1,10 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""Backend helper for `*.pretrained` model aliases.""" +"""DeepEval adapter for `*.pretrained` model aliases.""" from __future__ import ( annotations, ) -from functools import ( - lru_cache, -) from pathlib import ( Path, ) @@ -16,18 +13,16 @@ Any, ) +from deepmd.infer.deep_eval import ( + DeepEval, + DeepEvalBackend, +) from deepmd.pretrained.download import ( resolve_model_path, ) - if TYPE_CHECKING: import numpy as np - from deepmd.infer.deep_eval import ( - DeepEval, - DeepEvalBackend, - ) - def parse_pretrained_alias(model_file: str) -> str: """Extract model name from ``*.pretrained`` alias string.""" @@ -43,143 +38,133 @@ def parse_pretrained_alias(model_file: str) -> str: return model_name -@lru_cache(maxsize=1) -def get_pretrained_deep_eval_backend() -> type[DeepEvalBackend]: - """Build and cache the concrete DeepEval adapter lazily.""" - # Avoid circular import when deepmd backend entrypoints are loading. - from deepmd.infer.deep_eval import ( - DeepEvalBackend, - ) - - class PretrainedDeepEvalBackend(DeepEvalBackend): - """Resolve alias and delegate to backend selected by resolved model path.""" - - def __init__( - self, - model_file: str, - output_def: object, - *args: object, - auto_batch_size: object = True, - neighbor_list: object | None = None, - **kwargs: object, - ) -> None: - model_name = parse_pretrained_alias(model_file) - resolved = str(resolve_model_path(model_name)) - - # DeepEvalBackend.__new__ dispatches by resolved suffix (.pt/.pb/.dp...) - self._backend = DeepEvalBackend( - resolved, - output_def, - *args, - auto_batch_size=auto_batch_size, - neighbor_list=neighbor_list, - **kwargs, - ) - - def eval( - self, - coords: np.ndarray, - cells: np.ndarray | None, - atom_types: np.ndarray, - atomic: bool = False, - fparam: np.ndarray | None = None, - aparam: np.ndarray | None = None, - **kwargs: Any, - ) -> dict[str, np.ndarray]: - return self._backend.eval( - coords, - cells, - atom_types, - atomic, - fparam=fparam, - aparam=aparam, - **kwargs, - ) - - def eval_descriptor( - self, - coords: np.ndarray, - cells: np.ndarray | None, - atom_types: np.ndarray, - fparam: np.ndarray | None = None, - aparam: np.ndarray | None = None, - efield: np.ndarray | None = None, - mixed_type: bool = False, - **kwargs: Any, - ) -> np.ndarray: - return self._backend.eval_descriptor( - coords, - cells, - atom_types, - fparam=fparam, - aparam=aparam, - efield=efield, - mixed_type=mixed_type, - **kwargs, - ) - - def eval_fitting_last_layer( - self, - coords: np.ndarray, - cells: np.ndarray | None, - atom_types: np.ndarray, - fparam: np.ndarray | None = None, - aparam: np.ndarray | None = None, - **kwargs: Any, - ) -> np.ndarray: - return self._backend.eval_fitting_last_layer( - coords, - cells, - atom_types, - fparam=fparam, - aparam=aparam, - **kwargs, - ) - - def get_rcut(self) -> float: - return self._backend.get_rcut() - - def get_ntypes(self) -> int: - return self._backend.get_ntypes() - - def get_type_map(self) -> list[str]: - return self._backend.get_type_map() - - def get_dim_fparam(self) -> int: - return self._backend.get_dim_fparam() - - def has_default_fparam(self) -> bool: - return self._backend.has_default_fparam() - - def get_dim_aparam(self) -> int: - return self._backend.get_dim_aparam() - - @property - def model_type(self) -> type[DeepEval]: - return self._backend.model_type - - def get_sel_type(self) -> list[int]: - return self._backend.get_sel_type() - - def get_numb_dos(self) -> int: - return self._backend.get_numb_dos() - - def get_has_efield(self) -> bool: - return self._backend.get_has_efield() - - def get_has_spin(self) -> bool: - return self._backend.get_has_spin() - - def get_has_hessian(self) -> bool: - return self._backend.get_has_hessian() - - def get_var_name(self) -> str: - return self._backend.get_var_name() - - def get_ntypes_spin(self) -> int: - return self._backend.get_ntypes_spin() - - def get_model(self) -> Any: - return self._backend.get_model() - - return PretrainedDeepEvalBackend +class PretrainedDeepEvalBackend(DeepEvalBackend): + """Resolve alias and delegate to backend selected by resolved model path.""" + + def __init__( + self, + model_file: str, + output_def: object, + *args: object, + auto_batch_size: object = True, + neighbor_list: object | None = None, + **kwargs: object, + ) -> None: + model_name = parse_pretrained_alias(model_file) + resolved = str(resolve_model_path(model_name)) + + # DeepEvalBackend.__new__ dispatches by resolved suffix (.pt/.pb/.dp...) + self._backend = DeepEvalBackend( + resolved, + output_def, + *args, + auto_batch_size=auto_batch_size, + neighbor_list=neighbor_list, + **kwargs, + ) + + def eval( + self, + coords: np.ndarray, + cells: np.ndarray | None, + atom_types: np.ndarray, + atomic: bool = False, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, + **kwargs: Any, + ) -> dict[str, np.ndarray]: + return self._backend.eval( + coords, + cells, + atom_types, + atomic, + fparam=fparam, + aparam=aparam, + **kwargs, + ) + + def eval_descriptor( + self, + coords: np.ndarray, + cells: np.ndarray | None, + atom_types: np.ndarray, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, + efield: np.ndarray | None = None, + mixed_type: bool = False, + **kwargs: Any, + ) -> np.ndarray: + return self._backend.eval_descriptor( + coords, + cells, + atom_types, + fparam=fparam, + aparam=aparam, + efield=efield, + mixed_type=mixed_type, + **kwargs, + ) + + def eval_fitting_last_layer( + self, + coords: np.ndarray, + cells: np.ndarray | None, + atom_types: np.ndarray, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, + **kwargs: Any, + ) -> np.ndarray: + return self._backend.eval_fitting_last_layer( + coords, + cells, + atom_types, + fparam=fparam, + aparam=aparam, + **kwargs, + ) + + def get_rcut(self) -> float: + return self._backend.get_rcut() + + def get_ntypes(self) -> int: + return self._backend.get_ntypes() + + def get_type_map(self) -> list[str]: + return self._backend.get_type_map() + + def get_dim_fparam(self) -> int: + return self._backend.get_dim_fparam() + + def has_default_fparam(self) -> bool: + return self._backend.has_default_fparam() + + def get_dim_aparam(self) -> int: + return self._backend.get_dim_aparam() + + @property + def model_type(self) -> type[DeepEval]: + return self._backend.model_type + + def get_sel_type(self) -> list[int]: + return self._backend.get_sel_type() + + def get_numb_dos(self) -> int: + return self._backend.get_numb_dos() + + def get_has_efield(self) -> bool: + return self._backend.get_has_efield() + + def get_has_spin(self) -> bool: + return self._backend.get_has_spin() + + def get_has_hessian(self) -> bool: + return self._backend.get_has_hessian() + + def get_var_name(self) -> str: + return self._backend.get_var_name() + + def get_ntypes_spin(self) -> int: + return self._backend.get_ntypes_spin() + + def get_model(self) -> Any: + return self._backend.get_model() diff --git a/source/tests/common/test_pretrained_backend.py b/source/tests/common/test_pretrained_backend.py index 961296cf12..951b46cdfc 100644 --- a/source/tests/common/test_pretrained_backend.py +++ b/source/tests/common/test_pretrained_backend.py @@ -38,4 +38,8 @@ def test_parse_pretrained_alias_invalid(self) -> None: parse_pretrained_alias("DPA-3.2-5M.pt") def test_deep_eval_property(self) -> None: - self.assertIsNotNone(PretrainedBackend().deep_eval) + from deepmd.pretrained.deep_eval import ( + PretrainedDeepEvalBackend, + ) + + self.assertIs(PretrainedBackend().deep_eval, PretrainedDeepEvalBackend) From 397a451c75caa3530e0f2994f49e474b40dfe226 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 1 Mar 2026 22:02:40 +0000 Subject: [PATCH 06/11] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pretrained/deep_eval.py | 1 + 1 file changed, 1 insertion(+) diff --git a/deepmd/pretrained/deep_eval.py b/deepmd/pretrained/deep_eval.py index c722bcab5b..0a5db179f0 100644 --- a/deepmd/pretrained/deep_eval.py +++ b/deepmd/pretrained/deep_eval.py @@ -20,6 +20,7 @@ from deepmd.pretrained.download import ( resolve_model_path, ) + if TYPE_CHECKING: import numpy as np From 88db878e355302fb43af7dc3cd595ba6c4651ca9 Mon Sep 17 00:00:00 2001 From: njzjz-bot <48687836+njzjz-bot@users.noreply.github.com> Date: Sun, 1 Mar 2026 22:04:14 +0000 Subject: [PATCH 07/11] refactor(cli): move pretrained registry import to module scope Authored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.3-codex) --- deepmd/main.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/deepmd/main.py b/deepmd/main.py index d239ec80c2..3afcda8b4a 100644 --- a/deepmd/main.py +++ b/deepmd/main.py @@ -20,6 +20,9 @@ from deepmd.backend.backend import ( Backend, ) +from deepmd.pretrained.registry import ( + available_model_names, +) try: from deepmd._version import version as __version__ @@ -958,9 +961,6 @@ def main_parser() -> argparse.ArgumentParser: "download", help="Download one pretrained model", ) - from deepmd.pretrained.registry import ( - available_model_names, - ) parser_pretrained_download.add_argument( "MODEL", From e1c0f9a776e4eaa9a5eeb14fd5460c5b0b2f65a0 Mon Sep 17 00:00:00 2001 From: njzjz-bot <48687836+njzjz-bot@users.noreply.github.com> Date: Sun, 1 Mar 2026 22:33:05 +0000 Subject: [PATCH 08/11] fix(pretrained): accept case-insensitive .pretrained suffix - parse aliases with case-insensitive suffix check\n- add dedicated InvalidPretrainedAliasError\n- extend backend test to cover uppercase suffix\n\nAuthored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.3-codex) --- deepmd/pretrained/deep_eval.py | 13 ++++++++++--- source/tests/common/test_pretrained_backend.py | 4 ++++ 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/deepmd/pretrained/deep_eval.py b/deepmd/pretrained/deep_eval.py index 0a5db179f0..83853bbabe 100644 --- a/deepmd/pretrained/deep_eval.py +++ b/deepmd/pretrained/deep_eval.py @@ -25,16 +25,23 @@ import numpy as np +class InvalidPretrainedAliasError(ValueError): + """Raised when a pretrained alias string is malformed.""" + + def __init__(self, model_file: str) -> None: + super().__init__(f"Invalid pretrained alias: {model_file}") + + def parse_pretrained_alias(model_file: str) -> str: """Extract model name from ``*.pretrained`` alias string.""" alias = Path(model_file).name suffix = ".pretrained" - if not alias.endswith(suffix): - raise ValueError(f"Invalid pretrained alias: {model_file}") + if not alias.lower().endswith(suffix): + raise InvalidPretrainedAliasError(model_file) model_name = alias[: -len(suffix)] if not model_name: - raise ValueError(f"Invalid pretrained alias: {model_file}") + raise InvalidPretrainedAliasError(model_file) return model_name diff --git a/source/tests/common/test_pretrained_backend.py b/source/tests/common/test_pretrained_backend.py index 951b46cdfc..1daa9d26f8 100644 --- a/source/tests/common/test_pretrained_backend.py +++ b/source/tests/common/test_pretrained_backend.py @@ -32,6 +32,10 @@ def test_parse_pretrained_alias(self) -> None: parse_pretrained_alias("DPA-3.2-5M.pretrained"), "DPA-3.2-5M", ) + self.assertEqual( + parse_pretrained_alias("DPA-3.2-5M.PRETRAINED"), + "DPA-3.2-5M", + ) def test_parse_pretrained_alias_invalid(self) -> None: with self.assertRaises(ValueError): From 00a0ea2acc17671ad7052c6a9244baecbaf0cbbe Mon Sep 17 00:00:00 2001 From: "njzjz-bot (driven by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.3-codex))[bot]" <48687836+njzjz-bot@users.noreply.github.com> Date: Tue, 3 Mar 2026 09:34:47 +0000 Subject: [PATCH 09/11] docs(pretrained): add usage guide for dp pretrained - document dp pretrained download command and alias usage\n- include model/pretrained.md in docs index\n- keep CLI path output and add log message for visibility\n- clarify pretrained backend is an internal virtual alias backend\n\nAuthored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.3-codex) --- deepmd/backend/pretrained.py | 7 ++++- deepmd/pretrained/entrypoints.py | 2 ++ doc/model/index.rst | 1 + doc/model/pretrained.md | 48 ++++++++++++++++++++++++++++++++ 4 files changed, 57 insertions(+), 1 deletion(-) create mode 100644 doc/model/pretrained.md diff --git a/deepmd/backend/pretrained.py b/deepmd/backend/pretrained.py index f6f95d8c59..f165816eed 100644 --- a/deepmd/backend/pretrained.py +++ b/deepmd/backend/pretrained.py @@ -26,7 +26,12 @@ @Backend.register("pretrained") class PretrainedBackend(Backend): - """Backend for ``*.pretrained`` model aliases.""" + """Internal virtual backend for ``*.pretrained`` alias dispatch. + + This backend is not intended to be selected explicitly by users as a real + compute backend (such as TensorFlow/PyTorch/Paddle/JAX). It only bridges + ``*.pretrained`` aliases into the regular deep-eval loading path. + """ name = "Pretrained" features: ClassVar[Backend.Feature] = Backend.Feature.DEEP_EVAL diff --git a/deepmd/pretrained/entrypoints.py b/deepmd/pretrained/entrypoints.py index ba8a23de85..559f85e839 100644 --- a/deepmd/pretrained/entrypoints.py +++ b/deepmd/pretrained/entrypoints.py @@ -5,6 +5,7 @@ annotations, ) +import logging from pathlib import ( Path, ) @@ -25,6 +26,7 @@ def pretrained_entrypoint(args: argparse.Namespace) -> None: if args.pretrained_command == "download": cache_dir = Path(args.cache_dir) if args.cache_dir else None path = download_model(args.MODEL, cache_dir=cache_dir) + logging.getLogger(__name__).info("Pretrained model path: %s", path) print(path) # noqa: T201 return diff --git a/doc/model/index.rst b/doc/model/index.rst index 4896ccdc4e..a173732bbc 100644 --- a/doc/model/index.rst +++ b/doc/model/index.rst @@ -29,3 +29,4 @@ Model change-bias precision show-model-info + pretrained diff --git a/doc/model/pretrained.md b/doc/model/pretrained.md new file mode 100644 index 0000000000..cfee4a4afb --- /dev/null +++ b/doc/model/pretrained.md @@ -0,0 +1,48 @@ +# Use `dp pretrained` to download built-in models + +The `dp pretrained` command provides a simple way to download built-in pre-trained models and store them in a local cache. + +## Command syntax + +```bash +dp pretrained download [--cache-dir ] +``` + +- ``: the built-in model name. +- `--cache-dir `: optional cache directory. If omitted, DeePMD-kit uses the default cache path. + +## Available built-in models + +You can run `dp pretrained download -h` to see the currently supported model list in your installed version. + +Examples in this release include: + +- `DPA-3.2-5M` +- `DPA-3.1-3M` + +## Examples + +```bash +# Download to default cache directory +dp pretrained download DPA-3.2-5M + +# Download to a custom cache directory +dp pretrained download DPA-3.2-5M --cache-dir ./models +``` + +The command prints the local path of the downloaded model file on success. + +## Use downloaded models via alias + +After downloading, you can use the `.pretrained` alias directly in DeepEval/DeepPot workflows. + +For example: + +```python +from deepmd.infer import DeepPot + +# DeePMD-kit resolves this alias to the corresponding local model file +pot = DeepPot("DPA-3.2-5M.pretrained") +``` + +The `.pretrained` alias is designed for user-facing model selection, while backend details are handled internally. From 48ecbdf9aaf31e97433109a493179da5712c1d10 Mon Sep 17 00:00:00 2001 From: njzjz-bot <48687836+njzjz-bot@users.noreply.github.com> Date: Wed, 4 Mar 2026 15:32:54 +0000 Subject: [PATCH 10/11] feat(pretrained): support model-name aliases and clarify DeepPot docs - accept built-in model names without .pretrained suffix in DeepPot/DeepEval\n- register built-in model names as pretrained backend aliases\n- update docs: DeepPot does not require prior dp pretrained download\n- add tests for plain model-name alias resolution\n\nAuthored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.3-codex) --- deepmd/backend/pretrained.py | 12 +++++++- deepmd/pretrained/deep_eval.py | 30 ++++++++++++++----- doc/model/pretrained.md | 19 ++++++++---- .../tests/common/test_pretrained_backend.py | 8 +++++ 4 files changed, 56 insertions(+), 13 deletions(-) diff --git a/deepmd/backend/pretrained.py b/deepmd/backend/pretrained.py index f165816eed..50c66df913 100644 --- a/deepmd/backend/pretrained.py +++ b/deepmd/backend/pretrained.py @@ -10,6 +10,9 @@ from deepmd.backend.backend import ( Backend, ) +from deepmd.pretrained.registry import ( + available_model_names, +) if TYPE_CHECKING: from argparse import ( @@ -31,11 +34,18 @@ class PretrainedBackend(Backend): This backend is not intended to be selected explicitly by users as a real compute backend (such as TensorFlow/PyTorch/Paddle/JAX). It only bridges ``*.pretrained`` aliases into the regular deep-eval loading path. + + For convenience, all built-in pretrained model names are also registered as + suffix-like aliases, so users can pass model names directly, e.g. + ``DeepPot("DPA-3.2-5M")``. """ name = "Pretrained" features: ClassVar[Backend.Feature] = Backend.Feature.DEEP_EVAL - suffixes: ClassVar[list[str]] = [".pretrained"] + suffixes: ClassVar[list[str]] = [ + ".pretrained", + *[model_name.lower() for model_name in available_model_names()], + ] def is_available(self) -> bool: return True diff --git a/deepmd/pretrained/deep_eval.py b/deepmd/pretrained/deep_eval.py index 83853bbabe..410739000d 100644 --- a/deepmd/pretrained/deep_eval.py +++ b/deepmd/pretrained/deep_eval.py @@ -20,6 +20,9 @@ from deepmd.pretrained.download import ( resolve_model_path, ) +from deepmd.pretrained.registry import ( + MODEL_REGISTRY, +) if TYPE_CHECKING: import numpy as np @@ -33,17 +36,30 @@ def __init__(self, model_file: str) -> None: def parse_pretrained_alias(model_file: str) -> str: - """Extract model name from ``*.pretrained`` alias string.""" + """Extract model name from alias string. + + Accepted forms: + - ``.pretrained`` (case-insensitive suffix) + - ```` where ```` is a built-in registry name + """ alias = Path(model_file).name suffix = ".pretrained" - if not alias.lower().endswith(suffix): - raise InvalidPretrainedAliasError(model_file) - model_name = alias[: -len(suffix)] - if not model_name: - raise InvalidPretrainedAliasError(model_file) + if alias.lower().endswith(suffix): + model_name = alias[: -len(suffix)] + if not model_name: + raise InvalidPretrainedAliasError(model_file) + return model_name + + if alias in MODEL_REGISTRY: + return alias + + lowered = alias.lower() + for model_name in MODEL_REGISTRY: + if model_name.lower() == lowered: + return model_name - return model_name + raise InvalidPretrainedAliasError(model_file) class PretrainedDeepEvalBackend(DeepEvalBackend): diff --git a/doc/model/pretrained.md b/doc/model/pretrained.md index cfee4a4afb..abaa81d725 100644 --- a/doc/model/pretrained.md +++ b/doc/model/pretrained.md @@ -32,17 +32,26 @@ dp pretrained download DPA-3.2-5M --cache-dir ./models The command prints the local path of the downloaded model file on success. -## Use downloaded models via alias +## Use downloaded models in DeepPot -After downloading, you can use the `.pretrained` alias directly in DeepEval/DeepPot workflows. +Using `DeepPot`, you do **not** have to run `dp pretrained download` first. -For example: +You can pass either the model alias directly (recommended): + +```python +from deepmd.infer import DeepPot + +pot = DeepPot("DPA-3.2-5M") +``` + +or the explicit alias with suffix: ```python from deepmd.infer import DeepPot -# DeePMD-kit resolves this alias to the corresponding local model file pot = DeepPot("DPA-3.2-5M.pretrained") ``` -The `.pretrained` alias is designed for user-facing model selection, while backend details are handled internally. +If the model file is not already present in the local cache, DeePMD-kit will download and cache it automatically when resolving the alias. + +The `.pretrained` alias and plain model names are user-facing selectors. Backend details are handled internally. diff --git a/source/tests/common/test_pretrained_backend.py b/source/tests/common/test_pretrained_backend.py index 1daa9d26f8..dddde1271f 100644 --- a/source/tests/common/test_pretrained_backend.py +++ b/source/tests/common/test_pretrained_backend.py @@ -27,6 +27,10 @@ def test_detect_backend_by_pretrained_suffix(self) -> None: backend = Backend.detect_backend_by_model("DPA-3.2-5M.pretrained") self.assertIs(backend, PretrainedBackend) + def test_detect_backend_by_model_name(self) -> None: + backend = Backend.detect_backend_by_model("DPA-3.2-5M") + self.assertIs(backend, PretrainedBackend) + def test_parse_pretrained_alias(self) -> None: self.assertEqual( parse_pretrained_alias("DPA-3.2-5M.pretrained"), @@ -37,6 +41,10 @@ def test_parse_pretrained_alias(self) -> None: "DPA-3.2-5M", ) + def test_parse_pretrained_alias_plain_name(self) -> None: + self.assertEqual(parse_pretrained_alias("DPA-3.2-5M"), "DPA-3.2-5M") + self.assertEqual(parse_pretrained_alias("dpa-3.2-5m"), "DPA-3.2-5M") + def test_parse_pretrained_alias_invalid(self) -> None: with self.assertRaises(ValueError): parse_pretrained_alias("DPA-3.2-5M.pt") From 386ff68087ee42d077987023362fc8d98a3551f5 Mon Sep 17 00:00:00 2001 From: njzjz-bot <48687836+njzjz-bot@users.noreply.github.com> Date: Wed, 4 Mar 2026 15:56:34 +0000 Subject: [PATCH 11/11] refactor(pretrained): drop .pretrained compatibility --- deepmd/backend/pretrained.py | 7 +++---- deepmd/pretrained/deep_eval.py | 16 ++++------------ doc/model/pretrained.md | 14 +++----------- source/tests/common/test_pretrained_backend.py | 18 +++++------------- 4 files changed, 15 insertions(+), 40 deletions(-) diff --git a/deepmd/backend/pretrained.py b/deepmd/backend/pretrained.py index 50c66df913..f6233fc3a0 100644 --- a/deepmd/backend/pretrained.py +++ b/deepmd/backend/pretrained.py @@ -29,13 +29,13 @@ @Backend.register("pretrained") class PretrainedBackend(Backend): - """Internal virtual backend for ``*.pretrained`` alias dispatch. + """Internal virtual backend for pretrained model-name alias dispatch. This backend is not intended to be selected explicitly by users as a real compute backend (such as TensorFlow/PyTorch/Paddle/JAX). It only bridges - ``*.pretrained`` aliases into the regular deep-eval loading path. + built-in pretrained model names into the regular deep-eval loading path. - For convenience, all built-in pretrained model names are also registered as + For convenience, all built-in pretrained model names are registered as suffix-like aliases, so users can pass model names directly, e.g. ``DeepPot("DPA-3.2-5M")``. """ @@ -43,7 +43,6 @@ class PretrainedBackend(Backend): name = "Pretrained" features: ClassVar[Backend.Feature] = Backend.Feature.DEEP_EVAL suffixes: ClassVar[list[str]] = [ - ".pretrained", *[model_name.lower() for model_name in available_model_names()], ] diff --git a/deepmd/pretrained/deep_eval.py b/deepmd/pretrained/deep_eval.py index 410739000d..2dc671b0cc 100644 --- a/deepmd/pretrained/deep_eval.py +++ b/deepmd/pretrained/deep_eval.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""DeepEval adapter for `*.pretrained` model aliases.""" +"""DeepEval adapter for pretrained model-name aliases.""" from __future__ import ( annotations, @@ -32,24 +32,16 @@ class InvalidPretrainedAliasError(ValueError): """Raised when a pretrained alias string is malformed.""" def __init__(self, model_file: str) -> None: - super().__init__(f"Invalid pretrained alias: {model_file}") + super().__init__(f"Invalid pretrained model name: {model_file}") def parse_pretrained_alias(model_file: str) -> str: - """Extract model name from alias string. + """Extract built-in pretrained model name from alias string. - Accepted forms: - - ``.pretrained`` (case-insensitive suffix) + Accepted form: - ```` where ```` is a built-in registry name """ alias = Path(model_file).name - suffix = ".pretrained" - - if alias.lower().endswith(suffix): - model_name = alias[: -len(suffix)] - if not model_name: - raise InvalidPretrainedAliasError(model_file) - return model_name if alias in MODEL_REGISTRY: return alias diff --git a/doc/model/pretrained.md b/doc/model/pretrained.md index abaa81d725..e4fd3481de 100644 --- a/doc/model/pretrained.md +++ b/doc/model/pretrained.md @@ -36,7 +36,7 @@ The command prints the local path of the downloaded model file on success. Using `DeepPot`, you do **not** have to run `dp pretrained download` first. -You can pass either the model alias directly (recommended): +Pass the built-in model name directly: ```python from deepmd.infer import DeepPot @@ -44,14 +44,6 @@ from deepmd.infer import DeepPot pot = DeepPot("DPA-3.2-5M") ``` -or the explicit alias with suffix: +If the model file is not already present in the local cache, DeePMD-kit will download and cache it automatically when resolving the model name. -```python -from deepmd.infer import DeepPot - -pot = DeepPot("DPA-3.2-5M.pretrained") -``` - -If the model file is not already present in the local cache, DeePMD-kit will download and cache it automatically when resolving the alias. - -The `.pretrained` alias and plain model names are user-facing selectors. Backend details are handled internally. +Built-in model names are user-facing selectors; backend details are handled internally. diff --git a/source/tests/common/test_pretrained_backend.py b/source/tests/common/test_pretrained_backend.py index dddde1271f..e1502355b2 100644 --- a/source/tests/common/test_pretrained_backend.py +++ b/source/tests/common/test_pretrained_backend.py @@ -23,23 +23,13 @@ def setUpClass(cls) -> None: # ensure backend registration side effects are loaded importlib.import_module("deepmd.backend") - def test_detect_backend_by_pretrained_suffix(self) -> None: - backend = Backend.detect_backend_by_model("DPA-3.2-5M.pretrained") - self.assertIs(backend, PretrainedBackend) - def test_detect_backend_by_model_name(self) -> None: backend = Backend.detect_backend_by_model("DPA-3.2-5M") self.assertIs(backend, PretrainedBackend) - def test_parse_pretrained_alias(self) -> None: - self.assertEqual( - parse_pretrained_alias("DPA-3.2-5M.pretrained"), - "DPA-3.2-5M", - ) - self.assertEqual( - parse_pretrained_alias("DPA-3.2-5M.PRETRAINED"), - "DPA-3.2-5M", - ) + def test_detect_backend_by_pretrained_suffix_not_supported(self) -> None: + with self.assertRaises(ValueError): + Backend.detect_backend_by_model("DPA-3.2-5M.pretrained") def test_parse_pretrained_alias_plain_name(self) -> None: self.assertEqual(parse_pretrained_alias("DPA-3.2-5M"), "DPA-3.2-5M") @@ -48,6 +38,8 @@ def test_parse_pretrained_alias_plain_name(self) -> None: def test_parse_pretrained_alias_invalid(self) -> None: with self.assertRaises(ValueError): parse_pretrained_alias("DPA-3.2-5M.pt") + with self.assertRaises(ValueError): + parse_pretrained_alias("DPA-3.2-5M.pretrained") def test_deep_eval_property(self) -> None: from deepmd.pretrained.deep_eval import (