diff --git a/.bazelversion b/.bazelversion
new file mode 100644
index 00000000..c4d4fe2f
--- /dev/null
+++ b/.bazelversion
@@ -0,0 +1,2 @@
+7.7.0
+# TODO: b/390391579 - Upgrade once bazel 8+ works.
\ No newline at end of file
diff --git a/kokoro/github/ubuntu/cpu/build.sh b/kokoro/github/ubuntu/cpu/build.sh
index c0156085..5ec056d2 100644
--- a/kokoro/github/ubuntu/cpu/build.sh
+++ b/kokoro/github/ubuntu/cpu/build.sh
@@ -36,16 +36,18 @@ cd "${KOKORO_ARTIFACTS_DIR}/github/gnn/"
 
 PIP_TEST_PREFIX=bazel_pip
 
-python -m venv venv
-source venv/bin/activate
+python -m venv build_venv
+source build_venv/bin/activate
 
 # Check the python version
 python --version
-python3 --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"
@@ -59,14 +61,20 @@ if [[ -n "${USE_BAZEL_VERSION}" && $(bazel --version) != *${USE_BAZEL_VERSION}*
 fi
 
 bazel clean
-pip install -r requirements-dev.txt --progress-bar off
-pip install tf-keras-nightly tf-nightly --progress-bar off --upgrade
-# We need to remove the dependency on tensorflow to test nightly
-# The dependencies will be provided by tf-nightly
-perl  -i -lpe '$k+= s/tensorflow>=2\.[0-9]+\.[0-9]+(,<=?[0-9.]+)?;/tf-nightly;/g; END{exit($k != 1)}' setup.py
-python3 setup.py bdist_wheel
-pip uninstall -y tensorflow_gnn
+pip install --group test-nightly --progress-bar off --upgrade
+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
+
+pip install --upgrade pip
 pip install dist/tensorflow_gnn-*.whl
+pip install --group test-nightly --progress-bar off --upgrade
 
 echo "Final packages after all pip commands:"
 pip list
diff --git a/kokoro/github/ubuntu/cpu/build_versioned.sh b/kokoro/github/ubuntu/cpu/build_versioned.sh
index e2e6238c..ba38fd7d 100644
--- a/kokoro/github/ubuntu/cpu/build_versioned.sh
+++ b/kokoro/github/ubuntu/cpu/build_versioned.sh
@@ -19,6 +19,18 @@ set -x
 PYENV_ROOT="/home/kbuilder/.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/kbuilder/.local/bin:$PYENV_ROOT/bin:$PATH"
@@ -32,16 +44,18 @@ cd "${KOKORO_ARTIFACTS_DIR}/github/gnn/"
 
 PIP_TEST_PREFIX=bazel_pip
 
-python -m venv venv
-source venv/bin/activate
+python -m venv build_venv
+source build_venv/bin/activate
 
 # Debug messages to indicate the python version
 python --version
-python3 --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"
@@ -62,13 +76,27 @@ if [[ -n "${USE_BAZEL_VERSION}" && $(bazel --version) != *${USE_BAZEL_VERSION}*
 fi
 
 bazel clean
-pip install -r requirements-dev.txt --progress-bar off
-pip install tensorflow=="${TF_VERSION}" --progress-bar off --upgrade
+force_tensorflow_version
+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
+force_tensorflow_version
+
 if [[ "$TF_USE_LEGACY_KERAS" == 1 ]]; then
-  pip install tf-keras=="${TF_VERSION}" --progress-bar off --upgrade
+  pip install --group test-tf216plus --progress-bar off --upgrade
+else
+  pip install --group test-pre-tf216 --progress-bar off --upgrade
 fi
-python3 setup.py bdist_wheel
-pip uninstall -y tensorflow_gnn
+
 pip install dist/tensorflow_gnn-*.whl
 
 echo "Final packages after all pip commands:"
diff --git a/kokoro/github/ubuntu/cpu/continuous.cfg b/kokoro/github/ubuntu/cpu/continuous.cfg
index 37bda92b..dad6cf34 100644
--- a/kokoro/github/ubuntu/cpu/continuous.cfg
+++ b/kokoro/github/ubuntu/cpu/continuous.cfg
@@ -1,9 +1,5 @@
 build_file: "gnn/kokoro/github/ubuntu/cpu/build.sh"
 
-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"
diff --git a/kokoro/github/ubuntu/cpu/newest_stable/continuous.cfg b/kokoro/github/ubuntu/cpu/newest_stable/continuous.cfg
index 890a2916..fe95bbb4 100644
--- a/kokoro/github/ubuntu/cpu/newest_stable/continuous.cfg
+++ b/kokoro/github/ubuntu/cpu/newest_stable/continuous.cfg
@@ -1,9 +1,5 @@
 build_file: "gnn/kokoro/github/ubuntu/cpu/build_versioned.sh"
 
-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"
diff --git a/kokoro/github/ubuntu/cpu/newest_stable/presubmit.cfg b/kokoro/github/ubuntu/cpu/newest_stable/presubmit.cfg
index 890a2916..fe95bbb4 100644
--- a/kokoro/github/ubuntu/cpu/newest_stable/presubmit.cfg
+++ b/kokoro/github/ubuntu/cpu/newest_stable/presubmit.cfg
@@ -1,9 +1,5 @@
 build_file: "gnn/kokoro/github/ubuntu/cpu/build_versioned.sh"
 
-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"
diff --git a/kokoro/github/ubuntu/cpu/oldest/continuous.cfg b/kokoro/github/ubuntu/cpu/oldest/continuous.cfg
index aa98234d..d5dc9298 100644
--- a/kokoro/github/ubuntu/cpu/oldest/continuous.cfg
+++ b/kokoro/github/ubuntu/cpu/oldest/continuous.cfg
@@ -1,9 +1,5 @@
 build_file: "gnn/kokoro/github/ubuntu/cpu/build_versioned.sh"
 
-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"
diff --git a/kokoro/github/ubuntu/cpu/oldest/presubmit.cfg b/kokoro/github/ubuntu/cpu/oldest/presubmit.cfg
index aa98234d..d5dc9298 100644
--- a/kokoro/github/ubuntu/cpu/oldest/presubmit.cfg
+++ b/kokoro/github/ubuntu/cpu/oldest/presubmit.cfg
@@ -1,9 +1,5 @@
 build_file: "gnn/kokoro/github/ubuntu/cpu/build_versioned.sh"
 
-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"
diff --git a/kokoro/github/ubuntu/cpu/presubmit.cfg b/kokoro/github/ubuntu/cpu/presubmit.cfg
index 37bda92b..dad6cf34 100644
--- a/kokoro/github/ubuntu/cpu/presubmit.cfg
+++ b/kokoro/github/ubuntu/cpu/presubmit.cfg
@@ -1,9 +1,5 @@
 build_file: "gnn/kokoro/github/ubuntu/cpu/build.sh"
 
-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"
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 00000000..9d6fb624
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,88 @@
+[project]
+name = "tensorflow_gnn"
+description = "A library for building scalable graph neural networks in TensorFlow."
+readme = "README.md"
+keywords = ["tensorflow", "gnn", "graph"]
+dynamic = ["version"]
+license = "Apache-2.0"
+requires-python = ">=3.9, <4"
+dependencies = [
+  "google-vizier>=0.0.13, !=0.1.23, !=0.1.24",  # b/394062744
+  "ml-collections>=1.1.0",
+  "networkx>=3.6.1",
+  "pyarrow>=18.1.0",
+  "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'",
+  "apache-beam>=2.54",
+  "crcmod>=1.7, <2",
+]
+authors = [
+  {name = "Google LLC", email = "tensorflow-gnn@googlegroups.com"},
+]
+classifiers = [
+  "Development Status :: 5 - Production/Stable",
+  "Intended Audience :: Developers",
+  "Intended Audience :: Science/Research",
+  "Operating System :: MacOS :: MacOS X",
+  "Operating System :: POSIX :: Linux",
+  "Operating System :: Microsoft :: Windows",
+  "Programming Language :: Python",
+  "Programming Language :: Python :: 3",
+  "Programming Language :: Python :: 3.9",
+  "Programming Language :: Python :: 3 :: Only",
+  "Topic :: Scientific/Engineering",
+  "Topic :: Scientific/Engineering :: Artificial Intelligence",
+  "Topic :: Scientific/Engineering :: Mathematics",
+  "Topic :: Software Development",
+  "Topic :: Software Development :: Libraries",
+  "Topic :: Software Development :: Libraries :: Python Modules",
+]
+
+[project.urls]
+homepage = "https://github.com/tensorflow/gnn"
+source = "https://github.com/tensorflow/gnn.git"
+
+[dependency-groups]
+test-base = [
+  "mock",
+]
+test-pre-tf216 = [
+  "ai-edge-litert",
+  {include-group = "test-base"},
+]
+test-tf216plus = [
+  "tf-keras",
+  {include-group = "test-pre-tf216"},
+]
+test-nightly = [
+  "tf-nightly",
+  "tf-keras-nightly",
+  "ai-edge-litert-nightly",
+  {include-group = "test-base"},
+]
+
+[project.scripts]
+tfgnn_convert_ogb_dataset = 'tensorflow_gnn.converters.ogb.convert_ogb_dataset:main'
+tfgnn_sampler = 'tensorflow_gnn.experimental.sampler.beam.sampler:main'
+# copybara:comment_begin(NetworkX utils)
+tfgnn_nx_converter = 'tensorflow_gnn.sampler.nx_converter:main'
+tfgnn_nx_generator = 'tensorflow_gnn.sampler.nx_generator:main'
+# copybara:comment_end
+tfgnn_generate_training_data = 'tensorflow_gnn.tools.generate_training_data:main'
+tfgnn_print_training_data = 'tensorflow_gnn.tools.print_training_data:main'
+tfgnn_sampled_stats = 'tensorflow_gnn.tools.sampled_stats:main'
+tfgnn_validate_graph_schema = 'tensorflow_gnn.tools.validate_graph_schema:main'
+
+[build-system]
+requires = ["setuptools", "numpy"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools.dynamic]
+version = {attr = "tensorflow_gnn.version.__version__"}
+
+[tool.setuptools.packages.find]
+exclude = ["examples*", "package*", "testdata*", "build*", "bazel*"]
+
+[tool.setuptools.package-data]
+tensorflow_gnn = ["*.proto"]
+
diff --git a/requirements-dev.txt b/requirements-dev.txt
deleted file mode 100644
index 680bce7d..00000000
--- a/requirements-dev.txt
+++ /dev/null
@@ -1,4 +0,0 @@
-mock
-wheel
-# includes glibcxx older than 3.4.29
-ai-edge-litert-nightly
\ No newline at end of file
diff --git a/setup.py b/setup.py
index 23504f8e..df017025 100644
--- a/setup.py
+++ b/setup.py
@@ -16,20 +16,13 @@
 
 import os
 import platform
+import shutil
 import subprocess
 import sys
-
 import setuptools
-from setuptools import find_namespace_packages
-from setuptools import setup
-from setuptools.command.install import install
-from setuptools.dist import Distribution
-# pylint:disable=g-bad-import-order
-# setuptools must be imported prior to distutils.
-# TODO(b/317036825): Stop using distutils, Python 3.12 doesn't have it.
-from distutils import spawn
-from distutils.command import build
-# pylint:enable=g-bad-import-order
+from setuptools import dist
+from setuptools.command import build
+from setuptools.command import install
 
 
 class _BuildCommand(build.build):
@@ -57,7 +50,8 @@ def _build_cc_extensions(self):
   # sequentially when running a "build" command, if the second item in the tuple
   # (predicate method) is evaluated to true.
   sub_commands = [
-      ('bazel_build', _build_cc_extensions)] + build.build.sub_commands
+      ('bazel_build', _build_cc_extensions)
+  ] + build.build.sub_commands
 
 
 class _BazelBuildCommand(setuptools.Command):
@@ -71,7 +65,7 @@ def initialize_options(self):
     pass
 
   def finalize_options(self):
-    self._bazel_cmd = spawn.find_executable('bazel')
+    self._bazel_cmd = shutil.which('bazel')
     if not self._bazel_cmd:
       raise RuntimeError(
           'Could not find "bazel" binary. Please visit '
@@ -98,14 +92,14 @@ def run(self):
 # TFDV is not a purelib. However because of the extension module is not built
 # by setuptools, it will be incorrectly treated as a purelib. The following
 # works around that bug.
-class _InstallPlatlibCommand(install):
+class _InstallPlatlibCommand(install.install):
 
   def finalize_options(self):
-    install.finalize_options(self)
+    install.install.finalize_options(self)
     self.install_lib = self.install_platlib
 
 
-class _SourceDistributionWithProtos(Distribution):
+class _SourceDistributionWithProtos(dist.Distribution):
   """Proto-only deps don't require OS specific wheels."""
 
   def is_pure(self):
@@ -115,96 +109,11 @@ def has_ext_modules(self):
     return False
 
 
-def get_version():
-  """Get version from version module."""
-  version_path = os.path.join(os.path.dirname(__file__), 'tensorflow_gnn')
-  sys.path.insert(0, version_path)
-  # pytype: disable=import-error  # pylint: disable=g-import-not-at-top
-  from version import __version__ as v
-  return v
-
-
-# Get the long description from the README file.
-# TODO(b/316329189): The relative links from README.md don't work on PyPI.
-with open('README.md') as fp:
-  _LONG_DESCRIPTION = fp.read()
-
-
-console_scripts = [
-    'tensorflow_gnn.converters.ogb.convert_ogb_dataset',
-    'tensorflow_gnn.experimental.sampler.beam.sampler',
-    # copybara:uncomment_begin(NetworkX utils)
-    # 'tensorflow_gnn.sampler.nx_converter',
-    # 'tensorflow_gnn.sampler.nx_generator',
-    # copybara:uncomment_end
-    'tensorflow_gnn.tools.generate_training_data',
-    'tensorflow_gnn.tools.print_training_data',
-    'tensorflow_gnn.tools.sampled_stats',
-    'tensorflow_gnn.tools.validate_graph_schema',
-]
-
-
-setup(
-    name='tensorflow-gnn',
-    version=get_version(),
-    author='Google LLC',
-    author_email='tensorflow-gnn@googlegroups.com',
-    license='Apache 2.0',
-    classifiers=[
-        'Development Status :: 5 - Production/Stable',
-        'Intended Audience :: Developers',
-        'Intended Audience :: Science/Research',
-        'License :: OSI Approved :: Apache Software License',
-        'Operating System :: MacOS :: MacOS X',
-        'Operating System :: POSIX :: Linux',
-        'Operating System :: Microsoft :: Windows',
-        'Programming Language :: Python',
-        'Programming Language :: Python :: 3',
-        'Programming Language :: Python :: 3 :: Only',
-        'Topic :: Scientific/Engineering',
-        'Topic :: Scientific/Engineering :: Artificial Intelligence',
-        'Topic :: Scientific/Engineering :: Mathematics',
-        'Topic :: Software Development',
-        'Topic :: Software Development :: Libraries',
-        'Topic :: Software Development :: Libraries :: Python Modules',
-    ],
-    namespace_packages=[],
-    install_requires=[
-        'google-vizier>=0.0.13,!=0.1.23,!=0.1.24',  # b/394062744
-        'ml-collections',
-        'networkx',
-        'pyarrow',
-        # pylint:disable=g-line-too-long
-        '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>=2.54.0',
-        'crcmod>=1.7,<2',
-    ],
-    python_requires='>=3.9,<4',
-    packages=find_namespace_packages(
-        exclude=['examples*', 'package*', 'testdata*'],
-    ),
-    include_package_data=True,
-    package_data={'': ['*.proto']},
-    zip_safe=False,
+setuptools.setup(
     distclass=_SourceDistributionWithProtos,
-    description='A library for building scalable graph neural networks in TensorFlow.',
-    long_description=_LONG_DESCRIPTION,
-    long_description_content_type='text/markdown',
-    keywords='tensorflow gnn graph',
-    url='https://github.com/tensorflow/gnn',
-    download_url='https://github.com/tensorflow/gnn.git',
-    requires=[],
     cmdclass={
         'install': _InstallPlatlibCommand,
         'build': _BuildCommand,
         'bazel_build': _BazelBuildCommand,
     },
-    entry_points={
-        'console_scripts': [
-            'tfgnn_{}={}:main'.format(libname.split('.')[-1], libname)
-            for libname in console_scripts
-        ],
-    }
 )