Replace manyLinux compliance check action in TF wheel build rule with…

… macros for Linux platforms.

Example of usage:

```
load(
    "@tsl//:third_party/py/py_manylinux_compliance_test.bzl",
    "verify_manylinux_compliance_test",
)

verify_manylinux_compliance_test(
    name = "manylinux_compliance_test",
    aarch64_compliance_tag = "manylinux_2_17_aarch64",
    test_tags = [
        "mac_excluded",
        "windows_excluded",
    ],
    wheel = ":wheel",
    x86_64_compliance_tag = "manylinux_2_17_x86_64",
)
```

The test target is executed only when specified in Bazel command line. The test passes if `auditwheel show` results have the compliance tag value (depends on the machine type). The test fails otherwise and prints the `auditwheel show` results.

PiperOrigin-RevId: 708024471

opensource_only.files

@@ -99,17 +99,18 @@ third_party/nvtx/LICENSE:
 third_party/protobuf/BUILD:
 third_party/py/BUILD.tpl:
 third_party/py/BUILD:
+third_party/py/manylinux_compliance_test.py:
 third_party/py/ml_dtypes/BUILD:
 third_party/py/ml_dtypes/LICENSE:
 third_party/py/numpy/BUILD:
 third_party/py/py_import.bzl:
+third_party/py/py_manylinux_compliance_test.bzl:
 third_party/py/python_configure.bzl:
 third_party/py/python_init_pip.bzl:
 third_party/py/python_init_repositories.bzl:
 third_party/py/python_init_rules.bzl:
 third_party/py/python_init_toolchains.bzl:
 third_party/py/python_repo.bzl:
-third_party/py/verify_manylinux_compliance.py:
 third_party/pybind11.BUILD:
 third_party/pybind11_bazel/BUILD:
 third_party/python_runtime/BUILD:

third_party/py/BUILD

@@ -53,22 +53,8 @@ config_setting(
     },
 )
 
