From 8a4ab03361b9046dc5d0fbcc700faca04915721e Mon Sep 17 00:00:00 2001
From: Quang Tran <quangtr@google.com>
Date: Thu, 22 Jan 2026 17:40:43 +0000
Subject: [PATCH 1/7] test build config for kokoro instances

---
 kokoro/github/ubuntu/cpu/presubmit.cfg | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/kokoro/github/ubuntu/cpu/presubmit.cfg b/kokoro/github/ubuntu/cpu/presubmit.cfg
index 37bda92b..b02705fb 100644
--- a/kokoro/github/ubuntu/cpu/presubmit.cfg
+++ b/kokoro/github/ubuntu/cpu/presubmit.cfg
@@ -1,5 +1,9 @@
 build_file: "gnn/kokoro/github/ubuntu/cpu/build.sh"
 
+container_properties {
+  docker_image: "us-central1-docker.pkg.dev/kokoro-container-bakery/kokoro/ubuntu/ubuntu2004/full:current"
+}
+
 env_vars: {
   key: "USE_BAZEL_VERSION"
   value: "7.4.1"  # TODO - b/390391579: Unpin once bazel 8 works.
@@ -9,9 +13,9 @@ env_vars: {
   value: "3.11"
 }
 
-action {
-  define_artifacts {
-    regex: "**/sponge_log.log"
-    regex: "**/sponge_log.xml"
-  }
-}
\ No newline at end of file
+# action {
+#   define_artifacts {
+#     regex: "**/sponge_log.log"
+#     regex: "**/sponge_log.xml"
+#   }
+# }

From 9b1ecac5b1d6ef8dcac17cc8d5f08009adc949c8 Mon Sep 17 00:00:00 2001
From: Quang Tran <quangtr@google.com>
Date: Mon, 26 Jan 2026 20:15:51 +0000
Subject: [PATCH 2/7] tensorflow_gnn: Add Docker property to build configs for
 Kokoro Instaces migration

---
 kokoro/github/ubuntu/cpu/continuous.cfg              |  4 ++++
 .../github/ubuntu/cpu/newest_stable/continuous.cfg   |  4 ++++
 kokoro/github/ubuntu/cpu/newest_stable/presubmit.cfg |  4 ++++
 kokoro/github/ubuntu/cpu/oldest/continuous.cfg       |  4 ++++
 kokoro/github/ubuntu/cpu/oldest/presubmit.cfg        |  4 ++++
 kokoro/github/ubuntu/cpu/presubmit.cfg               | 12 ++++++------
 6 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/kokoro/github/ubuntu/cpu/continuous.cfg b/kokoro/github/ubuntu/cpu/continuous.cfg
index 37bda92b..73a1d641 100644
--- a/kokoro/github/ubuntu/cpu/continuous.cfg
+++ b/kokoro/github/ubuntu/cpu/continuous.cfg
@@ -1,5 +1,9 @@
 build_file: "gnn/kokoro/github/ubuntu/cpu/build.sh"
 
+container_properties {
+  docker_image: "us-central1-docker.pkg.dev/kokoro-container-bakery/kokoro/ubuntu/ubuntu2004/full:current"
+}
+
 env_vars: {
   key: "USE_BAZEL_VERSION"
   value: "7.4.1"  # TODO - b/390391579: Unpin once bazel 8 works.
diff --git a/kokoro/github/ubuntu/cpu/newest_stable/continuous.cfg b/kokoro/github/ubuntu/cpu/newest_stable/continuous.cfg
index 890a2916..9ca485cc 100644
--- a/kokoro/github/ubuntu/cpu/newest_stable/continuous.cfg
+++ b/kokoro/github/ubuntu/cpu/newest_stable/continuous.cfg
@@ -1,5 +1,9 @@
 build_file: "gnn/kokoro/github/ubuntu/cpu/build_versioned.sh"
 
+container_properties {
+  docker_image: "us-central1-docker.pkg.dev/kokoro-container-bakery/kokoro/ubuntu/ubuntu2004/full:current"
+}
+
 env_vars: {
   key: "USE_BAZEL_VERSION"
   value: "7.4.1"  # TODO - b/390391579: Unpin once bazel 8 works.
diff --git a/kokoro/github/ubuntu/cpu/newest_stable/presubmit.cfg b/kokoro/github/ubuntu/cpu/newest_stable/presubmit.cfg
index 890a2916..9ca485cc 100644
--- a/kokoro/github/ubuntu/cpu/newest_stable/presubmit.cfg
+++ b/kokoro/github/ubuntu/cpu/newest_stable/presubmit.cfg
@@ -1,5 +1,9 @@
 build_file: "gnn/kokoro/github/ubuntu/cpu/build_versioned.sh"
 
+container_properties {
+  docker_image: "us-central1-docker.pkg.dev/kokoro-container-bakery/kokoro/ubuntu/ubuntu2004/full:current"
+}
+
 env_vars: {
   key: "USE_BAZEL_VERSION"
   value: "7.4.1"  # TODO - b/390391579: Unpin once bazel 8 works.
diff --git a/kokoro/github/ubuntu/cpu/oldest/continuous.cfg b/kokoro/github/ubuntu/cpu/oldest/continuous.cfg
index aa98234d..4896bfb7 100644
--- a/kokoro/github/ubuntu/cpu/oldest/continuous.cfg
+++ b/kokoro/github/ubuntu/cpu/oldest/continuous.cfg
@@ -1,5 +1,9 @@
 build_file: "gnn/kokoro/github/ubuntu/cpu/build_versioned.sh"
 
+container_properties {
+  docker_image: "us-central1-docker.pkg.dev/kokoro-container-bakery/kokoro/ubuntu/ubuntu2004/full:current"
+}
+
 env_vars: {
   key: "USE_BAZEL_VERSION"
   value: "7.4.1"  # TODO - b/390391579: Unpin once bazel 8 works.
diff --git a/kokoro/github/ubuntu/cpu/oldest/presubmit.cfg b/kokoro/github/ubuntu/cpu/oldest/presubmit.cfg
index aa98234d..4896bfb7 100644
--- a/kokoro/github/ubuntu/cpu/oldest/presubmit.cfg
+++ b/kokoro/github/ubuntu/cpu/oldest/presubmit.cfg
@@ -1,5 +1,9 @@
 build_file: "gnn/kokoro/github/ubuntu/cpu/build_versioned.sh"
 
+container_properties {
+  docker_image: "us-central1-docker.pkg.dev/kokoro-container-bakery/kokoro/ubuntu/ubuntu2004/full:current"
+}
+
 env_vars: {
   key: "USE_BAZEL_VERSION"
   value: "7.4.1"  # TODO - b/390391579: Unpin once bazel 8 works.
diff --git a/kokoro/github/ubuntu/cpu/presubmit.cfg b/kokoro/github/ubuntu/cpu/presubmit.cfg
index b02705fb..d96beb68 100644
--- a/kokoro/github/ubuntu/cpu/presubmit.cfg
+++ b/kokoro/github/ubuntu/cpu/presubmit.cfg
@@ -13,9 +13,9 @@ env_vars: {
   value: "3.11"
 }
 
-# action {
-#   define_artifacts {
-#     regex: "**/sponge_log.log"
-#     regex: "**/sponge_log.xml"
-#   }
-# }
+action {
+  define_artifacts {
+    regex: "**/sponge_log.log"
+    regex: "**/sponge_log.xml"
+  }
+}

From e4813a6d0c2db3341bfb351899a57eb716f4947c Mon Sep 17 00:00:00 2001
From: Quang Tran <quangtr@google.com>
Date: Tue, 27 Jan 2026 20:31:52 +0000
Subject: [PATCH 3/7] add more tests

---
 .../cpu/newest_stable/presubmit-test.cfg      | 29 ++++++++++++++++
 .../ubuntu/cpu/oldest/presubmit-test.cfg      | 33 +++++++++++++++++++
 2 files changed, 62 insertions(+)
 create mode 100644 kokoro/github/ubuntu/cpu/newest_stable/presubmit-test.cfg
 create mode 100644 kokoro/github/ubuntu/cpu/oldest/presubmit-test.cfg

diff --git a/kokoro/github/ubuntu/cpu/newest_stable/presubmit-test.cfg b/kokoro/github/ubuntu/cpu/newest_stable/presubmit-test.cfg
new file mode 100644
index 00000000..9ca485cc
--- /dev/null
+++ b/kokoro/github/ubuntu/cpu/newest_stable/presubmit-test.cfg
@@ -0,0 +1,29 @@
+build_file: "gnn/kokoro/github/ubuntu/cpu/build_versioned.sh"
+
+container_properties {
+  docker_image: "us-central1-docker.pkg.dev/kokoro-container-bakery/kokoro/ubuntu/ubuntu2004/full:current"
+}
+
+env_vars: {
+  key: "USE_BAZEL_VERSION"
+  value: "7.4.1"  # TODO - b/390391579: Unpin once bazel 8 works.
+}
+env_vars: {
+  key: "PYTHON_VERSION"
+  value: "3.11"
+}
+env_vars: {
+  key: "TF_VERSION"
+  value: "2.20.*"
+}
+env_vars: {
+  key: "TF_USE_LEGACY_KERAS"
+  value: "1"
+}
+
+action {
+  define_artifacts {
+    regex: "**/sponge_log.log"
+    regex: "**/sponge_log.xml"
+  }
+}
\ No newline at end of file
diff --git a/kokoro/github/ubuntu/cpu/oldest/presubmit-test.cfg b/kokoro/github/ubuntu/cpu/oldest/presubmit-test.cfg
new file mode 100644
index 00000000..4896bfb7
--- /dev/null
+++ b/kokoro/github/ubuntu/cpu/oldest/presubmit-test.cfg
@@ -0,0 +1,33 @@
+build_file: "gnn/kokoro/github/ubuntu/cpu/build_versioned.sh"
+
+container_properties {
+  docker_image: "us-central1-docker.pkg.dev/kokoro-container-bakery/kokoro/ubuntu/ubuntu2004/full:current"
+}
+
+env_vars: {
+  key: "USE_BAZEL_VERSION"
+  value: "7.4.1"  # TODO - b/390391579: Unpin once bazel 8 works.
+}
+env_vars: {
+  key: "PYTHON_VERSION"
+  value: "3.11"
+}
+env_vars: {
+  key: "TF_VERSION"
+  value: "2.12.*"
+}
+env_vars: {
+  key: "TF_USE_LEGACY_KERAS"
+  value: "0"
+}
+env_vars: {
+  key: "TAG_FILTERS"
+  value: ",-tf_at_least_2_13"
+}
+
+action {
+  define_artifacts {
+    regex: "**/sponge_log.log"
+    regex: "**/sponge_log.xml"
+  }
+}
\ No newline at end of file

From 0680ac418d130cc76cc20c3b93515d9efbedfea3 Mon Sep 17 00:00:00 2001
From: Quang Tran <quangtr@google.com>
Date: Tue, 27 Jan 2026 20:49:40 +0000
Subject: [PATCH 4/7] modify build script

---
 kokoro/github/ubuntu/cpu/build_versioned.sh | 29 ++++++++++++++++++++-
 1 file changed, 28 insertions(+), 1 deletion(-)

diff --git a/kokoro/github/ubuntu/cpu/build_versioned.sh b/kokoro/github/ubuntu/cpu/build_versioned.sh
index e2e6238c..7bd64b59 100644
--- a/kokoro/github/ubuntu/cpu/build_versioned.sh
+++ b/kokoro/github/ubuntu/cpu/build_versioned.sh
@@ -16,6 +16,26 @@
 set -e
 set -x
 
+# --- START Bazel Version Management ---
+USE_BAZEL_VERSION=${USE_BAZEL_VERSION:-"7.4.1"}
+BAZEL_INSTALL_DIR="/tmp/bazel/${USE_BAZEL_VERSION}"
+BAZEL_BIN="${BAZEL_INSTALL_DIR}/bin/bazel"
+
+if ! [[ -x "${BAZEL_BIN}" && $(${BAZEL_BIN} --version) == *${USE_BAZEL_VERSION}* ]]; then
+  echo "Bazel ${USE_BAZEL_VERSION} not found or version mismatch. Installing..."
+  mkdir -p "${BAZEL_INSTALL_DIR}/bin"
+  curl -fLo "${BAZEL_INSTALL_DIR}/bin/bazel" "https://releases.bazel.build/${USE_BAZEL_VERSION}/release/bazel-${USE_BAZEL_VERSION}-linux-x86_64"
+  chmod +x "${BAZEL_INSTALL_DIR}/bin/bazel"
+  echo "Bazel ${USE_BAZEL_VERSION} installed."
+else
+  echo "Using cached Bazel ${USE_BAZEL_VERSION}"
+fi
+
+export PATH="${BAZEL_INSTALL_DIR}/bin:${PATH}"
+echo "Bazel version check:"
+bazel --version
+# --- END Bazel Version Management ---
+
 PYENV_ROOT="/home/kbuilder/.pyenv"
 PYTHON_VERSION=${PYTHON_VERSION:-"3.11"}
 
@@ -45,7 +65,14 @@ pip install --upgrade pip
 TEST_ROOT=$(pwd)/${PIP_TEST_PREFIX}
 rm -rf "$TEST_ROOT"
 mkdir -p "$TEST_ROOT"
-ln -s "$(pwd)"/tensorflow_gnn "$TEST_ROOT"/tensorflow_gnn
+
+# --- START Fixed Symlink ---
+# Create a relative symlink.
+# From $(pwd)/bazel_pip, we need to go up two levels to get to the parent of tensorflow_gnn
+ln -s ../../tensorflow_gnn "${TEST_ROOT}"/tensorflow_gnn
+echo "Created symlink:"
+ls -l "${TEST_ROOT}"/tensorflow_gnn
+# --- END Fixed Symlink ---
 
 # Print the OS version
 cat /etc/os-release

From cecf1327409a218c6b2a2244d232378bd466896c Mon Sep 17 00:00:00 2001
From: Quang Tran <quangtr@google.com>
Date: Tue, 27 Jan 2026 21:16:24 +0000
Subject: [PATCH 5/7] 2nd mod

---
 kokoro/github/ubuntu/cpu/build_versioned.sh | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/kokoro/github/ubuntu/cpu/build_versioned.sh b/kokoro/github/ubuntu/cpu/build_versioned.sh
index 7bd64b59..77c84871 100644
--- a/kokoro/github/ubuntu/cpu/build_versioned.sh
+++ b/kokoro/github/ubuntu/cpu/build_versioned.sh
@@ -68,8 +68,8 @@ mkdir -p "$TEST_ROOT"
 
 # --- START Fixed Symlink ---
 # Create a relative symlink.
-# From $(pwd)/bazel_pip, we need to go up two levels to get to the parent of tensorflow_gnn
-ln -s ../../tensorflow_gnn "${TEST_ROOT}"/tensorflow_gnn
+# From $(pwd)/bazel_pip, tensorflow_gnn is one level up.
+ln -s ../tensorflow_gnn "${TEST_ROOT}"/tensorflow_gnn
 echo "Created symlink:"
 ls -l "${TEST_ROOT}"/tensorflow_gnn
 # --- END Fixed Symlink ---
@@ -102,3 +102,8 @@ 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/...
+
+# --- START Remove Bazel Symlinks ---
+echo "Removing Bazel-generated symlinks..."
+rm -f bazel-bin bazel-out bazel-genfiles bazel-testlogs bazel-gnn
+# --- END Remove Bazel Symlinks ---

From 70229cc00ec0fb1e661b9f507fd05134a6dfb39f Mon Sep 17 00:00:00 2001
From: Quang Tran <quangtr@google.com>
Date: Tue, 27 Jan 2026 21:53:48 +0000
Subject: [PATCH 6/7] 3rd fix

---
 kokoro/github/ubuntu/cpu/build_versioned.sh | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/kokoro/github/ubuntu/cpu/build_versioned.sh b/kokoro/github/ubuntu/cpu/build_versioned.sh
index 77c84871..c9993e1d 100644
--- a/kokoro/github/ubuntu/cpu/build_versioned.sh
+++ b/kokoro/github/ubuntu/cpu/build_versioned.sh
@@ -51,9 +51,10 @@ pyenv global "$PYTHON_VERSION"
 cd "${KOKORO_ARTIFACTS_DIR}/github/gnn/"
 
 PIP_TEST_PREFIX=bazel_pip
+VENV_DIR="venv"
 
-python -m venv venv
-source venv/bin/activate
+python -m venv "${VENV_DIR}"
+source "${VENV_DIR}/bin/activate"
 
 # Debug messages to indicate the python version
 python --version
@@ -101,9 +102,16 @@ pip install dist/tensorflow_gnn-*.whl
 echo "Final packages after all pip commands:"
 pip list
 
+# Deactivate venv before running tests to avoid issues with test environments
+deactivate
+
+# Run tests
 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/...
 
-# --- START Remove Bazel Symlinks ---
+# --- START Cleanup Symlinks and venv ---
 echo "Removing Bazel-generated symlinks..."
 rm -f bazel-bin bazel-out bazel-genfiles bazel-testlogs bazel-gnn
-# --- END Remove Bazel Symlinks ---
+
+echo "Removing venv directory..."
+rm -rf "${VENV_DIR}"
+# --- END Cleanup Symlinks and venv ---

From f6acb693075062e4becd827a9b050be0386e32a0 Mon Sep 17 00:00:00 2001
From: Quang Tran <quangtr@google.com>
Date: Tue, 27 Jan 2026 22:41:16 +0000
Subject: [PATCH 7/7] 4th fix

---
 kokoro/github/ubuntu/cpu/build_versioned.sh   | 48 ++-----------------
 .../cpu/newest_stable/presubmit-test.cfg      |  2 +-
 2 files changed, 5 insertions(+), 45 deletions(-)

diff --git a/kokoro/github/ubuntu/cpu/build_versioned.sh b/kokoro/github/ubuntu/cpu/build_versioned.sh
index c9993e1d..378fa3d0 100644
--- a/kokoro/github/ubuntu/cpu/build_versioned.sh
+++ b/kokoro/github/ubuntu/cpu/build_versioned.sh
@@ -16,26 +16,6 @@
 set -e
 set -x
 
-# --- START Bazel Version Management ---
-USE_BAZEL_VERSION=${USE_BAZEL_VERSION:-"7.4.1"}
-BAZEL_INSTALL_DIR="/tmp/bazel/${USE_BAZEL_VERSION}"
-BAZEL_BIN="${BAZEL_INSTALL_DIR}/bin/bazel"
-
-if ! [[ -x "${BAZEL_BIN}" && $(${BAZEL_BIN} --version) == *${USE_BAZEL_VERSION}* ]]; then
-  echo "Bazel ${USE_BAZEL_VERSION} not found or version mismatch. Installing..."
-  mkdir -p "${BAZEL_INSTALL_DIR}/bin"
-  curl -fLo "${BAZEL_INSTALL_DIR}/bin/bazel" "https://releases.bazel.build/${USE_BAZEL_VERSION}/release/bazel-${USE_BAZEL_VERSION}-linux-x86_64"
-  chmod +x "${BAZEL_INSTALL_DIR}/bin/bazel"
-  echo "Bazel ${USE_BAZEL_VERSION} installed."
-else
-  echo "Using cached Bazel ${USE_BAZEL_VERSION}"
-fi
-
-export PATH="${BAZEL_INSTALL_DIR}/bin:${PATH}"
-echo "Bazel version check:"
-bazel --version
-# --- END Bazel Version Management ---
-
 PYENV_ROOT="/home/kbuilder/.pyenv"
 PYTHON_VERSION=${PYTHON_VERSION:-"3.11"}
 
@@ -51,10 +31,9 @@ pyenv global "$PYTHON_VERSION"
 cd "${KOKORO_ARTIFACTS_DIR}/github/gnn/"
 
 PIP_TEST_PREFIX=bazel_pip
-VENV_DIR="venv"
 
-python -m venv "${VENV_DIR}"
-source "${VENV_DIR}/bin/activate"
+python -m venv venv
+source venv/bin/activate
 
 # Debug messages to indicate the python version
 python --version
@@ -66,14 +45,7 @@ pip install --upgrade pip
 TEST_ROOT=$(pwd)/${PIP_TEST_PREFIX}
 rm -rf "$TEST_ROOT"
 mkdir -p "$TEST_ROOT"
-
-# --- START Fixed Symlink ---
-# Create a relative symlink.
-# From $(pwd)/bazel_pip, tensorflow_gnn is one level up.
-ln -s ../tensorflow_gnn "${TEST_ROOT}"/tensorflow_gnn
-echo "Created symlink:"
-ls -l "${TEST_ROOT}"/tensorflow_gnn
-# --- END Fixed Symlink ---
+ln -s "$(pwd)"/tensorflow_gnn "$TEST_ROOT"/tensorflow_gnn
 
 # Print the OS version
 cat /etc/os-release
@@ -102,16 +74,4 @@ pip install dist/tensorflow_gnn-*.whl
 echo "Final packages after all pip commands:"
 pip list
 
-# Deactivate venv before running tests to avoid issues with test environments
-deactivate
-
-# Run tests
-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/...
-
-# --- START Cleanup Symlinks and venv ---
-echo "Removing Bazel-generated symlinks..."
-rm -f bazel-bin bazel-out bazel-genfiles bazel-testlogs bazel-gnn
-
-echo "Removing venv directory..."
-rm -rf "${VENV_DIR}"
-# --- END Cleanup Symlinks and venv ---
+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/...
\ No newline at end of file
diff --git a/kokoro/github/ubuntu/cpu/newest_stable/presubmit-test.cfg b/kokoro/github/ubuntu/cpu/newest_stable/presubmit-test.cfg
index 9ca485cc..4d611591 100644
--- a/kokoro/github/ubuntu/cpu/newest_stable/presubmit-test.cfg
+++ b/kokoro/github/ubuntu/cpu/newest_stable/presubmit-test.cfg
@@ -6,7 +6,7 @@ container_properties {
 
 env_vars: {
   key: "USE_BAZEL_VERSION"
-  value: "7.4.1"  # TODO - b/390391579: Unpin once bazel 8 works.
+  value: "5.2.0"  # TODO - b/390391579: Unpin once bazel 8 works.
 }
 env_vars: {
   key: "PYTHON_VERSION"