Add linux/aarch64 build configuration

* Added matrix build with both aarch64 and x86_64 to github/actions
* Added dynamic platform wheel name generator script to build_python_package.sh

Author: markomitos

.bazelrc

@@ -33,10 +33,17 @@ build --define=allow_oversize_protos=true
 # explicit __init__.py files.
 # build --incompatible_default_to_explicit_init_py
 
+# TODO: find a way to automatically actvate os_arch specific flags
+# Configuration for Linux on AMD64 (x86_64)
 # Haswell processor and later optimizations. This covers most processors deployed
 # today, includin Colab CPU runtimes.
-build --copt=-march=haswell
-build --host_copt=-march=haswell
+build:linux_x86_64 --copt=-march=haswell
+build:linux_x86_64 --host_copt=-march=haswell
+
+# Configuration for Linux on ARM64 (aarch64)
+build:linux_aarch64 --copt=-march=armv8-a
+build:linux_aarch64 --host_copt=-march=armv8-a
+
 # Only use level three optimizations for target, not necessarily for host
 # since host artifacts don't need to be fast.
 build --copt=-O3

.devcontainer/devcontainer.json

@@ -0,0 +1,34 @@
+{
+  "name": "TensorFlow Federated (DevContainer)",
+  image: "mcr.microsoft.com/devcontainers/base:ubuntu-22.04",
+  "features": {
+    "ghcr.io/devcontainers-community/features/bazel": {
+      "bazelisk_version": "latest"
+    },
+    "ghcr.io/devcontainers/features/python" : {
+      "version": "3.10",
+      "toolsToInstall": "black,pytest,pylint,isort,twine"
+    },
+    "ghcr.io/devcontainers-community/features/llvm": {
+      "version": "17"
+    },
+    "ghcr.io/devcontainers/features/java": {
+      "jdkDistro": "open"
+    }
+  },
+  "mounts": [
+    { "target": "/root/.cache/bazel", "source": "jetbrains-federated-devcontainer-cache-bazel", "type": "volume" },
+    { "target": "/root/.cache/pip/", "source": "jetbrains-federated-devcontainer-cache-pip", "type": "volume" },
+  ],
+  "customizations": {
+    "jetbrains": {
+      "backend": "PyCharm",
+      "plugins": [
+        "com.google.idea.bazel.ijwb"
+      ]
+    }
+  },
+  "containerUser": "root",
+  "postCreateCommand": ".devcontainer/post-create.sh",
+  "waitFor": "postStartCommand",
+}

.devcontainer/post-create.sh

@@ -0,0 +1,14 @@
+#!/usr/bin/env bash
+set -e
+exec 2>&1
+
+PROJECT_DIR=$(pwd)
+
+echo "Installing required python packages..."
+
+if nvidia-smi &> /dev/null; then
+  GPU_EXTRA="tensorflow[and-cuda]"
+fi
+
+# shellcheck disable=SC2086
+python3 -m pip install --root-user-action=ignore -r $PROJECT_DIR/requirements.txt ${GPU_EXTRA}

.github/workflows/publish.yaml

@@ -37,10 +37,10 @@ jobs:
     # Only if:
     # * Repository is not a fork.
     # * Branch is `main` (for workflow_dispatch trigger).
-    if: |
-      github.repository == 'google-parfait/tensorflow-federated'
-      && github.ref == 'refs/heads/main'
-    runs-on: ubuntu-latest
+#    if: |
+#      github.repository == 'google-parfait/tensorflow-federated'
+#      && github.ref == 'refs/heads/main'
+    runs-on: ubuntu-22.04
     timeout-minutes: 5
     permissions:
       contents: write  # Required to create a release.
@@ -97,8 +97,16 @@ jobs:
   build-package:
     name: Build Package
     needs: [publish-release]
-    runs-on: ubuntu-20.04
-    timeout-minutes: 60
+    strategy:
+      matrix:
+        include:
+          - os: ubuntu-22.04
+            arch: x86_64
+          - os: ubuntu-22.04-arm
+            arch: aarch64
+    runs-on: ${{ matrix.os }}
+    timeout-minutes: 360
+    environment: release
     steps:
 
     - name: Checkout repository
@@ -114,11 +122,11 @@ jobs:
         create_credentials_file: true
 
     - name: Set up bazel repository cache
-      uses: actions/cache@v4.0.2
+      uses: actions/cache@v4.2.0
       with:
         path: "~/.cache/bazel/"
-        key: ${{ runner.os }}-bazel-${{ hashFiles('.bazelversion', '.bazelrc', 'WORKSPACE') }}
-        restore-keys: ${{ runner.os }}-bazel-
+        key: ${{ runner.os }}-${{ matrix.arch }}-bazel-${{ hashFiles('.bazelversion', '.bazelrc', 'WORKSPACE') }}
+        restore-keys: ${{ runner.os }}-${{ matrix.arch }}-bazel
 
     - name: Set up Python
       uses: actions/[email protected]
@@ -130,26 +138,39 @@ jobs:
       run: |
         pip install --upgrade "pip"
 
