ci(wheel): build PT OPs

[njzjz]

Build PT OPs for our wheels (against the CPU version of PT but work with the GPU version). Add PT to the build dependencies if it is not found. Bump MPICH to the latest version. Update documentation to address that PT support for LAMMPS and i-PI is not included.

Summary by CodeRabbit

• New Features

  • Introduced PyTorch support across various backend functionalities, enhancing compatibility and capabilities alongside existing TensorFlow support.

• Documentation

  • Updated installation guide to include warnings about TensorFlow backend support and removed outdated notes about macOS arm64 package support.

• Chores

  • Updated dependency versions and build configurations to include new PyTorch settings and improved compatibility for macOS and Linux.
  • Adjusted build scripts to conditionally link libraries based on new PyTorch-related parameters.

[coderabbitai]

Walkthrough

Walkthrough

This update introduces PyTorch support to the backend, dynamically adjusting the build and runtime requirements based on the availability of PyTorch. Key functions now include checks for PyTorch, and new functionalities are added to handle PyTorch-specific requirements and versions. The documentation and configuration files have been updated to reflect these changes, ensuring smooth integration with existing PyTorch and TensorFlow dependencies.

Changes

File(s)	Change Summary
backend/dp_backend.py	Imported find_pytorch, and updated get_requires_for_build_wheel and get_requires_for_build_editable to include find_pytorch().
backend/dynamic_metadata.py	Imported get_pt_requirement, and modified dynamic_metadata function to handle PyTorch version.
backend/find_pytorch.py	Updated find_pytorch to return PyTorch path and requirements. Added get_pt_requirement and get_pt_version functions.
backend/find_tensorflow.py	Updated get_tf_version to accept an optional path parameter.
backend/read_env.py	Imported get_pt_version, updated get_argument_from_env and set_scikit_build_env for PyTorch version handling.
doc/install/easy-install.md	Removed macOS arm64 note and added warning about TensorFlow backend support with LAMMPS and i-PI when installing from pip.
pyproject.toml	Updated dependency versions, test configurations, and build settings. Included PyTorch in test extras and added PyTorch-related environment variables.
source/api_cc/CMakeLists.txt	Introduced a condition that checks for BUILD_PY_IF before linking libraries and setting compile definitions when ENABLE_PYTORCH is enabled.

---

Recent review details

Configuration used: CodeRabbit UI
Review profile: CHILL

Commits

Files that changed from the base of the PR and between 4c38fad and 896e995.

Files selected for processing (3)

• backend/find_pytorch.py (4 hunks)
• backend/find_tensorflow.py (1 hunks)
• pyproject.toml (2 hunks)

Files skipped from review as they are similar to previous changes (1)

• pyproject.toml

Additional context used

Ruff

backend/find_tensorflow.py

105-105: Within an except clause, raise exceptions with raise ... from err or raise ... from None to distinguish them from errors in exception handling (B904)

---

Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media?

Share

• X
• Mastodon
• Reddit
• LinkedIn

Tips

Chat

There are 3 ways to chat with CodeRabbit:

• Review comments: Directly reply to a review comment made by CodeRabbit. Example:
  • I pushed a fix in commit <commit_id>.
  • Generate unit testing code for this file.
  • Open a follow-up GitHub issue for this discussion.
• Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
  • @coderabbitai generate unit testing code for this file.
  • @coderabbitai modularize this function.
• PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
  • @coderabbitai generate interesting stats about this repository and render them as a table.
  • @coderabbitai show all the console.log statements in this repository.
  • @coderabbitai read src/utils.ts and generate unit testing code.
  • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.
  • @coderabbitai help me debug CodeRabbit configuration file.

Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments.

CodeRabbit Commands (invoked as PR comments)

• @coderabbitai pause to pause the reviews on a PR.
• @coderabbitai resume to resume the paused reviews.
• @coderabbitai review to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
• @coderabbitai full review to do a full review from scratch and review all the files again.
• @coderabbitai summary to regenerate the summary of the PR.
• @coderabbitai resolve resolve all the CodeRabbit review comments.
• @coderabbitai configuration to show the current CodeRabbit configuration for the repository.
• @coderabbitai help to get help.

Additionally, you can add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.

CodeRabbit Configration File (.coderabbit.yaml)

• You can programmatically configure CodeRabbit by adding a .coderabbit.yaml file to the root of your repository.
• Please see the configuration documentation for more information.
• If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: # yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json

Documentation and Community

