feat(pretrained): add built-in pretrained downloader and alias backend#5277
feat(pretrained): add built-in pretrained downloader and alias backend#5277njzjz-bot wants to merge 11 commits intodeepmodeling:masterfrom
Conversation
- add dp pretrained download <MODEL> 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)
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)
Authored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.3-codex)
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds a new "pretrained" subsystem: backend registration, CLI subcommand and entrypoint, model registry, secure download/caching with checksum and URL ranking, and a DeepEval adapter resolving *.pretrained aliases to concrete backends. Includes unit tests for parser, backend, and download behavior. Changes
Sequence Diagram(s)sequenceDiagram
actor User
participant CLI as "CLI Parser"
participant Entrypoint as "pretrained_entrypoint"
participant Download as "download_model()"
participant Registry as "MODEL_REGISTRY"
participant URLRank as "URL Ranker"
participant HTTP as "HTTP Client"
participant Cache as "Cache Dir"
User->>CLI: deepmd pretrained download MODEL
CLI->>Entrypoint: invoke pretrained_entrypoint(args)
Entrypoint->>Download: request MODEL, cache_dir
Download->>Registry: lookup model metadata
Registry-->>Download: {urls, filename, sha256}
Download->>URLRank: probe & rank candidate URLs (parallel)
URLRank->>HTTP: probe URLs
HTTP-->>URLRank: latencies / reachability
URLRank-->>Download: sorted URLs
loop try URLs until success
Download->>HTTP: download from next URL
HTTP-->>Download: stream file
Download->>Download: verify SHA256
end
Download->>Cache: atomic write -> final path
Cache-->>Download: cached file path
Download-->>Entrypoint: return Path
Entrypoint-->>User: print model path
Estimated code review effort🎯 4 (Complex) | ⏱️ ~55 minutes Suggested reviewers
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Tip Try Coding Plans. Let us write the prompt for your AI agent so you can ship faster (with fewer bugs). Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
deepmd/pretrained/download.py (1)
196-214: Consider simplifying by delegating todownload_model.The cache check logic (lines 209-212) duplicates what
download_modelalready handles (lines 146-150). Sincedownload_modelreturns early when the cached file is valid,resolve_model_pathcould simply call it directly.♻️ Suggested simplification
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) + return download_model(model_name, cache_dir=cache_dir, logger=logger)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/pretrained/download.py` around lines 196 - 214, The cache existence/validation logic in resolve_model_path duplicates download_model; simplify by looking up the model with MODEL_REGISTRY inside resolve_model_path (raise the same ValueError if missing), compute target_dir = cache_dir or DEFAULT_CACHE_DIR, then return download_model(model_name, cache_dir=target_dir, logger=logger) directly so download_model handles the cached-file SHA256 check and download; keep symbols resolve_model_path, MODEL_REGISTRY, DEFAULT_CACHE_DIR, and download_model to locate and update the code.source/tests/common/test_pretrained_download.py (1)
83-107: Consider adding assertion that download is not triggered for cached files.The test validates the return path but doesn't verify that
download_modelwas never called. Adding a mock/patch assertion would strengthen coverage.♻️ Suggested improvement
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) + ), patch.object(dl, "download_model") as mock_download: + path = dl.resolve_model_path(model_name, cache_dir=cache_dir) + mock_download.assert_not_called() self.assertEqual(path, target)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/common/test_pretrained_download.py` around lines 83 - 107, Update the test_resolve_model_path_cached test to assert that the download path is not exercised by patching or mocking dl.download_model and asserting it was not called; specifically, inside the with patch.object(dl, "MODEL_REGISTRY", {...}) block wrap or patch dl.download_model (or use unittest.mock.patch.object on the download_model symbol) and after calling dl.resolve_model_path(model_name, cache_dir=cache_dir) assert download_model.assert_not_called() so the test verifies resolve_model_path returns the cached target and does not invoke download_model.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@deepmd/pretrained/backend.py`:
- Around line 54-145: PretrainedDeepEvalBackend is missing delegations for
eval_descriptor and eval_fitting_last_layer so calls get the default
NotImplementedError instead of the underlying backend implementation; add
methods eval_descriptor(...) and eval_fitting_last_layer(...) on
PretrainedDeepEvalBackend that forward all parameters and kwargs to
self._backend.eval_descriptor(...) and
self._backend.eval_fitting_last_layer(...) respectively (matching the same
signatures used by DeepEvalBackend) so that backends that implement these
methods (e.g., PyTorch/TensorFlow) are invoked via the _backend proxy.
---
Nitpick comments:
In `@deepmd/pretrained/download.py`:
- Around line 196-214: The cache existence/validation logic in
resolve_model_path duplicates download_model; simplify by looking up the model
with MODEL_REGISTRY inside resolve_model_path (raise the same ValueError if
missing), compute target_dir = cache_dir or DEFAULT_CACHE_DIR, then return
download_model(model_name, cache_dir=target_dir, logger=logger) directly so
download_model handles the cached-file SHA256 check and download; keep symbols
resolve_model_path, MODEL_REGISTRY, DEFAULT_CACHE_DIR, and download_model to
locate and update the code.
In `@source/tests/common/test_pretrained_download.py`:
- Around line 83-107: Update the test_resolve_model_path_cached test to assert
that the download path is not exercised by patching or mocking dl.download_model
and asserting it was not called; specifically, inside the with patch.object(dl,
"MODEL_REGISTRY", {...}) block wrap or patch dl.download_model (or use
unittest.mock.patch.object on the download_model symbol) and after calling
dl.resolve_model_path(model_name, cache_dir=cache_dir) assert
download_model.assert_not_called() so the test verifies resolve_model_path
returns the cached target and does not invoke download_model.
ℹ️ Review info
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (11)
deepmd/backend/pretrained.pydeepmd/entrypoints/main.pydeepmd/main.pydeepmd/pretrained/__init__.pydeepmd/pretrained/backend.pydeepmd/pretrained/download.pydeepmd/pretrained/entrypoints.pydeepmd/pretrained/registry.pysource/tests/common/test_pretrained_backend.pysource/tests/common/test_pretrained_download.pysource/tests/common/test_pretrained_parser.py
…ments - 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)
There was a problem hiding this comment.
🧹 Nitpick comments (1)
deepmd/pretrained/download.py (1)
196-209: Minor redundancy in model validation.The registry check at lines 204-207 duplicates the validation already performed in
download_model. While not incorrect, this could be simplified by removing the duplicate check here sincedownload_modelwill raise the same error.♻️ Optional simplification
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}") - return download_model(model_name, cache_dir=target_dir, logger=logger)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/pretrained/download.py` around lines 196 - 209, The resolve_model_path function currently duplicates model existence validation already performed by download_model; remove the MODEL_REGISTRY lookup and associated ValueError branch from resolve_model_path and simply delegate to download_model (keeping target_dir/cache_dir and logger forwarding) so download_model remains the single source of truth for verifying model_name (refer to resolve_model_path, download_model, and MODEL_REGISTRY to locate the logic).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@deepmd/pretrained/download.py`:
- Around line 196-209: The resolve_model_path function currently duplicates
model existence validation already performed by download_model; remove the
MODEL_REGISTRY lookup and associated ValueError branch from resolve_model_path
and simply delegate to download_model (keeping target_dir/cache_dir and logger
forwarding) so download_model remains the single source of truth for verifying
model_name (refer to resolve_model_path, download_model, and MODEL_REGISTRY to
locate the logic).
ℹ️ Review info
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
deepmd/backend/pretrained.pydeepmd/pretrained/deep_eval.pydeepmd/pretrained/download.pysource/tests/common/test_pretrained_backend.pysource/tests/common/test_pretrained_download.py
🚧 Files skipped from review as they are similar to previous changes (2)
- deepmd/backend/pretrained.py
- source/tests/common/test_pretrained_backend.py
…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)
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
source/tests/common/test_pretrained_backend.py (1)
36-39: Add a regression test for uppercase.PRETRAINEDsuffix.Please add a case like
parse_pretrained_alias("DPA-3.2-5M.PRETRAINED")to keep parser behavior aligned with backend suffix detection rules and prevent regressions.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/common/test_pretrained_backend.py` around lines 36 - 39, Add a regression case to the existing test_parse_pretrained_alias_invalid: call parse_pretrained_alias("DPA-3.2-5M.PRETRAINED") inside the same assertRaises(ValueError) context so the parser is validated to reject the uppercase .PRETRAINED suffix; update the test in source/tests/common/test_pretrained_backend.py alongside test_parse_pretrained_alias_invalid using the same pattern to ensure behavior matches backend suffix detection rules.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@deepmd/pretrained/deep_eval.py`:
- Around line 31-34: The suffix check for pretrained aliases is currently
case-sensitive; change it to compare lowercased values so aliases like
DPA-3.2-5M.PRETRAINED pass. Use alias_lower = alias.lower() (or equivalent) and
test alias_lower.endswith(suffix) while keeping the original alias when slicing
to produce model_name (alias, suffix, model_file, model_name are the identifiers
to update).
---
Nitpick comments:
In `@source/tests/common/test_pretrained_backend.py`:
- Around line 36-39: Add a regression case to the existing
test_parse_pretrained_alias_invalid: call
parse_pretrained_alias("DPA-3.2-5M.PRETRAINED") inside the same
assertRaises(ValueError) context so the parser is validated to reject the
uppercase .PRETRAINED suffix; update the test in
source/tests/common/test_pretrained_backend.py alongside
test_parse_pretrained_alias_invalid using the same pattern to ensure behavior
matches backend suffix detection rules.
ℹ️ Review info
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
deepmd/backend/pretrained.pydeepmd/pretrained/deep_eval.pysource/tests/common/test_pretrained_backend.py
🚧 Files skipped from review as they are similar to previous changes (1)
- deepmd/backend/pretrained.py
Authored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.3-codex)
There was a problem hiding this comment.
♻️ Duplicate comments (1)
deepmd/pretrained/deep_eval.py (1)
31-34:⚠️ Potential issue | 🟡 MinorMake alias suffix matching case-insensitive.
The suffix guard is still case-sensitive, so uppercase aliases can be detected upstream but rejected here.
💡 Proposed fix
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): + if not alias.lower().endswith(suffix): raise ValueError(f"Invalid pretrained alias: {model_file}")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/pretrained/deep_eval.py` around lines 31 - 34, The suffix check is case-sensitive; update the guard that uses alias and suffix to perform a case-insensitive comparison (e.g., compare alias.lower().endswith(suffix.lower())) and when raising the ValueError, report the offending alias variable (alias) rather than model_file to give the correct context; adjust the check around the existing suffix, alias, and ValueError usage in deep_eval.py accordingly.
🧹 Nitpick comments (1)
deepmd/pretrained/deep_eval.py (1)
33-37: Address Ruff TRY003 for repeated inline exception messages.Consider using a small dedicated exception class and raise that instead of repeating formatted messages inline.
♻️ Proposed refactor
+class InvalidPretrainedAliasError(ValueError): + """Raised when a pretrained alias 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: @@ - if not alias.lower().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)As per coding guidelines,
**/*.py: Always runruff check .andruff format .before committing changes or CI will fail.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/pretrained/deep_eval.py` around lines 33 - 37, Replace the repeated inline ValueError messages in deep_eval.py by defining a small dedicated exception class (e.g., PretrainedAliasError) near the top of the module and raising that class from the code paths that currently call raise ValueError(f"Invalid pretrained alias: {model_file}") (both where model_file validation fails and where model_name is empty); update the two occurrences to raise PretrainedAliasError(model_file) (or a descriptive message) and run ruff check . and ruff format . before committing to satisfy lint/format rules.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@deepmd/pretrained/deep_eval.py`:
- Around line 31-34: The suffix check is case-sensitive; update the guard that
uses alias and suffix to perform a case-insensitive comparison (e.g., compare
alias.lower().endswith(suffix.lower())) and when raising the ValueError, report
the offending alias variable (alias) rather than model_file to give the correct
context; adjust the check around the existing suffix, alias, and ValueError
usage in deep_eval.py accordingly.
---
Nitpick comments:
In `@deepmd/pretrained/deep_eval.py`:
- Around line 33-37: Replace the repeated inline ValueError messages in
deep_eval.py by defining a small dedicated exception class (e.g.,
PretrainedAliasError) near the top of the module and raising that class from the
code paths that currently call raise ValueError(f"Invalid pretrained alias:
{model_file}") (both where model_file validation fails and where model_name is
empty); update the two occurrences to raise PretrainedAliasError(model_file) (or
a descriptive message) and run ruff check . and ruff format . before committing
to satisfy lint/format rules.
- 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)
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5277 +/- ##
==========================================
+ Coverage 82.00% 82.17% +0.17%
==========================================
Files 750 760 +10
Lines 75215 76261 +1046
Branches 3615 3659 +44
==========================================
+ Hits 61680 62670 +990
- Misses 12372 12422 +50
- Partials 1163 1169 +6 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
wanghan-iapcm
left a comment
There was a problem hiding this comment.
Please add doc for how to use dp pretrained
njzjz-bot
left a comment
There was a problem hiding this comment.
Addressed the latest review feedback in commit 60882f5.
- Added end-user documentation for
dp pretrainedusage:doc/model/pretrained.mdand linked it fromdoc/model/index.rst. - Kept CLI
print(path)output intentionally for shell/script usage and added an info log inpretrained_entrypoint. - Clarified
deepmd/backend/pretrained.pyas an internal virtual backend used only for.pretrainedalias dispatch (not a user-selectable compute backend).
Thanks for the suggestions.
60882f5 to
6b19157
Compare
- 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)
6b19157 to
00a0ea2
Compare
- 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)
Summary
This PR integrates pretrained model support directly into
deepmd-kitunderdeepmd/pretrained, while keepingDeepPotusage unchanged.Added
dp pretrained download <MODEL>deepmd/pretrained/registry.py,download.py,backend.py,entrypoints.pyDPA-3.2-5MDPA-3.1-3M.partwrites.pretrainedbackend alias support viadeepmd/backend/pretrained.pyDeepPot("DPA-3.2-5M.pretrained")while keeping existing DeepPot API unchangedCLI wiring
pretrainedparser/subparser indeepmd/main.pydeepmd/entrypoints/main.pyTests
source/tests/common/test_pretrained_parser.pysource/tests/common/test_pretrained_download.pysource/tests/common/test_pretrained_backend.pyFormatting / lint
uvx prek run --all-filesand committed auto-format updates.Authored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.3-codex)
Summary by CodeRabbit
New Features
pretrained download(optional --cache-dir) to fetch pretrained models..pretrainedaliases work transparently with existing evaluation flow (some backend hooks intentionally unsupported).Downloads & Caching
Tests