automatically set pytorch requires

Signed-off-by: Jinzhe Zeng <[email protected]>

backend/dp_backend.py

@@ -7,6 +7,9 @@
 
 from scikit_build_core import build as _orig
 
+from .find_pytorch import (
+    find_pytorch,
+)
 from .find_tensorflow import (
     find_tensorflow,
 )
@@ -40,10 +43,18 @@ def __dir__() -> List[str]:
 def get_requires_for_build_wheel(
     config_settings: dict,
 ) -> List[str]:
-    return _orig.get_requires_for_build_wheel(config_settings) + find_tensorflow()[1]
+    return (
+        _orig.get_requires_for_build_wheel(config_settings)
+        + find_tensorflow()[1]
+        + find_pytorch()[1]
+    )
 
 
 def get_requires_for_build_editable(
     config_settings: dict,
 ) -> List[str]:
-    return _orig.get_requires_for_build_editable(config_settings) + find_tensorflow()[1]
+    return (
+        _orig.get_requires_for_build_editable(config_settings)
+        + find_tensorflow()[1]
+        + find_pytorch()[1]
+    )

backend/find_pytorch.py

@@ -17,12 +17,14 @@
     get_path,
 )
 from typing import (
+    List,
     Optional,
+    Tuple,
 )
 
 
 @lru_cache
-def find_pytorch() -> Optional[str]:
+def find_pytorch() -> Tuple[Optional[str], List[str]]:
     """Find PyTorch library.
 
     Tries to find PyTorch in the order of:
@@ -39,9 +41,12 @@ def find_pytorch() -> Optional[str]:
     -------
     str, optional
         PyTorch library path if found.
+    list of str
+        TensorFlow requirement if not found. Empty if found.
     """
     if os.environ.get("DP_ENABLE_PYTORCH", "1") == "0":
-        return None
+        return None, []
+    requires = []
     pt_spec = None
 
     if (pt_spec is None or not pt_spec) and os.environ.get("PYTORCH_ROOT") is not None:
@@ -73,4 +78,33 @@ def find_pytorch() -> Optional[str]:
         # IndexError if submodule_search_locations is an empty list
     except (AttributeError, TypeError, IndexError):
         pt_install_dir = None
-    return pt_install_dir
+        requires.extend(get_pt_requirement()["torch"])
+    return pt_install_dir, requires
+
+
+@lru_cache
+def get_pt_requirement(pt_version: str = "") -> dict:
+    """Get PyTorch requirement when PT is not installed.
+
+    If pt_version is not given and the environment variable `PYTORCH_VERSION` is set, use it as the requirement.
+
+    Parameters
+    ----------
+    pt_version : str, optional
+        PT version
+
+    Returns
+    -------
+    dict
+        PyTorch requirement.
+    """
+    if pt_version is None:
+        return {"torch": []}
+    if pt_version == "":
+        pt_version = os.environ.get("PYTORCH_VERSION", "")
+
+    return {
+        "torch": [
+            f"torch=={pt_version}" if pt_version != "" else "torch>=2a",
+        ],
+    }

backend/read_env.py

@@ -103,7 +103,7 @@ def get_argument_from_env() -> Tuple[str, list, list, dict, str]:
         tf_version = None
 
     if os.environ.get("DP_ENABLE_PYTORCH", "1") == "1":
-        pt_install_dir = find_pytorch()
+        pt_install_dir, _ = find_pytorch()
         if pt_install_dir is None:
             raise RuntimeError("Cannot find installed PyTorch.")
         cmake_args.extend(

pyproject.toml

@@ -126,9 +126,6 @@ cu12 = [
     "nvidia-cudnn-cu12<9",
     "nvidia-cuda-nvcc-cu12",
 ]
-torch = [
-    "torch>=2a",
-]
 
 [tool.deepmd_build_backend.scripts]
 dp = "deepmd.main:main"