Use hermetic python to build wheel.

This simplifies the build procedure, using only hermetic packages.
That way we do not need to rely on local python installations or
conflicting local package versions.

PiperOrigin-RevId: 731482365

WORKSPACE

@@ -67,7 +67,7 @@ python_init_repositories(
 load("@xla//third_party/py:python_init_toolchains.bzl", "python_init_toolchains")
 python_init_toolchains()
 
-load("@xla//third_party/py:python_init_pip.bzl", "python_init_pip")
+load("//third_party/bazel/python:python_init_pip.bzl", "python_init_pip")
 python_init_pip()
 
 load("@pypi//:requirements.bzl", "install_deps")

build/BUILD.bazel

@@ -11,6 +11,7 @@
 # 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.
+load("@pypi//:requirements.bzl", "requirement")
 load("@python//:defs.bzl", "compile_pip_requirements")
 load("@python_version_repo//:py_version.bzl", "REQUIREMENTS")
 
@@ -37,13 +38,17 @@ compile_pip_requirements(
 # PIP Package
 ###############################################################################
 
-sh_binary(
+py_binary(
     name = "build_pip_package",
-    srcs = ["build_pip_package.sh"],
+    srcs = ["build_pip_package.py"],
     data = [
         "//:LICENSE",
         "//:README.md",
         "//:pyproject.toml",
+    ],
+    deps = [
         "//jax_tpu_embedding/sparsecore",
+        requirement("build"),
+        requirement("setuptools"),
     ],
 )

build/build_pip_package.py

@@ -0,0 +1,46 @@
+# Copyright 2024 The JAX SC Authors.
+#
+# 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.
+"""Build script for pip package."""
+
+from collections.abc import Sequence
+import subprocess
+import sys
+
+from absl import app
+
+
+def main(argv: Sequence[str]) -> None:
+  if len(argv) < 2:
+    raise app.UsageError(f"Usage: \n  {argv[0]} <output path>")
+
+  output_path = argv[1]
+  print("Building wheels in", output_path)
+
+  python_command = sys.executable
+  try:
+    process = subprocess.run(
+        [python_command, "-m", "build", "--outdir", output_path],
+        check=True,
+        capture_output=True,
+        text=True,
+    )
+    print("Build output:")
+    print(process.stdout)
+  except subprocess.CalledProcessError as e:
+    print(f"Error executing build: {e}")
+    print(e.stderr)
+
+
+if __name__ == "__main__":
+  app.run(main)

build/requirements.in

@@ -6,6 +6,10 @@ dm-tree
 # Pre-release of JAX required for SparseCore TPUs.
 jax[tpu] --pre
 
+# Build utilities.
+build
+setuptools
+
 # Testing.
 clu
 einops