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/.github/workflows/build-reusable.yml b/.github/workflows/build-reusable.yml new file mode 100644 index 00000000..7c19d10b --- /dev/null +++ b/.github/workflows/build-reusable.yml @@ -0,0 +1,99 @@ +# Copyright 2026 The TensorFlow GNN Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +name: Reusable Build +on: + workflow_call: + inputs: + + python_version: + description: 'Python version (overridden by kokoro_config if set)' + required: false + type: string + default: '3.11' + bazel_version: + description: 'Bazel version (overridden by kokoro_config if set)' + required: false + type: string + default: '7.4.1' + tf_version_spec: + description: 'TensorFlow version specifier' + required: false + type: string + keras_version_spec: + description: 'Keras version specifier' + required: false + type: string + nightly_build: + description: 'Whether to use tf-nightly' + required: false + type: boolean + default: false + legacy_keras: + description: 'Value for TF_USE_LEGACY_KERAS' + required: false + type: string + default: '0' + tag_filters: + description: 'Bazel test tag filters' + required: false + type: string + default: '' + nightly_env: + description: 'Whether to run checks for tf-nightly' + required: false + type: boolean + default: false + run_tests: + description: 'Whether to run tests after building' + required: false + type: boolean + default: false + +permissions: + contents: read + +jobs: + build-and-test: + runs-on: linux-x86-n2-32 + defaults: + run: + shell: bash + container: + image: us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest + timeout-minutes: 60 + steps: + - name: Checkout + uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # ratchet:actions/checkout@v4 + with: + ref: ${{ github.sha }} + persist-credentials: false + + - name: Bazel Version Checkup + run: bazel version || sleep 5 && bazel version # work around potential flakiness on download + env: + USE_BAZEL_VERSION: ${{ inputs.bazel_version }} + + - name: Build and test wheel + env: + PYTHON_VERSION: ${{ inputs.python_version }} + USE_BAZEL_VERSION: ${{ inputs.bazel_version }} + TF_VERSION: ${{ inputs.tf_version_spec }} + KERAS_VERSION_SPEC: ${{ inputs.keras_version_spec }} + TAG_FILTERS: ${{ inputs.tag_filters }} + TF_USE_LEGACY_KERAS: ${{ inputs.legacy_keras }} + NIGHTLY_BUILD: ${{ inputs.nightly_build }} + NIGHTLY_ENV: ${{ inputs.nightly_env }} + + run: ./ci/build-gha.sh diff --git a/.github/workflows/ci-build.yml b/.github/workflows/ci-build.yml new file mode 100644 index 00000000..587aaeb2 --- /dev/null +++ b/.github/workflows/ci-build.yml @@ -0,0 +1,67 @@ +# Copyright 2026 The TensorFlow GNN Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +name: CI + +on: + push: + branches: + - main + pull_request: + branches: + - main + schedule: + - cron: '0 8 * * *' # Run daily at 08:00 UTC + +permissions: + contents: read + issues: write + +concurrency: + group: ${{ github.workflow }}-${{ github.head_ref || github.ref }} + # Don't cancel in-progress jobs for main branches. + cancel-in-progress: ${{ github.ref != 'main' }} + +jobs: + build-and-test-oldest: + uses: ./.github/workflows/build-reusable.yml + with: + python_version: '3.11' + tf_version_spec: '2.12.*' + legacy_keras: '0' + tag_filters: ',-tf_at_least_2_13' + + build-and-test-newest-stable: + uses: ./.github/workflows/build-reusable.yml + with: + python_version: '3.11' + tf_version_spec: '2.20.*' + legacy_keras: '1' + + build-and-test-nightly: + uses: ./.github/workflows/build-reusable.yml + with: + python_version: '3.11' + legacy_keras: '1' + nightly_build: true + + report-failure: + needs: [build-and-test-oldest, build-and-test-newest-stable, build-and-test-nightly] + if: always() && github.event_name == 'schedule' && contains(needs.*.result, 'failure') + runs-on: ubuntu-latest + steps: + - name: Notify failed build + uses: jayqi/failed-build-issue-action@1a893bbf43ef1c2a8705e2b115cd4f0fe3c5649b # v1.2.0 + with: + github-token: ${{ secrets.GITHUB_TOKEN }} diff --git a/BUILD b/BUILD index 78a7dcb4..87a7fcdf 100644 --- a/BUILD +++ b/BUILD @@ -3,6 +3,7 @@ py_library( # This is a dummy rule used as a absl dependency in open-source. # We expect absl to already be installed on the system, e.g. via # `pip install absl` + # strict_deps = False, visibility = ["//visibility:public"], deps = [], ) @@ -12,6 +13,7 @@ py_library( # This is a dummy rule used as a absl dependency in open-source. # We expect absl to already be installed on the system, e.g. via # `pip install absl` + # strict_deps = False, visibility = ["//visibility:public"], deps = [], ) @@ -21,6 +23,7 @@ py_library( # This is a dummy rule used as a absl dependency in open-source. # We expect absl to already be installed on the system, e.g. via # `pip install absl` + # strict_deps = False, visibility = ["//visibility:public"], deps = [], ) @@ -30,6 +33,7 @@ py_library( # This is a dummy rule used as a numpy dependency in open-source. # We expect numpy to already be installed on the system, e.g. via # `pip install numpy` + # strict_deps = False, visibility = ["//visibility:public"], deps = [], ) @@ -39,6 +43,7 @@ py_library( # This is a dummy rule used as a tensorflow dependency in open-source. # We expect tensorflow to already be installed on the system, e.g. via # `pip install tensorflow` + # strict_deps = False, visibility = ["//visibility:public"], deps = [], ) @@ -48,6 +53,7 @@ py_library( # This is a dummy rule used as a ai-edge-litert dependency in open-source. # We expect ai-edge-litert to already be installed on the system, e.g. via # `pip install ai-edge-litert` + # strict_deps = False, visibility = ["//visibility:public"], deps = [], ) @@ -57,6 +63,7 @@ py_library( # This is a dummy rule used as a tensorflow dependency in open-source. # We expect tensorflow to already be installed on the system, e.g. via # `pip install tensorflow` + # strict_deps = False, visibility = ["//visibility:public"], deps = [], ) @@ -66,6 +73,7 @@ py_library( # This is a dummy rule used as a dependency on vizier.service.pyvizier in open-source. # We expect Vizier to already be installed on the system, e.g. via # `pip install google-vizier` and extra steps if needed (b/254806045). + # strict_deps = False, visibility = ["//visibility:public"], deps = [], ) @@ -75,6 +83,7 @@ py_library( # This is a dummy rule used as a mock dependency in open-source tests. # We expect mock to already be installed on the system, e.g. via # `pip install mock` + # strict_deps = False, visibility = ["//visibility:public"], deps = [], ) diff --git a/ci/build-gha.sh b/ci/build-gha.sh new file mode 100755 index 00000000..2707253c --- /dev/null +++ b/ci/build-gha.sh @@ -0,0 +1,132 @@ +#!/bin/bash +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +set -e +set -x + +# For pyenv python installation on ml-build image +sudo apt-get update +sudo apt-get install -y libbz2-dev liblzma-dev libncurses-dev libffi-dev libssl-dev libreadline-dev libsqlite3-dev zlib1g-dev + +PYENV_ROOT="$HOME/.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/.local/bin:$PYENV_ROOT/bin:$PATH" +eval "$(pyenv init --path)" + +echo "Python setup..." +pyenv install -s "$PYTHON_VERSION" +pyenv global "$PYTHON_VERSION" + +PIP_TEST_PREFIX=bazel_pip + +python -m venv build_venv +source build_venv/bin/activate + +# Debug messages to indicate the python version +python --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" +ln -s "$(pwd)"/tensorflow_gnn "$TEST_ROOT"/tensorflow_gnn + +# Print the OS version +cat /etc/os-release + +# Prepend common tag filters to a defined env_var +# For example, tests for TF 2.8 shouldn't run RNG-dependent tests +# These tag filters are enforced to start with a comma for separation +tag_filters="-no_oss,-oss_excluded${TAG_FILTERS}" + +# Check that `bazel` does version selection as expected. +if [[ -n "${USE_BAZEL_VERSION}" && $(bazel --version) != *${USE_BAZEL_VERSION}* ]]; then + echo "Mismatch of configured and actual bazel version (see logged [[ command)" + exit 1 +fi + +bazel clean + +if [[ "$NIGHTLY_BUILD" == "true" ]]; then + pip install --group test-nightly --progress-bar off --upgrade +else + force_tensorflow_version +fi + +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 + +if [[ "$NIGHTLY_BUILD" == "true" ]]; then + pip install dist/tensorflow_gnn-*.whl + pip uninstall -y tensorflow tf-keras ai-edge-litert + pip install --group test-nightly --progress-bar off --upgrade + + # Check that tf-nightly is installed but tensorflow is not + # Also check that tf-keras-nightly is installed. + if [[ $(pip freeze | grep -q tf_nightly=; echo $?) -eq 0 && $(pip freeze | grep -q tensorflow=; echo $?) -eq 0 ]]; then + echo "Found tensorflow and tf_nightly in the environment." + exit 1 + fi + if [[ $(pip freeze | grep -q tf_keras-nightly=; echo $?) -eq 0 && $(pip freeze | grep -q tf_keras=; echo $?) -eq 0 ]]; then + echo "Found tf_keras and tf_keras-nightly in the environment." + exit 1 + fi + +else + force_tensorflow_version + + if [[ "$TF_USE_LEGACY_KERAS" == 1 ]]; then + pip install --group test-tf216plus --progress-bar off --upgrade + else + pip install --group test-pre-tf216 --progress-bar off --upgrade + fi + + pip install dist/tensorflow_gnn-*.whl +fi + +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/... diff --git a/kokoro/github/ubuntu/cpu/build.sh b/kokoro/github/ubuntu/cpu/build.sh index 980b7768..0e6a82dc 100644 --- a/kokoro/github/ubuntu/cpu/build.sh +++ b/kokoro/github/ubuntu/cpu/build.sh @@ -32,16 +32,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" @@ -55,25 +57,34 @@ 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 -# tf-nightly 2.21 lifted the protobuf version to what google-vizier 0.1.24 wants -# but what google-vizier 0.1.21 together with apache-beam does not allow. -# TODO: b/482003228 - Replace this patch by a proper dependency cleanup. -perl -i -lpe '$k+= s/google-vizier>=0.0.13,!=0.1.23,!=0.1.24/google-vizier>=0.0.13,!=0.1.23/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 uninstall -y tensorflow tf-keras ai-edge-litert +pip install --group test-nightly --progress-bar off --upgrade echo "Final packages after all pip commands:" pip list # Check that tf-nightly is installed but tensorflow is not # Also check that tf-keras-nightly is installed. -pip freeze | grep -q tf-nightly= && ! pip freeze | grep -q tensorflow= -pip freeze | grep -q tf-keras-nightly= && ! pip freeze | grep -q tf-keras= +if [[ $(pip freeze | grep -q tf_nightly=; echo $?) -eq 0 && $(pip freeze | grep -q tensorflow=; echo $?) -eq 0 ]]; then + echo "Found tensorflow and tf_nightly in the environment." + exit 1 +fi +if [[ $(pip freeze | grep -q tf_keras-nightly=; echo $?) -eq 0 && $(pip freeze | grep -q tf_keras=; echo $?) -eq 0 ]]; then + echo "Found tf_keras and tf_keras-nightly in the environment." + exit 1 +fi # The env variable is needed to ensure that TF keras still behaves like keras 2 bazel test --test_env="TF_USE_LEGACY_KERAS=1" --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/... 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..c95c9512 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,91 @@ +[project] +name = "tensorflow_gnn" +description = "A library for building scalable graph neural networks in TensorFlow." +# TODO(b/316329189): The relative links from README.md don't work on PyPI. +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, <=2.67", + # This dependency comes from apache beam implementing it as an extra for 2.70. + # To maintain current compatibility with pre and post 2.70, we need to include it here. b/465442394 + "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 :: 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", + "google-vizier==0.1.24", # b/394062744 + {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 - ], - } ) diff --git a/tensorflow_gnn/converters/networkx_io.py b/tensorflow_gnn/converters/networkx_io.py index 32004752..942f4d20 100644 --- a/tensorflow_gnn/converters/networkx_io.py +++ b/tensorflow_gnn/converters/networkx_io.py @@ -75,10 +75,7 @@ def to_edge_list(tf_gnn_graph): node_cnt = 0 for node_set_name, node_set in sorted(tf_gnn_graph.node_sets.items()): - if len(node_set.sizes) > 1: - size = int(node_set.sizes[0]) - else: - size = int(node_set.sizes) + size = int(node_set.sizes[0]) if '#id' in node_set.features and list(node_set.features['#id']): id_range = [str(x.numpy()).encode('utf-8') for x in node_set['#id']] else: diff --git a/tensorflow_gnn/converters/networkx_io_test.py b/tensorflow_gnn/converters/networkx_io_test.py index 4017f207..90ab2982 100644 --- a/tensorflow_gnn/converters/networkx_io_test.py +++ b/tensorflow_gnn/converters/networkx_io_test.py @@ -59,10 +59,10 @@ def test_two_node_sets(self): nx_graph = networkx_io.to_networkx_graph(tf_gnn_graph) - num_tf_nodes = int(tf_gnn_graph.node_sets['nodes'].sizes) + int( - tf_gnn_graph.node_sets['node_set2'].sizes + num_tf_nodes = int(tf_gnn_graph.node_sets['nodes'].sizes[0]) + int( + tf_gnn_graph.node_sets['node_set2'].sizes[0] ) - num_tf_edges = int(tf_gnn_graph.edge_sets['edges'].sizes) + num_tf_edges = int(tf_gnn_graph.edge_sets['edges'].sizes[0]) self.assertEqual(num_tf_nodes, nx_graph.number_of_nodes()) self.assertEqual(num_tf_edges, nx_graph.number_of_edges()) @@ -92,8 +92,8 @@ def test_three_components(self): nx_graph = networkx_io.to_networkx_graph(tf_gnn_graph) - num_tf_nodes = int(tf_gnn_graph.node_sets['nodes'].sizes) - num_tf_edges = int(tf_gnn_graph.edge_sets['edges'].sizes) + num_tf_nodes = int(tf_gnn_graph.node_sets['nodes'].sizes[0]) + num_tf_edges = int(tf_gnn_graph.edge_sets['edges'].sizes[0]) self.assertEqual(num_tf_nodes, nx_graph.number_of_nodes()) self.assertEqual(num_tf_edges, nx_graph.number_of_edges()) diff --git a/tensorflow_gnn/converters/triples_test.py b/tensorflow_gnn/converters/triples_test.py index 7f67f4db..01119f5b 100644 --- a/tensorflow_gnn/converters/triples_test.py +++ b/tensorflow_gnn/converters/triples_test.py @@ -30,8 +30,8 @@ def test_triples_from_array(self): ["nns", "are", "awesome"]] gt = triples.triple_to_graphtensor(triples=spos) - is_source = int(gt.edge_sets["is"].adjacency.source) - is_target = int(gt.edge_sets["is"].adjacency.target) + is_source = int(gt.edge_sets["is"].adjacency.source[0]) + is_target = int(gt.edge_sets["is"].adjacency.target[0]) self.assertLen(gt.node_sets, 1) self.assertLen(gt.edge_sets, 2) diff --git a/tensorflow_gnn/experimental/sampler/BUILD b/tensorflow_gnn/experimental/sampler/BUILD index 0d351af9..39a13321 100644 --- a/tensorflow_gnn/experimental/sampler/BUILD +++ b/tensorflow_gnn/experimental/sampler/BUILD @@ -6,6 +6,7 @@ py_test( name = "core_test", srcs = ["core_test.py"], python_version = "PY3", + # strict_deps = False, deps = [], ) @@ -13,6 +14,7 @@ py_test( name = "ext_ops_test", srcs = ["ext_ops_test.py"], python_version = "PY3", + # strict_deps = False, deps = [], ) @@ -20,5 +22,6 @@ py_test( name = "eval_dag_test", srcs = ["eval_dag_test.py"], python_version = "PY3", + # strict_deps = False, deps = [], -) \ No newline at end of file +) diff --git a/tensorflow_gnn/experimental/sampler/beam/BUILD b/tensorflow_gnn/experimental/sampler/beam/BUILD index 57feaeca..4789ca60 100644 --- a/tensorflow_gnn/experimental/sampler/beam/BUILD +++ b/tensorflow_gnn/experimental/sampler/beam/BUILD @@ -203,6 +203,7 @@ pytype_strict_contrib_test( py_binary( name = "sampler", srcs = ["sampler.py"], + # strict_deps = False, deps = [ ":sampler-lib", "//third_party/py/google/protobuf:use_fast_cpp_protos", # Automatically added go/proto_python_upb_flip diff --git a/tensorflow_gnn/experimental/sampler/proto/BUILD b/tensorflow_gnn/experimental/sampler/proto/BUILD index 6d73ac44..b51d6096 100644 --- a/tensorflow_gnn/experimental/sampler/proto/BUILD +++ b/tensorflow_gnn/experimental/sampler/proto/BUILD @@ -13,6 +13,7 @@ py_library( name = "proto", srcs = ["__init__.py"], srcs_version = "PY3", + # strict_deps = False, visibility = ["//visibility:public"], deps = [ ":eval_dag_py_proto", @@ -20,7 +21,6 @@ py_library( ], ) - proto_library( name = "eval_dag_proto", srcs = ["eval_dag.proto"], diff --git a/tensorflow_gnn/proto/BUILD b/tensorflow_gnn/proto/BUILD index ec3fd48d..076cadfd 100644 --- a/tensorflow_gnn/proto/BUILD +++ b/tensorflow_gnn/proto/BUILD @@ -13,6 +13,7 @@ py_library( name = "proto", srcs = ["__init__.py"], srcs_version = "PY3", + # strict_deps = False, visibility = ["//visibility:public"], deps = [ ":graph_schema", @@ -39,6 +40,7 @@ py_library( name = "graph_schema", srcs = ["graph_schema.py"], srcs_version = "PY3", + # strict_deps = False, deps = [ ":graph_schema_py_proto", ],