Refactor CI test environments / Add PET models - #90
Conversation
* Add PET-OAM * Apply suggestion from review * Apply one more suggestion * Add PET-MAD * Add package to optional dependencies * adjust deps group * ruff format * bump versions; fix tests * remove 3.10, add 3.13 * fix diatomics output path * move stability flow under mlip_arena.flows and benchmarks/stability.py * rollback dependencies * Refactor CI test environments to split conflicting models * Refactor CI test dependencies: move model deps to optional group 'all' and revise CI install scripts * ci: disable fail-fast in test matrix strategy * ci: exclude python 3.13 for GNN/compiled model test groups and skip torch-scatter/torch-sparse compilation * Refactor: remove torch_geometric and torch_scatter dependencies, implement local scatter helpers, and mark MACE test workflows * move torch_dftd and matscipy to base dep * Fix default CI test marker selection * refactor: apply PyTorch/e3nn compatibility patches at model calculator level * refactor: consolidate asymptotes and shifts benchmarks into submit.py * refactor: implement modular model-level reporting and refactor stability aggregation * refactor: extract model name and calculator resolution to DRY helpers in stability.py * refactor: simplify calculator resolution and restrict e3nn monkey-patch to import phase * test stability benchmark run * move stability from tasks to flows * handle cache separately for each test group in CI * test external calculators in different groups separately in CI * mark fairchem uma as xfail * separate orb-models into a separate test group and allow version <0.8.0 * add to() method to custom Data class in collate.py * fix parameter binding for EOSFlow and EVFlow in submit.py * pin orb-models to 0.4.0 in pyproject.toml * gracefully handle pytest exit code 5 (no tests collected) in CI * set torch to 2.2 for orb test group and exclude python 3.13 in CI * monkeypatch typing.io in nequip model and exclude orb test group on py3.13 * remove e3nn optimization defaults workaround in nequip.py * remove nequip patch and exclude nequip+py3.13 in CI * pin orb-models to 0.5.5 and update torch dep to 2.8 in install-linux.sh * increase test coverage * disable cache * refactor(deps): loosen prefect constraint and comment out ipykernel/ipywidgets * ci: disable caching in setup-uv step * test: switch prefect tests skip condition from python 3.11 to 3.12 * test: run test_stability.py flows inside prefect_test_harness on python 3.12 * test: isolate PREFECT_HOME in conftest.py to prevent concurrent database lock race conditions * ci: add Free Disk Space step and update fairchem-core version to 2.21.0 in pyproject.toml * test: remove xfail handler for UMA-S-1P1 to enable testing on CPU in CI * test: set PREFECT_HOME env variables at top-level of conftest.py to ensure isolation before prefect is imported * test: classify uma family under fairchem test group * refactor(deps): restore prefect version pin constraint <3.7.0 * ci: restore -n 5 flag to pytest command to run tests in parallel processes * refactor(deps): pin fastapi<0.115.0 to fix PrefectRouter AttributeError in Prefect < 3.7.0 * ci: disable pytest-xdist and upgrade prefect dependency to avoid fastapi pin * fix(stability): disable prefect persistence and refresh cache; update prefect/fastapi constraints --------- Co-authored-by: Filippo Bigi <98903385+frostedoyster@users.noreply.github.com> Co-authored-by: frostedoyster <bigi.f@libero.it> Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
📝 WalkthroughWalkthroughThis PR adds PET-OAM and PET-MAD calculator models to the registry, consolidates stability task utilities from ChangesMLIP Arena — PET models, stability consolidation, dependency removal, CI overhaul
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Possibly related PRs
Suggested labels
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 Biome (2.4.16)benchmarks/diatomics/nequip/NequIP-OAM-L.jsonFile contains syntax errors that prevent linting: Line 1: String values must be double quoted.; Line 1: String values must be double quoted.; Line 1: End of file expected; Line 3: End of file expected 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 10
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/test_external_calculators.py (1)
38-53:⚠️ Potential issue | 🟠 Major | ⚡ Quick winHandle
ImportErrorfromMLIPEnum.load()in the skip path.
MLIPEnum.load()can raiseImportErrorfor missing optional model packages, but this branch now falls into the genericExceptionhandler and fails the test instead of skipping. That makes this test brittle in dependency-grouped environments.Proposed fix
except ( + ImportError, GatedRepoError, HTTPError, HTTPStatusError, FileNotFoundError, ) as e: pytest.skip(str(e))🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/test_external_calculators.py` around lines 38 - 53, The ImportError exception raised by MLIPEnum.load() for missing optional model packages is not being caught in the skip path and instead falls through to the generic Exception handler, causing the test to fail or xfail instead of being skipped. Add ImportError to the except clause that currently catches GatedRepoError, HTTPError, HTTPStatusError, and FileNotFoundError so that ImportError is also handled in the skip path with pytest.skip().
🧹 Nitpick comments (1)
mlip_arena/models/classicals/zbl.py (1)
18-32: 💤 Low valueConsider documenting the
dimdefault difference fromtorch_scatter.The implementation is correct. Note that the default
dim=0differs fromtorch_scatter.scatter_add's default ofdim=-1. Since all current call sites explicitly passdim=0, this isn't an issue now, but could surprise future callers expectingtorch_scatter-compatible defaults.A brief docstring noting this local function's contract would help maintainability.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@mlip_arena/models/classicals/zbl.py` around lines 18 - 32, Add a docstring to the scatter_add function to document its behavior and clarify that the default parameter dim=0 differs from torch_scatter.scatter_add (which defaults to dim=-1). The docstring should explain what the function does, document the parameters (src, index, dim, dim_size), and note this behavioral difference to help future maintainers and callers understand the local implementation's contract.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In @.github/workflows/ci.yaml:
- Around line 16-17: Replace the mutable version tag `@v4` in the actions/checkout
step with a full commit SHA to ensure supply-chain security and consistency.
Additionally, add persist-credentials: false to the checkout step to prevent
GitHub tokens from being passed to subsequent workflow steps. After fixing the
checkout step, audit all other action uses entries in the workflow file
(including any `@v6`, `@v5`, `@v1.0` references) and replace their mutable version
tags with specific commit SHAs as well.
In `@mlip_arena/flows/diatomics.py`:
- Around line 26-28: The `calculator_kwargs` parameter in both the
`homonuclear_diatomic` function (at lines 26-28) and the corresponding function
at lines 259-261 should be optional with a default value of None. Change the
parameter signature from `calculator_kwargs: dict | None` to `calculator_kwargs:
dict | None = None` in both function definitions to make the parameter truly
optional and maintain consistency with the optional-kwargs contract used
elsewhere in the codebase.
In `@mlip_arena/flows/eos_bulk.py`:
- Line 83: The line 83 in the eos_bulk.py file silently ignores
calculator_kwargs when calculator is already a BaseCalculator instance, which
can lead to misleading behavior. Add validation to check if calculator is a
BaseCalculator instance and calculator_kwargs is provided; if both conditions
are true, raise an error to fail fast and alert the user that they cannot pass
calculator_kwargs with an already-instantiated BaseCalculator. If calculator is
a BaseCalculator instance without calculator_kwargs, use it as-is; otherwise,
call get_calculator with the provided kwargs to instantiate a new calculator.
In `@mlip_arena/flows/stability.py`:
- Around line 196-205: The download logic in the stability.py file is
conditioned only on whether the database file exists locally, completely
ignoring the force_download parameter. Modify the condition that guards the
hf_hub_download call to also check the force_download flag, so that the download
executes when either the file does not exist OR when force_download is True.
This ensures that callers can force a fresh download of the database by passing
force_download=True, even if a local copy already exists.
- Line 295: The code accessing target_steps uses off-by-one indexing with
traj[1], which causes an IndexError when the trajectory has only a single frame
(only index 0 exists), resulting in the file being silently dropped during
processing. Change the indexing from traj[1] to traj[0] in the line that reads
target_steps = traj[1].info["target_steps"] to correctly access the first frame
of the trajectory instead of attempting to access a second frame that may not
exist.
In `@mlip_arena/models/externals/ani.py`:
- Around line 8-11: The __new__ method accepts **kwargs but does not forward
them to the model.ase() call on the return statement. Update the return
statement to pass **kwargs to the ase() method so that callers can configure
parameters like overwrite and stress_kind as intended.
In `@mlip_arena/models/registry.yaml`:
- Around line 441-446: The PET-MAD checkpoint name in the registry.yaml file
declares pet-mad-s-v1.0.2.ckpt while the version hardcoded in
mlip_arena/models/externals/pet.py at line 30 shows version="1.5.0", creating a
provenance mismatch. Update the checkpoint value in the PET-MAD registry entry
to align with the actual version string from the pet.py file, ensuring the
declared checkpoint artifact metadata matches the runtime model version being
loaded.
- Line 391: The DeepMD entry in the model registry has a deprecated flag set to
true, but it lacks an explanatory comment documenting the reason for
deprecation, unlike other deprecated entries such as ORB which include
descriptive comments. Add a comment next to the deprecated: true flag in the
DeepMD entry explaining the reason for its deprecation to maintain consistency
with other deprecated model entries and improve maintainability.
In `@scripts/install-linux.sh`:
- Line 7: The extras specifiers in the `uv pip install` commands are not quoted,
which allows the shell to interpret the square brackets `[]` as glob character
classes. Quote each extras argument (the part containing `.[test,extra,nequip]`
and similar patterns) by wrapping it in single or double quotes so that the
brackets are passed literally to `uv pip` rather than being expanded as shell
glob patterns. This applies to all occurrences of the `uv pip install` command
where extras are specified with unquoted bracket notation.
In `@tests/conftest.py`:
- Around line 8-12: Replace the deterministic temp_dir path construction
(currently using tempfile.gettempdir() with worker_id) with tempfile.mkdtemp()
to generate a unique directory for each test run, ensuring test isolation and
preventing stale Prefect state leakage from previous runs. Update the temp_dir
assignment to use mkdtemp(), which will automatically create a unique temporary
directory and eliminate the reuse of old prefect.db content across test runs.
---
Outside diff comments:
In `@tests/test_external_calculators.py`:
- Around line 38-53: The ImportError exception raised by MLIPEnum.load() for
missing optional model packages is not being caught in the skip path and instead
falls through to the generic Exception handler, causing the test to fail or
xfail instead of being skipped. Add ImportError to the except clause that
currently catches GatedRepoError, HTTPError, HTTPStatusError, and
FileNotFoundError so that ImportError is also handled in the skip path with
pytest.skip().
---
Nitpick comments:
In `@mlip_arena/models/classicals/zbl.py`:
- Around line 18-32: Add a docstring to the scatter_add function to document its
behavior and clarify that the default parameter dim=0 differs from
torch_scatter.scatter_add (which defaults to dim=-1). The docstring should
explain what the function does, document the parameters (src, index, dim,
dim_size), and note this behavioral difference to help future maintainers and
callers understand the local implementation's contract.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 9f8f35f7-7479-4d83-94a9-cfca7c3ce19e
📒 Files selected for processing (34)
.github/workflows/ci.yamlbenchmarks/diatomics/nequip/NequIP-OAM-L.jsonbenchmarks/stability/aggregate.pybenchmarks/stability/run.pybenchmarks/submit.pymlip_arena/data/collate.pymlip_arena/flows/diatomics.pymlip_arena/flows/eos_bulk.pymlip_arena/flows/stability.pymlip_arena/models/classicals/zbl.pymlip_arena/models/externals/ani.pymlip_arena/models/externals/nequip.pymlip_arena/models/externals/pet.pymlip_arena/models/registry.yamlmlip_arena/tasks/stability/__init__.pymlip_arena/tasks/stability/analysis.pymlip_arena/tasks/stability/data.pymlip_arena/tasks/stability/flow.pypyproject.tomlscripts/install-linux.shscripts/install-perlmutter.shserve/ranks/combustion.pyserve/ranks/homonuclear-diatomics.pytests/conftest.pytests/test_data_integrity.pytests/test_data_local.pytests/test_elasticity.pytests/test_eos.pytests/test_external_calculators.pytests/test_md.pytests/test_mof.pytests/test_neb.pytests/test_stability.pytests/test_utils.py
💤 Files with no reviewable changes (6)
- scripts/install-perlmutter.sh
- mlip_arena/tasks/stability/flow.py
- mlip_arena/tasks/stability/init.py
- mlip_arena/tasks/stability/data.py
- mlip_arena/tasks/stability/analysis.py
- benchmarks/stability/run.py
| - name: Checkout code | ||
| uses: actions/checkout@v4 |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Verify workflow action pinning and checkout credential persistence settings.
# Expected:
# 1) Every "uses:" line should end with a full commit SHA.
# 2) Each actions/checkout step should explicitly set persist-credentials: false.
echo "== uses entries =="
rg -n '^\s*uses:\s*' .github/workflows/ci.yaml
echo
echo "== checkout blocks (with nearby context) =="
rg -n -C3 'actions/checkout@' .github/workflows/ci.yaml
echo
echo "== persist-credentials settings =="
rg -n 'persist-credentials' .github/workflows/ci.yamlRepository: atomind-ai/mlip-arena
Length of output: 773
Pin workflow actions to commit SHAs and add persist-credentials: false to checkout steps.
Multiple actions use mutable version tags (@v4, @v6, @v5, @v1.0) instead of immutable commit SHAs, creating supply-chain risk. Additionally, checkout steps lack persist-credentials: false, leaving GitHub tokens available to all subsequent steps. Audit and pin all uses: entries to full commit SHAs and set persist-credentials: false for checkout steps.
🧰 Tools
🪛 zizmor (1.25.2)
[warning] 16-17: credential persistence through GitHub Actions artifacts (artipacked): does not set persist-credentials: false
(artipacked)
[error] 17-17: unpinned action reference (unpinned-uses): action is not pinned to a hash (required by blanket policy)
(unpinned-uses)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In @.github/workflows/ci.yaml around lines 16 - 17, Replace the mutable version
tag `@v4` in the actions/checkout step with a full commit SHA to ensure
supply-chain security and consistency. Additionally, add persist-credentials:
false to the checkout step to prevent GitHub tokens from being passed to
subsequent workflow steps. After fixing the checkout step, audit all other
action uses entries in the workflow file (including any `@v6`, `@v5`, `@v1.0`
references) and replace their mutable version tags with specific commit SHAs as
well.
Source: Linters/SAST tools
| def homonuclear_diatomic( | ||
| symbol: str, calculator: str | MLIPEnum | BaseCalculator, calculator_kwargs: dict | None, out_dir: Path | ||
| ): |
There was a problem hiding this comment.
Make calculator_kwargs optional in both public signatures.
Line 27 and Line 260 currently require calculator_kwargs, which introduces a breaking API change and contradicts the optional-kwargs contract used elsewhere in this PR.
Proposed fix
def homonuclear_diatomic(
- symbol: str, calculator: str | MLIPEnum | BaseCalculator, calculator_kwargs: dict | None, out_dir: Path
+ symbol: str,
+ calculator: str | MLIPEnum | BaseCalculator,
+ out_dir: Path,
+ calculator_kwargs: dict | None = None,
):
@@
def homonuclear_diatomics(
- calculator: BaseCalculator | str, calculator_kwargs: dict | None, run_dir: Path | None = None
+ calculator: BaseCalculator | str,
+ run_dir: Path | None = None,
+ calculator_kwargs: dict | None = None,
):Also applies to: 259-261
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@mlip_arena/flows/diatomics.py` around lines 26 - 28, The `calculator_kwargs`
parameter in both the `homonuclear_diatomic` function (at lines 26-28) and the
corresponding function at lines 259-261 should be optional with a default value
of None. Change the parameter signature from `calculator_kwargs: dict | None` to
`calculator_kwargs: dict | None = None` in both function definitions to make the
parameter truly optional and maintain consistency with the optional-kwargs
contract used elsewhere in the codebase.
| pd.DataFrame: A DataFrame containing the raw EOS results. | ||
| """ | ||
| calculator = model if isinstance(model, BaseCalculator) else get_calculator(model) | ||
| calculator = calculator if isinstance(calculator, BaseCalculator) else get_calculator(calculator, calculator_kwargs) |
There was a problem hiding this comment.
Avoid silently ignoring calculator_kwargs for calculator instances.
Line 83 drops calculator_kwargs when calculator is already a BaseCalculator instance. With the new parameter exposed, this should fail fast to prevent misleading runs.
Proposed fix
- calculator = calculator if isinstance(calculator, BaseCalculator) else get_calculator(calculator, calculator_kwargs)
+ if isinstance(calculator, BaseCalculator):
+ if calculator_kwargs:
+ raise ValueError("`calculator_kwargs` cannot be used with a calculator instance.")
+ else:
+ calculator = get_calculator(calculator, calculator_kwargs)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@mlip_arena/flows/eos_bulk.py` at line 83, The line 83 in the eos_bulk.py file
silently ignores calculator_kwargs when calculator is already a BaseCalculator
instance, which can lead to misleading behavior. Add validation to check if
calculator is a BaseCalculator instance and calculator_kwargs is provided; if
both conditions are true, raise an error to fail fast and alert the user that
they cannot pass calculator_kwargs with an already-instantiated BaseCalculator.
If calculator is a BaseCalculator instance without calculator_kwargs, use it
as-is; otherwise, call get_calculator with the provided kwargs to instantiate a
new calculator.
| if not db_path.exists(): | ||
| db_path = hf_hub_download( | ||
| repo_id=repo_id, | ||
| repo_type=repo_type, | ||
| subfolder=subfolder, | ||
| # local_dir=db_path.parent, | ||
| filename=db_path.name, | ||
| token=hf_token, | ||
| force_download=force_download, | ||
| ) |
There was a problem hiding this comment.
force_download is currently ignored when the local DB file exists.
On Line 196, the download branch only runs when the file is missing, so callers passing force_download=True still read stale local data.
Proposed fix
- if not db_path.exists():
+ if force_download or not db_path.exists():
db_path = hf_hub_download(
repo_id=repo_id,
repo_type=repo_type,
subfolder=subfolder,
# local_dir=db_path.parent,
filename=db_path.name,
token=hf_token,
force_download=force_download,
)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if not db_path.exists(): | |
| db_path = hf_hub_download( | |
| repo_id=repo_id, | |
| repo_type=repo_type, | |
| subfolder=subfolder, | |
| # local_dir=db_path.parent, | |
| filename=db_path.name, | |
| token=hf_token, | |
| force_download=force_download, | |
| ) | |
| if force_download or not db_path.exists(): | |
| db_path = hf_hub_download( | |
| repo_id=repo_id, | |
| repo_type=repo_type, | |
| subfolder=subfolder, | |
| # local_dir=db_path.parent, | |
| filename=db_path.name, | |
| token=hf_token, | |
| force_download=force_download, | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@mlip_arena/flows/stability.py` around lines 196 - 205, The download logic in
the stability.py file is conditioned only on whether the database file exists
locally, completely ignoring the force_download parameter. Modify the condition
that guards the hf_hub_download call to also check the force_download flag, so
that the download executes when either the file does not exist OR when
force_download is True. This ensures that callers can force a fresh download of
the database by passing force_download=True, even if a local copy already
exists.
| total_time_seconds += block_time.total_seconds() | ||
| total_steps += steps[indices][-1] - steps[indices][0] | ||
|
|
||
| target_steps = traj[1].info["target_steps"] |
There was a problem hiding this comment.
target_steps indexing is off-by-one and can crash for short trajectories.
Line 295 reads traj[1]; a valid single-frame trajectory raises IndexError, and this silently drops the file in gather_results.
Proposed fix
- target_steps = traj[1].info["target_steps"]
+ target_steps = traj[0].info.get("target_steps", steps[-1] if len(steps) else 0)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| target_steps = traj[1].info["target_steps"] | |
| target_steps = traj[0].info.get("target_steps", steps[-1] if len(steps) else 0) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@mlip_arena/flows/stability.py` at line 295, The code accessing target_steps
uses off-by-one indexing with traj[1], which causes an IndexError when the
trajectory has only a single frame (only index 0 exists), resulting in the file
being silently dropped during processing. Change the indexing from traj[1] to
traj[0] in the line that reads target_steps = traj[1].info["target_steps"] to
correctly access the first frame of the trajectory instead of attempting to
access a second frame that may not exist.
| def __new__(cls, device: str | None = None, **kwargs): | ||
| device = device or str(get_freer_device()) | ||
| model = torchani.models.ANI2x().to(device) | ||
| return model.ase() |
There was a problem hiding this comment.
🧩 Analysis chain
🌐 Web query:
For torchani==2.2.4, what keyword arguments are supported by ANI2x().ase(), and are they intended to be passed through from wrapper constructors?
💡 Result:
In torchani 2.2.4 (consistent with the 2.2.x series), the ase method on an ANI model instance (like ANI2x) serves as a convenience wrapper to instantiate the torchani.ase.Calculator class [1][2]. The keyword arguments supported by the ase method are passed directly to the constructor of the underlying torchani.ase.Calculator class [1][3]. While specific signatures can evolve across versions, common keyword arguments include: 1. overwrite (bool): Determines whether to replace original positions in the ASE Atoms object with wrapped positions after periodic boundary condition handling (defaults to False) [1][3]. 2. stress_kind (str): Defines the strategy for stress calculation; valid options typically include 'scaling', 'fdotr', and 'numerical' (defaulting to 'scaling') [1][3]. Additionally, users may pass other arguments (such as dtype) that the implementation may handle during the calculator setup process [4]. Because the ase method acts as a factory for the calculator, these arguments are explicitly intended to be passed through to configure the resulting ASE calculator instance [1][3]. Always check the documentation specific to your local version by inspecting the docstring of the method directly if in doubt, as internal argument handling can vary slightly between minor versions [1][2].
Citations:
- 1: https://aiqm.github.io/torchani/_modules/torchani/ase.html
- 2: https://aiqm.github.io/torchani-test-docs/_modules/torchani/ase.html
- 3: https://aiqm.github.io/torchani/api_autogen/torchani.ase.html
- 4: https://aiqm.github.io/torchani-test-docs/examples/ase_interface.html
🏁 Script executed:
# First, let's locate and examine the actual file
fd -t f "ani.py" | head -20Repository: atomind-ai/mlip-arena
Length of output: 100
🏁 Script executed:
# Also check the structure to understand the file better
cat -n mlip_arena/models/externals/ani.py | head -50Repository: atomind-ai/mlip-arena
Length of output: 450
🏁 Script executed:
# Get the full context of the file
cat mlip_arena/models/externals/ani.pyRepository: atomind-ai/mlip-arena
Length of output: 373
🏁 Script executed:
# Check for similar wrapper patterns in the externals directory
ls -la mlip_arena/models/externals/Repository: atomind-ai/mlip-arena
Length of output: 1181
🏁 Script executed:
# Look at other model files to understand the pattern
head -30 mlip_arena/models/externals/*.pyRepository: atomind-ai/mlip-arena
Length of output: 13121
Forward kwargs to model.ase().
The constructor accepts **kwargs but silently discards them at line 11. The ase() method explicitly supports keyword arguments like overwrite and stress_kind that are intended to be configurable by callers. Update the return statement to forward these arguments.
Proposed fix
class ANI2x:
def __new__(cls, device: str | None = None, **kwargs):
device = device or str(get_freer_device())
model = torchani.models.ANI2x().to(device)
- return model.ase()
+ return model.ase(**kwargs)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def __new__(cls, device: str | None = None, **kwargs): | |
| device = device or str(get_freer_device()) | |
| model = torchani.models.ANI2x().to(device) | |
| return model.ase() | |
| def __new__(cls, device: str | None = None, **kwargs): | |
| device = device or str(get_freer_device()) | |
| model = torchani.models.ANI2x().to(device) | |
| return model.ase(**kwargs) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@mlip_arena/models/externals/ani.py` around lines 8 - 11, The __new__ method
accepts **kwargs but does not forward them to the model.ase() call on the return
statement. Update the return statement to pass **kwargs to the ase() method so
that callers can configure parameters like overwrite and stress_kind as
intended.
| nvt: true | ||
| npt: true | ||
| license: | ||
| deprecated: true |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
rg -n -C3 'DeepMD:|deprecated:\s*true' mlip_arena/models/registry.yaml
rg -n -C5 'meta\.get\("deprecated",\s*False\)|Skipping deprecated models|members\[model_name\]' mlip_arena/models/__init__.pyRepository: atomind-ai/mlip-arena
Length of output: 1880
Add explanatory comment for DeepMD deprecation.
The deprecated: true flag at line 391 is intentional and correctly excludes DeepMD from model registration via the deprecation check in __init__.py (lines 63–69). However, unlike other deprecated entries (e.g., ORB: "Superseded by ORBv2"), DeepMD's deprecation lacks an explanatory comment. Add a comment documenting the reason for deprecation for consistency and maintainability.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@mlip_arena/models/registry.yaml` at line 391, The DeepMD entry in the model
registry has a deprecated flag set to true, but it lacks an explanatory comment
documenting the reason for deprecation, unlike other deprecated entries such as
ORB which include descriptive comments. Add a comment next to the deprecated:
true flag in the DeepMD entry explaining the reason for its deprecation to
maintain consistency with other deprecated model entries and improve
maintainability.
| PET-MAD: | ||
| module: externals | ||
| class: PET_MAD | ||
| family: pet | ||
| package: upet==0.1.1 | ||
| checkpoint: pet-mad-s-v1.0.2.ckpt |
There was a problem hiding this comment.
Align PET-MAD registry artifact metadata with runtime model version.
Line [446] declares checkpoint pet-mad-s-v1.0.2.ckpt, while mlip_arena/models/externals/pet.py Line [30] hardcodes version="1.5.0". This creates a provenance mismatch between reported artifact metadata and the actually loaded model version.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@mlip_arena/models/registry.yaml` around lines 441 - 446, The PET-MAD
checkpoint name in the registry.yaml file declares pet-mad-s-v1.0.2.ckpt while
the version hardcoded in mlip_arena/models/externals/pet.py at line 30 shows
version="1.5.0", creating a provenance mismatch. Update the checkpoint value in
the PET-MAD registry entry to align with the actual version string from the
pet.py file, ensuring the declared checkpoint artifact metadata matches the
runtime model version being loaded.
| TORCH=2.5 | ||
| CUDA=cu124 | ||
| uv pip install torch==${TORCH}.0 | ||
| uv pip install -e .[test,extra,nequip] --no-cache |
There was a problem hiding this comment.
Quote extras specifiers to prevent shell glob expansion.
Lines 7/12/17/22/27/34 pass .[...] unquoted; [] is a glob character class in Bash. Quote each extras argument so it is always passed literally to uv pip.
Proposed fix
- uv pip install -e .[test,extra,nequip] --no-cache
+ uv pip install -e ".[test,extra,nequip]" --no-cache
...
- uv pip install -e .[test,extra,sevennet] --no-cache
+ uv pip install -e ".[test,extra,sevennet]" --no-cache
...
- uv pip install -e .[test,extra,mace] --no-cache
+ uv pip install -e ".[test,extra,mace]" --no-cache
...
- uv pip install -e .[test,extra,fairchem] --no-cache
+ uv pip install -e ".[test,extra,fairchem]" --no-cache
...
- uv pip install -e .[test,extra,orb] --no-cache
+ uv pip install -e ".[test,extra,orb]" --no-cache
...
- uv pip install -e .[matgl] --no-cache
- uv pip install -e .[test,extra,all] --no-cache
+ uv pip install -e ".[matgl]" --no-cache
+ uv pip install -e ".[test,extra,all]" --no-cacheAlso applies to: 12-12, 17-17, 22-22, 27-27, 34-34
🧰 Tools
🪛 Shellcheck (0.11.0)
[info] 7-7: Ranges can only match single chars (mentioned due to duplicates).
(SC2102)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@scripts/install-linux.sh` at line 7, The extras specifiers in the `uv pip
install` commands are not quoted, which allows the shell to interpret the square
brackets `[]` as glob character classes. Quote each extras argument (the part
containing `.[test,extra,nequip]` and similar patterns) by wrapping it in single
or double quotes so that the brackets are passed literally to `uv pip` rather
than being expanded as shell glob patterns. This applies to all occurrences of
the `uv pip install` command where extras are specified with unquoted bracket
notation.
Source: Linters/SAST tools
| temp_dir = os.path.join(tempfile.gettempdir(), f"prefect_home_{worker_id}") | ||
| os.makedirs(temp_dir, exist_ok=True) | ||
|
|
||
| os.environ["PREFECT_HOME"] = temp_dir | ||
| os.environ["PREFECT_API_DATABASE_CONNECTION_URL"] = f"sqlite+aiosqlite:///{temp_dir}/prefect.db" |
There was a problem hiding this comment.
Use a unique temp directory per test run to avoid stale Prefect state leakage.
Line 8 builds a deterministic path per worker, so repeated runs can reuse old prefect.db content. Prefer tempfile.mkdtemp(...) (or a session-scoped fixture creating unique dirs) to isolate runs as well as workers.
Proposed fix
-worker_id = os.environ.get("PYTEST_XDIST_WORKER", "master")
-temp_dir = os.path.join(tempfile.gettempdir(), f"prefect_home_{worker_id}")
-os.makedirs(temp_dir, exist_ok=True)
+worker_id = os.environ.get("PYTEST_XDIST_WORKER", "master")
+temp_dir = tempfile.mkdtemp(prefix=f"prefect_home_{worker_id}_")📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| temp_dir = os.path.join(tempfile.gettempdir(), f"prefect_home_{worker_id}") | |
| os.makedirs(temp_dir, exist_ok=True) | |
| os.environ["PREFECT_HOME"] = temp_dir | |
| os.environ["PREFECT_API_DATABASE_CONNECTION_URL"] = f"sqlite+aiosqlite:///{temp_dir}/prefect.db" | |
| temp_dir = tempfile.mkdtemp(prefix=f"prefect_home_{worker_id}_") | |
| os.environ["PREFECT_HOME"] = temp_dir | |
| os.environ["PREFECT_API_DATABASE_CONNECTION_URL"] = f"sqlite+aiosqlite:///{temp_dir}/prefect.db" |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/conftest.py` around lines 8 - 12, Replace the deterministic temp_dir
path construction (currently using tempfile.gettempdir() with worker_id) with
tempfile.mkdtemp() to generate a unique directory for each test run, ensuring
test isolation and preventing stale Prefect state leakage from previous runs.
Update the temp_dir assignment to use mkdtemp(), which will automatically create
a unique temporary directory and eliminate the reuse of old prefect.db content
across test runs.
1 participant
Organize model dependency groups in pyproject.toml and split pytest groups.
This PR also includes PET models in the dependency and will close #78 first. We will include benchmark results in a separate PR.
Summary by CodeRabbit
New Features
Bug Fixes
Chores