Add jax_source_package macros and target to generate a source package .tar.gz.

Refactor `jax_wheel` macros, so it outputs a `.whl` file only.

When the macros returns one output object only, it allows all downstream dependencies consume it easily without the need to filter the macros outputs.

The previous implementation design (when `jax_wheel` returned `.tar.gz` and `.whl` files) required one of two options: either create a new target that produces `.whl` only, or to implement filename filtering in the downstream rules. With the new implementation we can just depend on `//:jax_wheel` target that produces the `.whl`.

PiperOrigin-RevId: 738547491

Author: Google-ML-Automation

BUILD.bazel

@@ -15,6 +15,7 @@
 load("@xla//third_party/py:python_wheel.bzl", "collect_data_files", "transitive_py_deps")
 load(
     "//jaxlib:jax.bzl",
+    "jax_source_package",
     "jax_wheel",
 )
 
@@ -67,7 +68,6 @@ py_binary(
 
 jax_wheel(
     name = "jax_wheel",
-    build_wheel_only = False,
     platform_independent = True,
     source_files = [
         ":transitive_py_data",
@@ -82,3 +82,19 @@ jax_wheel(
     wheel_binary = ":build_wheel",
     wheel_name = "jax",
 )
+
+jax_source_package(
+    name = "jax_source_package",
+    source_files = [
+        ":transitive_py_data",
+        ":transitive_py_deps",
+        "//jax:py.typed",
+        "AUTHORS",
+        "LICENSE",
+        "README.md",
+        "pyproject.toml",
+        "setup.py",
+    ],
+    source_package_binary = ":build_wheel",
+    source_package_name = "jax",
+)

build/build.py

@@ -68,6 +68,7 @@
 # rule as the default.
 WHEEL_BUILD_TARGET_DICT_NEW = {
     "jax": "//:jax_wheel",
+    "jax_source_package": "//:jax_source_package",
     "jaxlib": "//jaxlib/tools:jaxlib_wheel",
     "jax-cuda-plugin": "//jaxlib/tools:jax_cuda_plugin_wheel",
     "jax-cuda-pjrt": "//jaxlib/tools:jax_cuda_pjrt_wheel",
@@ -661,6 +662,8 @@ async def main():
       # Append the build target to the Bazel command.
       build_target = wheel_build_targets[wheel]
       wheel_build_command.append(build_target)
+      if args.use_new_wheel_build_rule and wheel == "jax":
+        wheel_build_command.append(wheel_build_targets["jax_source_package"])
 
       if not args.use_new_wheel_build_rule:
         wheel_build_command.append("--")

build_wheel.py

@@ -47,6 +47,20 @@
 parser.add_argument(
     "--srcs", help="source files for the wheel", action="append"
 )
+parser.add_argument(
+    "--build-wheel-only",
+    default=False,
+    help=(
+        "Whether to build the wheel only. Optional."
+    ),
+)
+parser.add_argument(
+    "--build-source-package-only",
+    default=False,
+    help=(
+        "Whether to build the source package only. Optional."
+    ),
+)
 args = parser.parse_args()
 
 
@@ -94,7 +108,8 @@ def prepare_srcs(deps: list[str], srcs_dir: str) -> None:
       args.output_path,
       package_name="jax",
       git_hash=args.jaxlib_git_hash,
-      build_wheel_only=False,
+      build_wheel_only=args.build_wheel_only,
+      build_source_package_only=args.build_source_package_only,
   )
 finally:
   if tmpdir:

jaxlib/jax.bzl

@@ -362,7 +362,7 @@ def _get_full_wheel_name(
         free_threaded_suffix = "t" if py_freethreaded.lower() == "yes" else "",
     )
 
-def _get_source_distribution_name(package_name, wheel_version):
+def _get_source_package_name(package_name, wheel_version):
     return "{package_name}-{wheel_version}.tar.gz".format(
         package_name = package_name,
         wheel_version = wheel_version,
@@ -394,37 +394,47 @@ def _jax_wheel_impl(ctx):
     no_abi = ctx.attr.no_abi
     platform_independent = ctx.attr.platform_independent
     build_wheel_only = ctx.attr.build_wheel_only
+    build_source_package_only = ctx.attr.build_source_package_only
     editable = ctx.attr.editable
     platform_name = ctx.attr.platform_name
+
+    output_dir_path = ""
+    outputs = []
     if editable:
         output_dir = ctx.actions.declare_directory(output_path + "/" + ctx.attr.wheel_name)
-        wheel_dir = output_dir.path
+        output_dir_path = output_dir.path
         outputs = [output_dir]
         args.add("--editable")
     else:
-        wheel_name = _get_full_wheel_name(
-            package_name = ctx.attr.wheel_name,
-            no_abi = no_abi,
-            platform_independent = platform_independent,
-            platform_name = platform_name,
-            cpu_name = cpu,
-            wheel_version = full_wheel_version,
-            py_freethreaded = py_freethreaded,
-        )
-        wheel_file = ctx.actions.declare_file(output_path +
-                                              "/" + wheel_name)
-        wheel_dir = wheel_file.path[:wheel_file.path.rfind("/")]
-        outputs = [wheel_file]
-        if not build_wheel_only:
-            source_distribution_name = _get_source_distribution_name(
+        if build_wheel_only:
+            wheel_name = _get_full_wheel_name(
                 package_name = ctx.attr.wheel_name,
+                no_abi = no_abi,
+                platform_independent = platform_independent,
+                platform_name = platform_name,
+                cpu_name = cpu,
                 wheel_version = full_wheel_version,
+                py_freethreaded = py_freethreaded,
             )
-            source_distribution_file = ctx.actions.declare_file(output_path +
-                                                                "/" + source_distribution_name)
-            outputs.append(source_distribution_file)
-
-    args.add("--output_path", wheel_dir)  # required argument
+            wheel_file = ctx.actions.declare_file(output_path +
+                                                  "/" + wheel_name)
+            output_dir_path = wheel_file.path[:wheel_file.path.rfind("/")]
+            outputs = [wheel_file]
+            if ctx.attr.wheel_name == "jax":
+                args.add("--build-wheel-only", "True")
+        if build_source_package_only:
+            source_package_name = _get_source_package_name(
+                package_name = ctx.attr.wheel_name,
+                wheel_version = full_wheel_version,
+            )
+            source_package_file = ctx.actions.declare_file(output_path +
+                                                           "/" + source_package_name)
+            output_dir_path = source_package_file.path[:source_package_file.path.rfind("/")]
+            outputs = [source_package_file]
+            if ctx.attr.wheel_name == "jax":
+                args.add("--build-source-package-only", "True")
+
+    args.add("--output_path", output_dir_path)  # required argument
     if not platform_independent:
         args.add("--cpu", cpu)
     args.add("--jaxlib_git_hash", git_hash)  # required argument
@@ -472,16 +482,17 @@ _jax_wheel = rule(
         "wheel_name": attr.string(mandatory = True),
         "no_abi": attr.bool(default = False),
         "platform_independent": attr.bool(default = False),
-        "build_wheel_only": attr.bool(default = True),
+        "build_wheel_only": attr.bool(mandatory = True, default = True),
+        "build_source_package_only": attr.bool(mandatory = True, default = False),
         "editable": attr.bool(default = False),
-        "cpu": attr.string(mandatory = True),
-        "platform_name": attr.string(mandatory = True),
+        "cpu": attr.string(),
+        "platform_name": attr.string(),
         "git_hash": attr.label(default = Label("//jaxlib/tools:jaxlib_git_hash")),
         "source_files": attr.label_list(allow_files = True),
         "output_path": attr.label(default = Label("//jaxlib/tools:output_path")),
         "enable_cuda": attr.bool(default = False),
         # A cuda/rocm version is required for gpu wheels; for cpu wheels, it can be an empty string.
-        "platform_version": attr.string(mandatory = True, default = ""),
+        "platform_version": attr.string(),
         "skip_gpu_kernels": attr.bool(default = False),
         "enable_rocm": attr.bool(default = False),
         "include_cuda_libs": attr.label(default = Label("@local_config_cuda//cuda:include_cuda_libs")),
@@ -498,7 +509,6 @@ def jax_wheel(
         wheel_name,
         no_abi = False,
         platform_independent = False,
-        build_wheel_only = True,
         editable = False,
         enable_cuda = False,
         enable_rocm = False,
@@ -509,11 +519,10 @@ def jax_wheel(
     Common artifact attributes are grouped within a single macro.
 
     Args:
-      name: the name of the wheel
+      name: the target name
       wheel_binary: the binary to use to build the wheel
       wheel_name: the name of the wheel
       no_abi: whether to build a wheel without ABI
-      build_wheel_only: whether to build a wheel without source distribution
       editable: whether to build an editable wheel
       platform_independent: whether to build a wheel without platform tag
       enable_cuda: whether to build a cuda wheel
@@ -522,15 +531,16 @@ def jax_wheel(
       source_files: the source files to include in the wheel
 
     Returns:
-      A directory containing the wheel
+      A wheel file or a wheel directory.
     """
     _jax_wheel(
         name = name,
         wheel_binary = wheel_binary,
         wheel_name = wheel_name,
         no_abi = no_abi,
         platform_independent = platform_independent,
-        build_wheel_only = build_wheel_only,
+        build_wheel_only = True,
+        build_source_package_only = False,
         editable = editable,
         enable_cuda = enable_cuda,
         enable_rocm = enable_rocm,
@@ -554,6 +564,34 @@ def jax_wheel(
         source_files = source_files,
     )
 
+def jax_source_package(
+        name,
+        source_package_binary,
+        source_package_name,
+        source_files = []):
+    """Create jax source package.
+
+    Common artifact attributes are grouped within a single macro.
+
+    Args:
+      name: the target name
+      source_package_binary: the binary to use to build the package
+      source_package_name: the name of the source package
+      source_files: the source files to include in the package
+
+    Returns:
+      A jax source package file.
+    """
+    _jax_wheel(
+        name = name,
+        wheel_binary = source_package_binary,
+        wheel_name = source_package_name,
+        build_source_package_only = True,
+        build_wheel_only = False,
+        platform_independent = True,
+        source_files = source_files,
+    )
+
 jax_test_file_visibility = []
 
 jax_export_file_visibility = []

jaxlib/tools/build_utils.py

@@ -65,6 +65,7 @@ def build_wheel(
     package_name: str,
     git_hash: str = "",
     build_wheel_only: bool = True,
+    build_source_package_only: bool = False,
 ) -> None:
   """Builds a wheel in `output_path` using the source tree in `sources_path`."""
   env = dict(os.environ)
@@ -78,7 +79,8 @@ def build_wheel(
     env["USERPROFILE"] = env.get("SYSTEMDRIVE", "C:")
   subprocess.run(
       [sys.executable, "-m", "build", "-n"]
-      + (["-w"] if build_wheel_only else []),
+      + (["-w"] if build_wheel_only else [])
+      + (["-s"] if build_source_package_only else []),
       check=True,
       cwd=sources_path,
       env=env,
@@ -97,10 +99,10 @@ def build_wheel(
     sys.stderr.write("  bazel run //build:requirements.update" +
                      f" --repo_env=HERMETIC_PYTHON_VERSION={py_version}\n\n")
     shutil.copy(wheel, output_path)
-  if not build_wheel_only:
+  if build_source_package_only:
     for dist in glob.glob(os.path.join(sources_path, "dist", "*.tar.gz")):
       output_file = os.path.join(output_path, os.path.basename(dist))
-      sys.stderr.write(f"Output source distribution: {output_file}\n\n")
+      sys.stderr.write(f"Output source package: {output_file}\n\n")
       shutil.copy(dist, output_path)