• Visit our Documentation for detailed information on how to use CodeRabbit.
• Join our Discord Community to get help, request features, and share feedback.
• Follow us on X/Twitter for updates and announcements.

[coderabbitai]

Actionable comments posted: 4

Outside diff range and nitpick comments (6)

backend/find_pytorch.py (1)

Line range hint 33-88: Update function documentation to reflect actual return values.

The function find_pytorch documentation incorrectly mentions returning TensorFlow requirements. It should state that it returns PyTorch requirements when PyTorch is not found.

-        TensorFlow requirement if not found. Empty if found.
+        PyTorch requirement if not found. Empty if found.

backend/read_env.py (3)

18-18: Add a comment for clarity about the new import.

The import of get_pt_version is new and directly related to the integration of PyTorch. It would be beneficial to add a comment explaining its purpose and usage in this file.

---

Line range hint 56-56: Correct the casing of the environment variable for consistency.

The environment variable CUDAToolkit_ROOT should be CUDATOOLKIT_ROOT to maintain consistency and avoid potential issues on case-sensitive systems.

- CUDAToolkit_ROOT
+ CUDATOOLKIT_ROOT

---

137-137: Clarify the use of unpacking in function calls.

The unpacking in the function call to get_argument_from_env is not immediately clear and could benefit from a comment explaining why only certain values are used.

doc/install/easy-install.md (2)

Line range hint 40-40: Correct the verb form for clarity.

The verb "setup" should be "set up" when used as a verb. This improves the clarity and correctness of the documentation.

- for how to setup a conda environment.
+ for how to set up a conda environment.

---

Line range hint 129-129: Clarify the instruction for adding extras.

The use of "as" in the sentence about adding extras is confusing. Clarifying this could improve the readability of the documentation.

- add `lmp` and/or `ipi` to extras:
+ add `lmp` and/or `ipi` as extras:

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 82.74%. Comparing base (b6f0fa3) to head (896e995).
Report is 119 commits behind head on devel.

Additional details and impacted files

@@           Coverage Diff           @@
##            devel    #3894   +/-   ##
=======================================
  Coverage   82.74%   82.74%           
=======================================
  Files         519      519           
  Lines       50491    50491           
  Branches     3015     3015           
=======================================
  Hits        41781    41781           
  Misses       7773     7773           
  Partials      937      937           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[coderabbitai]

Actionable comments posted: 3

Outside diff range and nitpick comments (2)

backend/find_pytorch.py (1)

Line range hint 33-88: Refactor find_pytorch function for clarity and efficiency.

The find_pytorch function attempts to locate the PyTorch library and determine the necessary TensorFlow requirements if PyTorch is not found. While the logic is sound, the function could benefit from refactoring to improve readability and maintainability. Specifically:

• The use of multiple if conditions to check pt_spec can be streamlined.
• Exception handling in the function is broad and might suppress important errors unintentionally.

Consider restructuring the code to reduce duplication and improve error handling.

@lru_cache
def find_pytorch() -> Tuple[Optional[str], List[str]]:
    if os.environ.get("DP_ENABLE_PYTORCH", "0") == "0":
        return None, []
    pt_spec = find_spec("torch") or find_spec_in_custom_path()

    if not pt_spec and site.ENABLE_USER_SITE:
        pt_spec = find_spec_in_site_packages(site.getusersitepackages())

    if not pt_spec:
        pt_spec = find_spec_in_site_packages(get_path("purelib"))

    pt_install_dir = get_install_dir_from_spec(pt_spec)
    requires = get_pt_requirement()["torch"] if not pt_install_dir else []
    return pt_install_dir, requires

def find_spec_in_custom_path():
    site_packages = Path(os.environ.get("PYTORCH_ROOT", '')).parent.absolute()
    return FileFinder(str(site,packages)).find_spec("torch") if site_packages else None

def find_spec_in_site_packages(site_packages: str):
    return FileFinder(site_packages).find_spec("torch") if site_packages else None

def get_install_dir_from_spec(pt_spec):
    try:
        return pt_spec.submodule_search_locations[0]
    except (AttributeError, TypeError, IndexError):
        return None

backend/find_tensorflow.py (1)

Line range hint 105-105: Improve exception handling in TensorFlow version extraction.

The current implementation catches broad exceptions which might suppress important errors. Use raise ... from err to provide more context about the original exception.

    try:
        tf_install_dir = tf_spec.submodule_search_locations[0]
    except (AttributeError, TypeError, IndexError) as err:
        raise RuntimeError("Failed to locate TensorFlow installation directory") from err