From a2197a764e6d4974f506eff950fec18e4ac17223 Mon Sep 17 00:00:00 2001 From: Younghee Kwon Date: Thu, 23 Jul 2026 12:10:35 -0700 Subject: [PATCH] Fixed installation from source. - `setup.py` now builds proto files. - it also builds pybind11 modules. - replaced absl string library to std string in the pybind extension to avoid adding the dependency in the setup tool stage. PiperOrigin-RevId: 952888201 --- .../experimental/index_shuffle/python/BUILD | 5 +- .../python/index_shuffle_module.cc | 17 ++-- grain/oss/build_whl.sh | 2 +- pyproject.toml | 2 +- setup.py | 87 ++++++++++++++++++- 5 files changed, 97 insertions(+), 16 deletions(-) diff --git a/grain/_src/python/experimental/index_shuffle/python/BUILD b/grain/_src/python/experimental/index_shuffle/python/BUILD index 2ae25f6d5..351ee45e3 100644 --- a/grain/_src/python/experimental/index_shuffle/python/BUILD +++ b/grain/_src/python/experimental/index_shuffle/python/BUILD @@ -7,10 +7,7 @@ licenses(["notice"]) pybind_extension( name = "index_shuffle_module", srcs = ["index_shuffle_module.cc"], - deps = [ - "//grain/_src/python/experimental/index_shuffle", - "@abseil-cpp//absl/strings", - ], + deps = ["//grain/_src/python/experimental/index_shuffle"], ) py_test( diff --git a/grain/_src/python/experimental/index_shuffle/python/index_shuffle_module.cc b/grain/_src/python/experimental/index_shuffle/python/index_shuffle_module.cc index d2ab15c4a..faefa4861 100644 --- a/grain/_src/python/experimental/index_shuffle/python/index_shuffle_module.cc +++ b/grain/_src/python/experimental/index_shuffle/python/index_shuffle_module.cc @@ -1,8 +1,8 @@ #include #include +#include -#include "absl/strings/str_cat.h" #include "grain/_src/python/experimental/index_shuffle/index_shuffle.h" namespace py = pybind11; @@ -16,15 +16,18 @@ PYBIND11_MODULE(index_shuffle_module, m) { "index_shuffle", [](int64_t index, int64_t max_index, uint32_t seed, uint32_t rounds) { if (rounds < 4 || rounds % 2 != 0 || rounds > 1024) { + // Using std::to_string to avoid extra dependency in setup.py. throw py::value_error( - absl::StrCat("rounds must be an even integer between 4 and 1024, " - "but got rounds = ", - rounds)); + "rounds must be an even integer between 4 and 1024, " + "but got rounds = " + + std::to_string(rounds)); } if (index < 0 || index > max_index) { - throw py::value_error(absl::StrCat( - "index must be in [0, max_index], but got index = ", index, - " and max_index = ", max_index)); + // Using std::to_string to avoid extra dependency in setup.py. + throw py::value_error( + "index must be in [0, max_index], but got index = " + + std::to_string(index) + + " and max_index = " + std::to_string(max_index)); } return grain::random::index_shuffle(index, max_index, seed, rounds); }, diff --git a/grain/oss/build_whl.sh b/grain/oss/build_whl.sh index 4d3f78e52..94af5863a 100644 --- a/grain/oss/build_whl.sh +++ b/grain/oss/build_whl.sh @@ -112,7 +112,7 @@ main() { if [ "$(uname)" == "Darwin" ]; then WHEEL_BLD_ARGS="${WHEEL_BLD_ARGS} --plat-name macosx_11_0_$(uname -m)" fi - "$PYTHON_BIN" setup.py $WHEEL_BLD_ARGS + GRAIN_SKIP_EXTRA_BUILD=1 "$PYTHON_BIN" setup.py $WHEEL_BLD_ARGS if [ -n "${AUDITWHEEL_PLATFORM}" ]; then printf '%s : "=== Auditing wheel\n' "$(date)" diff --git a/pyproject.toml b/pyproject.toml index 4ab35874e..feca1dbcd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,5 +1,5 @@ [build-system] -requires = ["setuptools"] +requires = ["setuptools", "grpcio-tools", "pybind11"] build-backend = "setuptools.build_meta" [project] diff --git a/setup.py b/setup.py index c93fc3afe..9ae3781fe 100644 --- a/setup.py +++ b/setup.py @@ -2,10 +2,19 @@ Most project configs are in `pyproject.toml` -- prefer to modify `pyproject.toml` over this file if possible. + +There are extra building steps in this script: + 1. Generates proto bindings. + 2. Compiles C++ extensions using pybind11. +In case those are pre-built (e.g. when building wheels), you can skip this +extra build using `GRAIN_SKIP_EXTRA_BUILD=1`. """ +import importlib.resources +import os import setuptools from setuptools import dist +from setuptools.command import build_py class BinaryDistribution(dist.Distribution): @@ -15,6 +24,78 @@ def has_ext_modules(self): return True -setuptools.setup( - distclass=BinaryDistribution, -) +# If GRAIN_SKIP_EXTRA_BUILD=1, we are packaging pre-built binaries (e.g. in +# build_whl.sh). We can skip all build logic (C++ compilation and Proto +# generation). +if os.environ.get("GRAIN_SKIP_EXTRA_BUILD", "0") == "1": + setuptools.setup( + distclass=BinaryDistribution, + ) +else: + # Full build logic for installing from source + + class GenerateProtosCommand(setuptools.Command): + """Command to generate Python protobuf bindings.""" + + description = "Generate Python protobuf bindings" + user_options = [] + + def initialize_options(self): + pass + + def finalize_options(self): + pass + + def run(self): + from grpc_tools import protoc # pylint: disable=g-import-not-at-top + + root_dir = os.path.dirname(os.path.abspath(__file__)) + + proto_file = os.path.join( + root_dir, "grain", "proto", "execution_summary.proto" + ) + + grpc_protos_include = str( + importlib.resources.files("grpc_tools").joinpath("_proto") + ) + + proto_args = [ + "grpc_tools.protoc", + f"--proto_path={grpc_protos_include}", + f"--proto_path={root_dir}", + f"--python_out={root_dir}", + proto_file, + ] + + if protoc.main(proto_args) != 0: + raise RuntimeError(f"Error compiling proto: {proto_args}") + + class BuildPyCommand(build_py.build_py): + """Run proto generation before build_py.""" + + def run(self): + self.run_command("generate_protos") + super().run() + + from pybind11.setup_helpers import Pybind11Extension # pylint: disable=g-import-not-at-top + + ext_modules = [ + Pybind11Extension( + name="grain._src.python.experimental.index_shuffle.python.index_shuffle_module", + sources=[ + "grain/_src/python/experimental/index_shuffle/python/index_shuffle_module.cc", + "grain/_src/python/experimental/index_shuffle/index_shuffle.cc", + ], + include_dirs=["."], + ) + ] + + setuptools.setup( + distclass=BinaryDistribution, + cmdclass={ + "build_py": BuildPyCommand, + "generate_protos": GenerateProtosCommand, + }, + ext_modules=ext_modules, + setup_requires=["grpcio-tools"], + )