diff --git a/kokoro/github/ubuntu/cpu/build_versioned.sh b/kokoro/github/ubuntu/cpu/build_versioned.sh
index 52389159..283fee92 100644
--- a/kokoro/github/ubuntu/cpu/build_versioned.sh
+++ b/kokoro/github/ubuntu/cpu/build_versioned.sh
@@ -22,6 +22,7 @@ sudo apt-get install -y "python${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
 
 PYTHON_BINARY="/usr/bin/python${PYTHON_VERSION}"
 PIP_TEST_PREFIX=bazel_pip
diff --git a/kokoro/github/ubuntu/cpu/newest_stable/continuous.cfg b/kokoro/github/ubuntu/cpu/newest_stable/continuous.cfg
index f04d5e99..92ca15d7 100644
--- a/kokoro/github/ubuntu/cpu/newest_stable/continuous.cfg
+++ b/kokoro/github/ubuntu/cpu/newest_stable/continuous.cfg
@@ -2,7 +2,7 @@ build_file: "gnn/kokoro/github/ubuntu/cpu/build_versioned.sh"
 
 env_vars: {
   key: "PYTHON_VERSION"
-  value: "3.10"
+  value: "3.11"
 }
 env_vars: {
   key: "TF_VERSION"
diff --git a/kokoro/github/ubuntu/cpu/newest_stable/presubmit.cfg b/kokoro/github/ubuntu/cpu/newest_stable/presubmit.cfg
index f04d5e99..92ca15d7 100644
--- a/kokoro/github/ubuntu/cpu/newest_stable/presubmit.cfg
+++ b/kokoro/github/ubuntu/cpu/newest_stable/presubmit.cfg
@@ -2,7 +2,7 @@ build_file: "gnn/kokoro/github/ubuntu/cpu/build_versioned.sh"
 
 env_vars: {
   key: "PYTHON_VERSION"
-  value: "3.10"
+  value: "3.11"
 }
 env_vars: {
   key: "TF_VERSION"
diff --git a/setup.py b/setup.py
index ed6225c9..c2391968 100644
--- a/setup.py
+++ b/setup.py
@@ -169,8 +169,11 @@ def get_version():
     namespace_packages=[],
     install_requires=[
         'google-vizier>=0.0.13',
+        # TODO: b/306688501 - Unpin grpcio when we figure out issue with TF 2.14
+        'grpcio~=1.59.0',
         'ml-collections',
         'networkx',
+        'protobuf~=3.20.3',
         'pyarrow',
         # pylint:disable=g-line-too-long
         'tensorflow>=2.10.0; platform_machine != "arm64" or platform_system != "Darwin"',