build wheels for Linux ARM64

.github/workflows/build_wheel.yml

@@ -16,30 +16,42 @@ jobs:
           - os: ubuntu-latest
             python: 310
             platform_id: manylinux_x86_64
+            dp_variant: cuda
           # macos-x86-64
           - os: macos-latest
             python: 310
             platform_id: macosx_x86_64
+            dp_variant: cpu
           # win-64
           - os: windows-2019
             python: 310
             platform_id: win_amd64
+            dp_variant: cpu
+          # linux-aarch64
+          - os: ubuntu-latest
+            python: 310
+            platform_id: manylinux_aarch64
+            dp_variant: cpu
     steps:
       - uses: actions/checkout@v2
         with:
           submodules: true
+      - uses: docker/setup-qemu-action@v2
+        name: Setup QEMU
+        if: matrix.platform_id == 'manylinux_aarch64'
       - uses: actions/setup-python@v2
         name: Install Python
         with:
           python-version: '3.8'
 
-      - run: python -m pip install cibuildwheel==2.11.2
+      - run: python -m pip install cibuildwheel==2.11.3
       - name: Build wheels
         run: python -m cibuildwheel --output-dir wheelhouse
         env:
           CIBW_BUILD_VERBOSITY: 1
           CIBW_ARCHS: all
           CIBW_BUILD: cp${{ matrix.python }}-${{ matrix.platform_id }}
+          DP_VARIANT: ${{ matrix.dp_variant }}
       - uses: actions/upload-artifact@v2
         with:
           path: ./wheelhouse/*.whl

backend/find_tensorflow.py

@@ -88,18 +88,18 @@ def get_tf_requirement(tf_version: str = "") -> dict:
 
     if tf_version == "":
         return {
-            "cpu": ["tensorflow-cpu"],
-            "gpu": ["tensorflow"],
+            "cpu": ["tensorflow-cpu; platform_machine!='aarch64'", "tensorflow; platform_machine=='aarch64'"],
+            "gpu": ["tensorflow; platform_machine!='aarch64'", "tensorflow; platform_machine=='aarch64'"],
         }
     elif tf_version in SpecifierSet("<1.15") or tf_version in SpecifierSet(">=2.0,<2.1"):
         return {
-            "cpu": [f"tensorflow=={tf_version}"],
-            "gpu": [f"tensorflow-gpu=={tf_version}"],
+            "cpu": [f"tensorflow=={tf_version}; platform_machine!='aarch64'", f"tensorflow=={tf_version}; platform_machine=='aarch64'"],
+            "gpu": [f"tensorflow-gpu=={tf_version}; platform_machine!='aarch64'", f"tensorflow=={tf_version}; platform_machine=='aarch64'"],
         }
     else:
         return {
-            "cpu": [f"tensorflow-cpu=={tf_version}"],
-            "gpu": [f"tensorflow=={tf_version}"],
+            "cpu": [f"tensorflow-cpu=={tf_version}; platform_machine!='aarch64'", f"tensorflow=={tf_version}; platform_machine=='aarch64'"],
+            "gpu": [f"tensorflow=={tf_version}; platform_machine!='aarch64'", f"tensorflow=={tf_version}; platform_machine=='aarch64'"],
         }
 
 

doc/install/easy-install.md

@@ -95,7 +95,7 @@ Or install the CPU version without CUDA supported:
 pip install deepmd-kit[cpu]
 ```
 
-The supported platform includes Linux x86-64 with GNU C Library 2.28 or above, macOS x86-64, and Windows x86-64.
+The supported platform includes Linux x86-64 and aarch64 with GNU C Library 2.28 or above, macOS x86-64, and Windows x86-64.
 A specific version of TensorFlow which is compatible with DeePMD-kit will be also installed.
 
 :::{Warning}

pyproject.toml

@@ -55,15 +55,19 @@ build = ["cp310-*"]
 skip = ["*-win32", "*-manylinux_i686", "*-musllinux*"]
 # TODO: bump to "latest" tag when CUDA supports GCC 12
 manylinux-x86_64-image = "quay.io/pypa/manylinux_2_28_x86_64:2022-11-19-1b19e81"
+manylinux-aarch64-image = "quay.io/pypa/manylinux_2_28_aarch64:2022-11-19-1b19e81"
 
 [tool.cibuildwheel.macos]
 repair-wheel-command = "delocate-wheel --require-archs {delocate_archs} -w {dest_dir} -v {wheel} --ignore-missing-dependencies"
 
 [tool.cibuildwheel.linux]
 repair-wheel-command = "auditwheel repair --exclude libtensorflow_framework.so.2 --exclude libtensorflow_framework.so.1 --exclude libtensorflow_framework.so -w {dest_dir} {wheel}"
-environment-pass = ["CIBW_BUILD"]
-environment = { DP_VARIANT="cuda" }
-before-all = "yum config-manager --add-repo http://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/cuda-rhel8.repo && yum install -y cuda-11-8"
+environment-pass = ["CIBW_BUILD", "DP_VARIANT"]
+before-all = """
+if [ "$(uname -m)" = "x86_64" ]; then
+    yum config-manager --add-repo http://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/cuda-rhel8.repo && yum install -y cuda-11-8
+fi
+"""
 
 # selectively turn of lintner warnings, always include reasoning why any warning should
 # be silenced