Add backend kwargs to xla tests.

PiperOrigin-RevId: 698163185

opensource_only.files

@@ -30,12 +30,14 @@ third_party/py/BUILD:
 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/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/python_runtime/BUILD:
 third_party/repo.bzl:
 third_party/spirv_llvm_translator/spirv_llvm_translator.BUILD:

third_party/py/py_import.bzl

@@ -0,0 +1,73 @@
+""" Macros to unpack a wheel and use its content as a py_library. """
+
+def _unpacked_wheel_impl(ctx):
+    output_dir = ctx.actions.declare_directory(ctx.label.name)
+    libs = []
+    for dep in ctx.attr.cc_deps:
+        linker_inputs = dep[CcInfo].linking_context.linker_inputs.to_list()
+        for linker_input in linker_inputs:
+            if linker_input.libraries and linker_input.libraries[0].dynamic_library:
+                lib = linker_input.libraries[0].dynamic_library
+                libs.append(lib)
+    wheel = None
+    for w in ctx.files.wheel_rule_outputs:
+        if w.basename.endswith(".whl"):
+            wheel = w
+            break
+    script = """
+    {zipper} x {wheel} -d {output}
+    for lib in {libs}; do
+        cp $lib {output}/tensorflow
+    done
+    """.format(
+        zipper = ctx.executable.zipper.path,
+        wheel = wheel.path,
+        output = output_dir.path,
+        libs = " ".join(["'%s'" % lib.path for lib in libs]),
+    )
+    ctx.actions.run_shell(
+        inputs = ctx.files.wheel_rule_outputs + libs,
+        command = script,
+        outputs = [output_dir],
+        tools = [ctx.executable.zipper],
+    )
+
+    return [
+        DefaultInfo(files = depset([output_dir])),
+    ]
+
+_unpacked_wheel = rule(
+    implementation = _unpacked_wheel_impl,
+    attrs = {
+        "wheel_rule_outputs": attr.label(mandatory = True, allow_files = True),
+        "zipper": attr.label(
+            default = Label("@bazel_tools//tools/zip:zipper"),
+            cfg = "exec",
+            executable = True,
+        ),
+        "cc_deps": attr.label_list(providers = [CcInfo]),
+    },
+)
+
+def py_import(name, wheel, deps = [], cc_deps = []):
+    unpacked_wheel_name = name + "_unpacked_wheel"
+    _unpacked_wheel(
+        name = unpacked_wheel_name,
+        wheel_rule_outputs = wheel,
+        cc_deps = cc_deps,
+    )
+    native.py_library(
+        name = name,
+        data = [":" + unpacked_wheel_name],
+        imports = [unpacked_wheel_name],
+        deps = deps,
+        visibility = ["//visibility:public"],
+    )
+
+"""Unpacks the wheel and uses its content as a py_library.
+Args:
+  wheel: wheel file to unpack.
+  deps: dependencies of the py_library.
+  cc_deps: dependencies that will be copied in the folder
+           with the unpacked wheel content.
+"""  # buildifier: disable=no-effect

third_party/py/verify_manylinux_compliance.py

@@ -0,0 +1,108 @@
+# Copyright 2024 The Tensorflow 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
+#
+#     https://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 re
+import sys
+from auditwheel import main_show
+
+
+def parse_args() -> argparse.Namespace:
+  """Arguments parser."""
+  parser = argparse.ArgumentParser(
+      description="Helper for auditwheel", fromfile_prefix_chars="@"
+  )
+  parser.add_argument(
+      "--wheel_path", required=True, help="Path of the wheel, mandatory"
+  )
+  parser.add_argument(
+      "--compliance-tag", help="ManyLinux compliance tag", required=False
+  )
+  parser.add_argument(
+      "--auditwheel-show-log-path",
+      help="Path to file with auditwheel show results, mandatory",
+      required=True,
+  )
+  return parser.parse_args()
+
+
+def get_auditwheel_output(wheel_path: str) -> None:
+  """Run "auditwheel show" on the wheel and return the output.
+
+  Args:
+    wheel_path: path of the wheel file
+
+  Returns:
+    "auditwheel show" output
+  """
+  stringio = io.StringIO()
+  previous_stdout = sys.stdout
+  sys.stdout = stringio
+
+  auditwheel_parser = argparse.ArgumentParser(
+      description="Cross-distro Python wheels."
+  )
+  sub_parsers = auditwheel_parser.add_subparsers(metavar="command", dest="cmd")
+  main_show.configure_parser(sub_parsers)
+  auditwheel_args = argparse.Namespace(
+      WHEEL_FILE=wheel_path,
+      verbose=1,
+  )
+  main_show.execute(args=auditwheel_args, p=auditwheel_parser)
+
+  sys.stdout = previous_stdout
+  return stringio.getvalue()
+
+
+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)
+    )
+
+
+if __name__ == "__main__":
+  args = parse_args()
+  auditwheel_output = get_auditwheel_output(args.wheel_path)
+  verify_manylinux_compliance(
+      auditwheel_output,
+      args.compliance_tag,
+      args.auditwheel_show_log_path,
+  )

xla/tests/BUILD

@@ -567,6 +567,7 @@ xla_test(
 xla_test(
     name = "while_test",
     srcs = ["while_test.cc"],
+    # placeholder for extra args for while_test
     tags = ["test_xla_cpu_thunks"],
     deps = [
         ":client_library_test_base",
@@ -1773,6 +1774,7 @@ xla_test(
     name = "reduce_window_test",
     timeout = "long",
     srcs = [],
+    # placeholder for extra args for reduce_window_test
     shard_count = 40,
     tags = [
         "optonly",
@@ -3367,6 +3369,7 @@ bzl_library(
         ":plugin_bzl",
         "//xla:xla_bzl",
         "//xla/stream_executor:build_defs_bzl",
+        "//xla/tsl:tsl_bzl",
         "@tsl//tsl/platform:build_config_root_bzl",
     ] + tests_build_defs_bzl_deps(),
 )