install mpi

njzjz · Sep 12, 2023 · 9f0aac1 · 9f0aac1
1 parent fa58e69
commit 9f0aac1
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 1 deletion.
diff --git a/.github/workflows/test_python.yml b/.github/workflows/test_python.yml
@@ -22,12 +22,15 @@ jobs:
       with:
         python-version: ${{ matrix.python }}
         cache: 'pip'
+    - uses: mpi4py/setup-mpi@v1
     # https://github.com/pypa/pip/issues/11770
     - run: python -m pip install -U "pip>=21.3.1,!=23.0.0"
-    - run: pip install -e .[cpu,test]
+    - run: pip install -e .[cpu,test,mpi]
       env:
         TENSORFLOW_VERSION: ${{ matrix.tf }}
         DP_BUILD_TESTING: 1
+        HOROVOD_WITH_TENSORFLOW: 1
+        HOROVOD_WITHOUT_GLOO: 1
     - run: dp --version
     - run: pytest --cov=deepmd --cov=deepmd_cli source/tests --durations=0
     - uses: codecov/codecov-action@v3

diff --git a/backend/find_tensorflow.py b/backend/find_tensorflow.py
@@ -113,8 +113,16 @@ def get_tf_requirement(tf_version: str = "") -> dict:
         tf_version = os.environ.get("TENSORFLOW_VERSION", "")
 
     extra_requires = []
+    extra_select = {}
     if not (tf_version == "" or tf_version in SpecifierSet(">=2.12")):
         extra_requires.append("protobuf<3.20")
+    if tf_version == "" or tf_version in SpecifierSet(">=1.15"):
+        extra_select["mpi"] = [
+            "horovod",
+            "mpi4py",
+        ]
+    else:
+        extra_select["mpi"] = []
 
     if tf_version == "":
         return {
@@ -128,6 +136,7 @@ def get_tf_requirement(tf_version: str = "") -> dict:
                 "tensorflow-metal; platform_machine=='arm64' and platform_system == 'Darwin'",
                 *extra_requires,
             ],
+            **extra_select,
         }
     elif tf_version in SpecifierSet("<1.15") or tf_version in SpecifierSet(
         ">=2.0,<2.1"
@@ -142,6 +151,7 @@ def get_tf_requirement(tf_version: str = "") -> dict:
                 f"tensorflow=={tf_version}; platform_machine=='aarch64'",
                 *extra_requires,
             ],
+            **extra_select,
         }
     else:
         return {
@@ -155,6 +165,7 @@ def get_tf_requirement(tf_version: str = "") -> dict:
                 "tensorflow-metal; platform_machine=='arm64' and platform_system == 'Darwin'",
                 *extra_requires,
             ],
+            **extra_select,
         }