Migrate package metadata to pyproject.toml.

.bazelversion

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

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

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:"

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"

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"

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"

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"

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"

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"

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 = "[email protected]"},
+]
+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"]
+