Initial changes for moving to ML-Actions GHA workflows

.bazelversion

@@ -0,0 +1,2 @@
+7.7.0
+# TODO: b/390391579 - Upgrade once bazel 8+ works.

.github/workflows/build-reusable.yml

@@ -0,0 +1,88 @@
+# Copyright 2026 The TensorFlow GNN Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+name: Reusable Build
+on:
+  workflow_call:
+    inputs:
+
+      python_version:
+        description: 'Python version'
+        required: false
+        type: string
+        default: '3.11'
+      bazel_version:
+        description: 'Bazel version'
+        required: false
+        type: string
+        default: '7.4.1'
+      tf_version_spec:
+        description: 'TensorFlow version specifier'
+        required: false
+        type: string
+      keras_version_spec:
+        description: 'Keras version specifier'
+        required: false
+        type: string
+      test_tf_nightly:
+        description: 'Whether to use tf-nightly'
+        required: false
+        type: boolean
+        default: false
+      legacy_keras:
+        description: 'Value for TF_USE_LEGACY_KERAS'
+        required: false
+        type: string
+        default: '0'
+      tag_filters:
+        description: 'Bazel test tag filters'
+        required: false
+        type: string
+        default: ''
+
+permissions:
+  contents: read
+
+jobs:
+  build-and-test:
+    runs-on: linux-x86-n2-32
+    defaults:
+      run:
+        shell: bash
+    container:
+      image: us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest
+    timeout-minutes: 60
+    steps:
+      - name: Checkout
+        uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # ratchet:actions/checkout@v4
+        with:
+          ref: ${{ github.sha }}
+          persist-credentials: false
+
+      - name: Bazel Version Checkup
+        run: bazel version || sleep 5 && bazel version # work around potential flakiness on download
+        env:
+          USE_BAZEL_VERSION: ${{ inputs.bazel_version }}
+
+      - name: Build and test wheel
+        env:
+          PYTHON_VERSION: ${{ inputs.python_version }}
+          USE_BAZEL_VERSION: ${{ inputs.bazel_version }}
+          TF_VERSION: ${{ inputs.tf_version_spec }}
+          KERAS_VERSION_SPEC: ${{ inputs.keras_version_spec }}
+          TAG_FILTERS: ${{ inputs.tag_filters }}
+          TF_USE_LEGACY_KERAS: ${{ inputs.legacy_keras }}
+          TEST_TF_NIGHTLY: ${{ inputs.test_tf_nightly }}
+
+        run: ./ci/build-gha.sh

.github/workflows/ci-build.yml

@@ -0,0 +1,54 @@
+# Copyright 2026 The TensorFlow GNN Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+name: CI
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+    branches:
+      - main
+
+permissions:
+  contents: read
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
+  # Don't cancel in-progress jobs for main branches.
+  cancel-in-progress: ${{ github.ref != 'main' }}
+
+jobs:
+  build-and-test-oldest:
+    uses: ./.github/workflows/build-reusable.yml
+    with:
+      python_version: '3.11'
+      tf_version_spec: '2.12.*'
+      legacy_keras: '0'
+      tag_filters: ',-tf_at_least_2_13'
+
+  build-and-test-newest-stable:
+    uses: ./.github/workflows/build-reusable.yml
+    with:
+      python_version: '3.11'
+      tf_version_spec: '2.20.*'
+      legacy_keras: '1'
+
+  build-and-test-nightly:
+    uses: ./.github/workflows/build-reusable.yml
+    with:
+      python_version: '3.11'
+      legacy_keras: '1'
+      test_tf_nightly: true

BUILD

@@ -3,6 +3,7 @@ py_library(
     # This is a dummy rule used as a absl dependency in open-source.
     # We expect absl to already be installed on the system, e.g. via
     # `pip install absl`
+    # strict_deps = False,
     visibility = ["//visibility:public"],
     deps = [],
 )
@@ -12,6 +13,7 @@ py_library(
     # This is a dummy rule used as a absl dependency in open-source.
     # We expect absl to already be installed on the system, e.g. via
     # `pip install absl`
+    # strict_deps = False,
     visibility = ["//visibility:public"],
     deps = [],
 )
@@ -21,6 +23,7 @@ py_library(
     # This is a dummy rule used as a absl dependency in open-source.
     # We expect absl to already be installed on the system, e.g. via
     # `pip install absl`
+    # strict_deps = False,
     visibility = ["//visibility:public"],
     deps = [],
 )
@@ -30,6 +33,7 @@ py_library(
     # This is a dummy rule used as a numpy dependency in open-source.
     # We expect numpy to already be installed on the system, e.g. via
     # `pip install numpy`
+    # strict_deps = False,
     visibility = ["//visibility:public"],
     deps = [],
 )
@@ -39,6 +43,7 @@ py_library(
     # This is a dummy rule used as a tensorflow dependency in open-source.
     # We expect tensorflow to already be installed on the system, e.g. via
     # `pip install tensorflow`
+    # strict_deps = False,
     visibility = ["//visibility:public"],
     deps = [],
 )
@@ -48,6 +53,7 @@ py_library(
     # This is a dummy rule used as a ai-edge-litert dependency in open-source.
     # We expect ai-edge-litert to already be installed on the system, e.g. via
     # `pip install ai-edge-litert`
+    # strict_deps = False,
     visibility = ["//visibility:public"],
     deps = [],
 )
@@ -57,6 +63,7 @@ py_library(
     # This is a dummy rule used as a tensorflow dependency in open-source.
     # We expect tensorflow to already be installed on the system, e.g. via
     # `pip install tensorflow`
+    # strict_deps = False,
     visibility = ["//visibility:public"],
     deps = [],
 )