+    - name: Install LLVM and Clang
+      uses: KyleMayes/install-llvm-action@v2
+      with:
+        version: "14.0"
+
     - name: Build Python package
       run: |
         pip install --upgrade "numpy~=1.25"
         bazelisk run //tools/python_package:build_python_package \
             --build_tag_filters="-nokokoro,-nopresubmit,-requires-gpu-nvidia" \
             --google_credentials="${{ steps.auth.outputs.credentials_file_path }}" \
-            --remote_cache="https://storage.googleapis.com/tensorflow-federated-bazel-cache/${{ github.job }}" \
+            --remote_cache="https://storage.googleapis.com/${{ vars.BAZEL_CACHE_BUCKET }}/${{ github.job }}" \
+            --config="linux_${{ matrix.arch }}" \
             -- \
             --output_dir="${{ github.workspace }}/dist/"
 
     - name: Upload Python package
       uses: actions/[email protected]
       with:
-        name: python-package-distributions
+        name: python-package-distributions-${{ matrix.arch }}
         path: dist/*.whl
 
   test-package:
     name: Test Package
     needs: [build-package]
-    runs-on: ubuntu-20.04
+    strategy:
+      matrix:
+        include:
+          - os: ubuntu-22.04
+            arch: x86_64
+          - os: ubuntu-22.04-arm
+            arch: aarch64
+    runs-on: ${{ matrix.os }}
     timeout-minutes: 5
     steps:
 
@@ -161,7 +182,7 @@ jobs:
     - name: Download Python package
       uses: actions/[email protected]
       with:
-        name: python-package-distributions
+        name: python-package-distributions-${{ matrix.arch }}
         path: dist/
 
     - name: Set up Python
@@ -190,17 +211,24 @@ jobs:
   publish-package:
     name: Publish Package
     needs: [build-package, test-package]
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-22.04
     timeout-minutes: 5
     permissions:
       id-token: write  # Required for trusted publishing.
+    environment: release
     steps:
 
     - name: Download Python package
       uses: actions/[email protected]
       with:
-        name: python-package-distributions
+        pattern: python-package-distributions-*
+        merge-multiple: true
         path: dist/
 
     - name: Publish Python package
       uses: pypa/[email protected]
+      with:
+        password: ${{ secrets.PYTHON_REPOSITORY_TOKEN }}
+        repository-url: ${{ vars.PYTHON_REPOSITORY_URL }}
+        verbose: 'true'
+        user: ${{ secrets.PYTHON_REPOSITORY_USER }}

pyproject.toml

@@ -39,9 +39,8 @@ dependencies = [
   'dp-accounting==0.4.3',
   'google-vizier==0.1.11',
   'grpcio~=1.46',
-  'jaxlib==0.4.14',
-  'jax==0.4.14',
-  'ml_dtypes>=0.2.0,==0.2.*',
+  'jaxlib==0.4.18',
+  'jax==0.4.18',
   'numpy~=1.25',
   'portpicker~=1.6',
   'scipy~=1.9.3',

requirements.txt

@@ -32,9 +32,8 @@ dm-tree==0.1.8
 dp-accounting==0.4.3
 google-vizier==0.1.11
 grpcio~=1.46
-jaxlib==0.4.14
-jax==0.4.14
-ml_dtypes>=0.2.0,==0.2.*
+jaxlib==0.4.18
+jax==0.4.18
 numpy~=1.25
 portpicker~=1.6
 scipy~=1.9.3

tools/python_package/build_python_package.sh

@@ -49,13 +49,30 @@ main() {
   fi
 
   # Check the GLIBC version.
-  local expected_glibc="2.31"
-  if ! ldd --version | grep --quiet "${expected_glibc}"; then
-    echo "error: expected GLIBC version to be '${expected_glibc}', found:" 1>&2
-    ldd --version 1>&2
+  glibc_version=$(ldd --version 2>&1 | grep "GLIBC" | awk '{print $NF}')
+
+  # Error handling if GLIBC version couldn't be determined.
+  if [[ -z "$glibc_version" ]]; then
+    echo "error: Could not determine GLIBC version." 1>&2
     exit 1
   fi
 
+  echo "Detected GLIBC version: $glibc_version"
+
+  # Extract major and minor version numbers for manylinux tag.
+  IFS='.' read -r glibc_major glibc_minor <<< "$glibc_version"
+  manylinux_version="${glibc_major}_${glibc_minor}"
+
+  # Detect architecture.
+  arch=$(uname -m)
+  case "$arch" in
+    aarch64|x86_64) ;; # Supported architectures
+    *) echo "error: Unsupported architecture: $arch" >&2; exit 1 ;;
+  esac
+
+  plat_name="manylinux_${manylinux_version}_${arch}"
+
+
   # Create a temp directory.
   local temp_dir="$(mktemp --directory)"
   trap "rm -rf ${temp_dir}" EXIT
@@ -68,7 +85,10 @@ main() {
   pip --version
 
   # Build the Python package.
-  pip install --upgrade "build"
+  pip install --upgrade "build" "toml-cli"
+
+  # Update wheel platform
+  toml set --toml-path "pyproject.toml" "tool.distutils.bdist_wheel.plat-name" "$plat_name"
   pip freeze
   python -m build --outdir "${output_dir}"
 }