diff --git a/.github/workflows/test_python.yml b/.github/workflows/test_python.yml index 3c79b31be5..1ef17d0f67 100644 --- 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 index 9bfc472f1c..aa75d5ecb4 100644 --- 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, }