@@ -66,6 +73,7 @@ py_library(
     # This is a dummy rule used as a dependency on vizier.service.pyvizier in open-source.
     # We expect Vizier to already be installed on the system, e.g. via
     # `pip install google-vizier` and extra steps if needed (b/254806045).
+    # strict_deps = False,
     visibility = ["//visibility:public"],
     deps = [],
 )
@@ -75,6 +83,7 @@ py_library(
     # This is a dummy rule used as a mock dependency in open-source tests.
     # We expect mock to already be installed on the system, e.g. via
     # `pip install mock`
+    # strict_deps = False,
     visibility = ["//visibility:public"],
     deps = [],
 )

ci/build-gha.sh

@@ -0,0 +1,132 @@
+#!/bin/bash
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+set -e
+set -x
+
+# For pyenv python installation on ml-build image
+sudo apt-get update
+sudo apt-get install -y libbz2-dev liblzma-dev libncurses-dev libffi-dev libssl-dev libreadline-dev libsqlite3-dev zlib1g-dev
+
+PYENV_ROOT="$HOME/.pyenv"
+PYTHON_VERSION=${PYTHON_VERSION:-"3.11"}
+
+function force_tensorflow_version() {
+  if [[ -z "${TF_VERSION}" ]]; then
+    echo "TF_VERSION is not set. Not forcing tensorflow version."
+    return
+  fi
+
+  pip install tensorflow=="${TF_VERSION}" --progress-bar off --upgrade
+  if [[ "$TF_USE_LEGACY_KERAS" == 1 ]]; then
+    pip install tf-keras=="${TF_VERSION}" --progress-bar off --upgrade
+  fi
+}
+
+echo "Installing pyenv.."
+git clone https://github.com/pyenv/pyenv.git "$PYENV_ROOT"
+export PATH="$HOME/.local/bin:$PYENV_ROOT/bin:$PATH"
+eval "$(pyenv init --path)"
+
+echo "Python setup..."
+pyenv install -s "$PYTHON_VERSION"
+pyenv global "$PYTHON_VERSION"
+
+PIP_TEST_PREFIX=bazel_pip
+
+python -m venv build_venv
+source build_venv/bin/activate
+
+# Debug messages to indicate the python version
+python --version
+
+# update pip
+pip install --upgrade pip
+
+# Install build
+pip install build
+
+TEST_ROOT=$(pwd)/${PIP_TEST_PREFIX}
+rm -rf "$TEST_ROOT"
+mkdir -p "$TEST_ROOT"
+ln -s "$(pwd)"/tensorflow_gnn "$TEST_ROOT"/tensorflow_gnn
+
+# Print the OS version
+cat /etc/os-release
+
+# Prepend common tag filters to a defined env_var
+# For example, tests for TF 2.8 shouldn't run RNG-dependent tests
+# These tag filters are enforced to start with a comma for separation
+tag_filters="-no_oss,-oss_excluded${TAG_FILTERS}"
+
+# Check that `bazel` does version selection as expected.
+if [[ -n "${USE_BAZEL_VERSION}" && $(bazel --version) != *${USE_BAZEL_VERSION}* ]]; then
+  echo "Mismatch of configured and actual bazel version (see logged [[ command)"
+  exit 1
+fi
+
+bazel clean
+
+if [[ "$TEST_TF_NIGHTLY" == "true" ]]; then
+  pip install --group test-nightly --progress-bar off --upgrade
+else
+  force_tensorflow_version
+fi
+
+python3 -m build --wheel
+deactivate
+
+# Start the test environment.
+python3 -m venv test_venv
+source test_venv/bin/activate
+
+# Check the python version
+python --version
+
+# update pip
+pip install --upgrade pip
+
+if [[ "$TEST_TF_NIGHTLY" == "true" ]]; then
+  pip install dist/tensorflow_gnn-*.whl
+  pip uninstall -y tensorflow tf-keras ai-edge-litert
+  pip install --group test-nightly --progress-bar off --upgrade
+
+  # Check that tf-nightly is installed but tensorflow is not
+  # Also check that tf-keras-nightly is installed.
+  if [[ $(pip freeze | grep -q tf_nightly=; echo $?) -eq 0 && $(pip freeze | grep -q tensorflow=; echo $?) -eq 0 ]]; then
+    echo "Found tensorflow and tf_nightly in the environment."
+    exit 1
+  fi
+  if [[ $(pip freeze | grep -q tf_keras-nightly=; echo $?) -eq 0 && $(pip freeze | grep -q tf_keras=; echo $?) -eq 0 ]]; then
+    echo "Found tf_keras and tf_keras-nightly in the environment."
+    exit 1
+  fi
+
+else
+  force_tensorflow_version
+
+  if [[ "$TF_USE_LEGACY_KERAS" == 1 ]]; then
+    pip install --group test-tf216plus --progress-bar off --upgrade
+  else
+    pip install --group test-pre-tf216 --progress-bar off --upgrade
+  fi
+
+  pip install dist/tensorflow_gnn-*.whl
+fi
+
+echo "Final packages after all pip commands:"
+pip list
+
+bazel test --test_env=TF_USE_LEGACY_KERAS --build_tag_filters="${tag_filters}" --test_tag_filters="${tag_filters}" --test_output=errors --verbose_failures=true --build_tests_only --define=no_tfgnn_py_deps=true --keep_going --experimental_repo_remote_exec //bazel_pip/tensorflow_gnn/...