diff --git a/kokoro/github/ubuntu/cpu/build.sh b/kokoro/github/ubuntu/cpu/build.sh
index e0aac7a0..e00d89df 100644
--- a/kokoro/github/ubuntu/cpu/build.sh
+++ b/kokoro/github/ubuntu/cpu/build.sh
@@ -16,20 +16,35 @@
 set -e
 set -x
 
-cd "${KOKORO_ARTIFACTS_DIR}/github/gnn/"
+echo "Build disabled due to impossibility of using tf-nightly for testing."
+exit 0
+
+PYENV_ROOT="/home/kbuilder/.pyenv"
+PYTHON_VERSION=${PYTHON_VERSION:-"3.11"}
+
+echo "Installing pyenv.."
+git clone https://github.com/pyenv/pyenv.git "$PYENV_ROOT"
+export PATH="/home/kbuilder/.local/bin:$PYENV_ROOT/bin:$PATH"
+eval "$(pyenv init --path)"
 
-sudo update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.9 1
+echo "Python setup..."
+pyenv install -s "$PYTHON_VERSION"
+pyenv global "$PYTHON_VERSION"
+
+cd "${KOKORO_ARTIFACTS_DIR}/github/gnn/"
 
-PYTHON_BINARY="/usr/bin/python3.9"
 PIP_TEST_PREFIX=bazel_pip
 
-"${PYTHON_BINARY}" -m venv venv
+python -m venv venv
 source venv/bin/activate
 
 # Check the python version
 python --version
 python3 --version
 
+# update pip
+pip install --upgrade pip
+
 TEST_ROOT=$(pwd)/${PIP_TEST_PREFIX}
 rm -rf "$TEST_ROOT"
 mkdir -p "$TEST_ROOT"
diff --git a/kokoro/github/ubuntu/cpu/build_versioned.sh b/kokoro/github/ubuntu/cpu/build_versioned.sh
index 47ab64bb..e2e6238c 100644
--- a/kokoro/github/ubuntu/cpu/build_versioned.sh
+++ b/kokoro/github/ubuntu/cpu/build_versioned.sh
@@ -16,24 +16,32 @@
 set -e
 set -x
 
-cd "${KOKORO_ARTIFACTS_DIR}/github/gnn/"
+PYENV_ROOT="/home/kbuilder/.pyenv"
+PYTHON_VERSION=${PYTHON_VERSION:-"3.11"}
+
+echo "Installing pyenv.."
+git clone https://github.com/pyenv/pyenv.git "$PYENV_ROOT"
+export PATH="/home/kbuilder/.local/bin:$PYENV_ROOT/bin:$PATH"
+eval "$(pyenv init --path)"
 
-sudo apt-get install -y "python${PYTHON_VERSION}"
+echo "Python setup..."
+pyenv install -s "$PYTHON_VERSION"
+pyenv global "$PYTHON_VERSION"
 
-# Update alternatives, taken from corresponding Keras OSS test script
-sudo update-alternatives --install /usr/bin/python3 python3 "/usr/bin/python$PYTHON_VERSION" 1
-sudo apt-get install -y python$PYTHON_VERSION-venv
+cd "${KOKORO_ARTIFACTS_DIR}/github/gnn/"
 
-PYTHON_BINARY="/usr/bin/python${PYTHON_VERSION}"
 PIP_TEST_PREFIX=bazel_pip
 
-"${PYTHON_BINARY}" -m venv venv
+python -m venv venv
 source venv/bin/activate
 
 # Debug messages to indicate the python version
 python --version
 python3 --version
 
+# update pip
+pip install --upgrade pip
+
 TEST_ROOT=$(pwd)/${PIP_TEST_PREFIX}
 rm -rf "$TEST_ROOT"
 mkdir -p "$TEST_ROOT"
diff --git a/kokoro/github/ubuntu/cpu/continuous.cfg b/kokoro/github/ubuntu/cpu/continuous.cfg
index 88bf40cb..37bda92b 100644
--- a/kokoro/github/ubuntu/cpu/continuous.cfg
+++ b/kokoro/github/ubuntu/cpu/continuous.cfg
@@ -4,6 +4,10 @@ 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"
+}
 
 action {
   define_artifacts {
diff --git a/kokoro/github/ubuntu/cpu/oldest/continuous.cfg b/kokoro/github/ubuntu/cpu/oldest/continuous.cfg
index e64bd26c..aa98234d 100644
--- a/kokoro/github/ubuntu/cpu/oldest/continuous.cfg
+++ b/kokoro/github/ubuntu/cpu/oldest/continuous.cfg
@@ -6,7 +6,7 @@ env_vars: {
 }
 env_vars: {
   key: "PYTHON_VERSION"
-  value: "3.9"
+  value: "3.11"
 }
 env_vars: {
   key: "TF_VERSION"
diff --git a/kokoro/github/ubuntu/cpu/oldest/presubmit.cfg b/kokoro/github/ubuntu/cpu/oldest/presubmit.cfg
index e64bd26c..aa98234d 100644
--- a/kokoro/github/ubuntu/cpu/oldest/presubmit.cfg
+++ b/kokoro/github/ubuntu/cpu/oldest/presubmit.cfg
@@ -6,7 +6,7 @@ env_vars: {
 }
 env_vars: {
   key: "PYTHON_VERSION"
-  value: "3.9"
+  value: "3.11"
 }
 env_vars: {
   key: "TF_VERSION"
diff --git a/kokoro/github/ubuntu/cpu/presubmit.cfg b/kokoro/github/ubuntu/cpu/presubmit.cfg
index 88bf40cb..37bda92b 100644
--- a/kokoro/github/ubuntu/cpu/presubmit.cfg
+++ b/kokoro/github/ubuntu/cpu/presubmit.cfg
@@ -4,6 +4,10 @@ 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"
+}
 
 action {
   define_artifacts {
diff --git a/setup.py b/setup.py
index ddcddcd5..5948d883 100644
--- a/setup.py
+++ b/setup.py
@@ -178,7 +178,7 @@ def get_version():
         'tensorflow>=2.12.0,<3; platform_machine != "arm64" or platform_system != "Darwin"',
         'tensorflow-macos>=2.12.0,<3; platform_machine == "arm64" and platform_system == "Darwin"',
         # pylint:enable=g-line-too-long
-        'apache-beam',
+        'apache-beam>=2.54.0',
     ],
     python_requires='>=3.9,<4',
     packages=find_namespace_packages(