-# Flag indicating if the target requires manylinux compliance verification.
-bool_flag(
-    name = "verify_manylinux",
-    # TODO(ybaturina): Enable the flag by default when hermetic C++ toolchain is ready.
-    build_setting_default = False,
+filegroup(
+    name = "manylinux_compliance_test",
+    srcs = ["manylinux_compliance_test.py"],
     visibility = ["//visibility:public"],
 )
-
-py_binary(
-    name = "verify_manylinux_compliance",
-    srcs = [
-        "verify_manylinux_compliance.py",
-    ],
-    main = "verify_manylinux_compliance.py",
-    visibility = ["//visibility:public"],
-    deps = [
-        "@pypi_auditwheel//:pkg",
-    ],
-)

third_party/py/verify_manylinux_compliance.py → third_party/py/manylinux_compliance_test.py

@@ -1,40 +1,44 @@
-# Copyright 2024 The Tensorflow Authors.
+# Copyright 2024 The TensorFlow 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
 #
-#     https://www.apache.org/licenses/LICENSE-2.0
+#     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.
-"""Tool to verify wheel manylinux compliance."""
+# ==============================================================================
 
 import argparse
 import io
+import platform
 import re
 import sys
 from auditwheel import main_show
 
 
-def parse_args() -> argparse.Namespace:
+def parse_args():
   """Arguments parser."""
   parser = argparse.ArgumentParser(
-      description="Helper for auditwheel", fromfile_prefix_chars="@"
+      description="Helper for manylinux compliance verification",
+      fromfile_prefix_chars="@",
   )
   parser.add_argument(
-      "--wheel_path", required=True, help="Path of the wheel, mandatory"
+      "--wheel-path", required=True, help="Path of the wheel, mandatory"
   )
   parser.add_argument(
-      "--compliance-tag", help="ManyLinux compliance tag", required=False
+      "--aarch64-compliance-tag",
+      required=True,
+      help="ManyLinux compliance tag for aarch64",
   )
   parser.add_argument(
-      "--auditwheel-show-log-path",
-      help="Path to file with auditwheel show results, mandatory",
+      "--x86_64-compliance-tag",
       required=True,
+      help="ManyLinux compliance tag for x86_64",
   )
   return parser.parse_args()
 
@@ -70,39 +74,37 @@ def get_auditwheel_output(wheel_path: str) -> None:
 def verify_manylinux_compliance(
     auditwheel_log: str,
     compliance_tag: str,
-    auditwheel_show_log_path: str,
 ) -> None:
   """Verify manylinux compliance.
 
   Args:
     auditwheel_log: "auditwheel show" execution results
     compliance_tag: manyLinux compliance tag
-    auditwheel_show_log_path: path to file with auditwheel show results
 
   Raises:
     RuntimeError: if the wheel is not manyLinux compliant.
   """
-  with open(auditwheel_show_log_path, "w") as auditwheel_show_log:
-    auditwheel_show_log.write(auditwheel_log)
-  if not compliance_tag:
-    return
   regex = 'following platform tag: "{}"'.format(compliance_tag)
   if not re.search(regex, auditwheel_log):
     raise RuntimeError(
-        (
-            "The wheel is not compliant with tag {tag}."
-            + " If you want to disable this check, please provide"
-            + " `--@tsl//third_party/py:verify_manylinux=false`."
-            + "\n{result}"
-        ).format(tag=compliance_tag, result=auditwheel_log)
+        ("The wheel is not compliant with the tag {tag}.\n{result}").format(
+            tag=compliance_tag, result=auditwheel_log
+        )
     )
 
 
-if __name__ == "__main__":
-  args = parse_args()
+def test_manylinux_compliance(args):
+  machine_type = platform.uname().machine
+  if machine_type == "x86_64":
+    compliance_tag = args.x86_64_compliance_tag
+  else:
+    compliance_tag = args.aarch64_compliance_tag
   auditwheel_output = get_auditwheel_output(args.wheel_path)
   verify_manylinux_compliance(
       auditwheel_output,
-      args.compliance_tag,
-      args.auditwheel_show_log_path,
+      compliance_tag,
   )
+
+
+if __name__ == "__main__":
+  test_manylinux_compliance(parse_args())

third_party/py/py_import.bzl

@@ -2,11 +2,7 @@
 
 def _unpacked_wheel_impl(ctx):
     output_dir = ctx.actions.declare_directory(ctx.label.name)
-    wheel = None
-    for w in ctx.files.wheel_rule_outputs:
-        if w.basename.endswith(".whl"):
-            wheel = w
-            break
+    wheel = ctx.file.wheel
     script = """
     {zipper} x {wheel} -d {output}
     for wheel_dep in {wheel_deps}; do
@@ -22,7 +18,7 @@ def _unpacked_wheel_impl(ctx):
         ]),
     )
     ctx.actions.run_shell(
-        inputs = ctx.files.wheel_rule_outputs + ctx.files.wheel_deps,
+        inputs = ctx.files.wheel + ctx.files.wheel_deps,
         command = script,
         outputs = [output_dir],
         tools = [ctx.executable.zipper],
@@ -35,7 +31,7 @@ def _unpacked_wheel_impl(ctx):
 _unpacked_wheel = rule(
     implementation = _unpacked_wheel_impl,
     attrs = {
-        "wheel_rule_outputs": attr.label(mandatory = True, allow_files = True),
+        "wheel": attr.label(mandatory = True, allow_single_file = True),
         "zipper": attr.label(
             default = Label("@bazel_tools//tools/zip:zipper"),
             cfg = "exec",
@@ -53,7 +49,7 @@ def py_import(
     unpacked_wheel_name = name + "_unpacked_wheel"
     _unpacked_wheel(
         name = unpacked_wheel_name,
-        wheel_rule_outputs = wheel,
+        wheel = wheel,
         wheel_deps = wheel_deps,
     )
     native.py_library(

third_party/py/py_manylinux_compliance_test.bzl

@@ -0,0 +1,25 @@
+""" Macros for manylinux compliance verification test. """
+
+load("@rules_python//python:py_test.bzl", "py_test")
+
+def verify_manylinux_compliance_test(
+        name,
+        wheel,
+        aarch64_compliance_tag,
+        x86_64_compliance_tag,
+        test_tags = []):
+    py_test(
+        name = name,
+        srcs = [Label("//third_party/py:manylinux_compliance_test")],
+        data = [
+            wheel,
+        ],
+        deps = ["@pypi_auditwheel//:pkg"],
+        args = [
+            "--wheel-path=$(location {})".format(wheel),
+            "--aarch64-compliance-tag={}".format(aarch64_compliance_tag),
+            "--x86_64-compliance-tag={}".format(x86_64_compliance_tag),
+        ],
+        main = "manylinux_compliance_test.py",
+        tags = ["manual"] + test_tags,
+    )