diff --git a/.github/workflows/build_wheel.yml b/.github/workflows/build_wheel.yml index 0173803cc0..c053766b35 100644 --- a/.github/workflows/build_wheel.yml +++ b/.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 diff --git a/backend/find_tensorflow.py b/backend/find_tensorflow.py index 74b50d40a0..d0998d0460 100644 --- a/backend/find_tensorflow.py +++ b/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'"], } diff --git a/doc/install/easy-install.md b/doc/install/easy-install.md index 0d9577613f..97bc0aa223 100644 --- a/doc/install/easy-install.md +++ b/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} diff --git a/pyproject.toml b/pyproject.toml index 66b1c724ba..46d88afc67 100644 --- a/pyproject.toml +++ b/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