wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 5 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/tools/pip_package/BUILD
+558
View File
@@ -0,0 +1,558 @@
# Description:
# Tools for building the TensorFlow pip package.
load("@cuda_cudart//:version.bzl", _cudart_version = "VERSION")
load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda")
load("@local_config_syslibs//:build_defs.bzl", "if_not_system_lib")
load("@nvidia_wheel_versions//:versions.bzl", "NVIDIA_WHEEL_VERSIONS")
load(
"@xla//third_party/py:py_import.bzl",
"py_import",
)
load(
"@xla//third_party/py:py_manylinux_compliance_test.bzl",
"verify_manylinux_compliance_test",
)
load("@xla//third_party/py:python_wheel.bzl", "collect_data_files", "transitive_py_deps")
load("@xla//third_party/rules_python/python:defs.bzl", "py_binary", "py_test")
load("@xla//xla/tsl/mkl:build_defs.bzl", "if_enable_mkl", "if_mkl", "if_mkl_ml")
load("//tensorflow:tensorflow.bzl", "if_wheel_dependency", "if_with_tpu_support", "transitive_hdrs")
load("//tensorflow:tf_version.bzl", "TF_SEMANTIC_VERSION_SUFFIX", "TF_VERSION")
load(
"//tensorflow/core/platform:build_config_root.bzl",
"if_pywrap",
"tf_additional_license_deps",
"tf_cuda_tests_tags",
"tf_exec_properties",
)
load("//tensorflow/tools/pip_package/utils:tf_wheel.bzl", "tf_wheel", "tf_wheel_dep")
package(default_visibility = ["//visibility:public"])
# This returns a list of headers of all public header libraries (e.g.,
# framework, lib), and all of the transitive dependencies of those
# public headers. Not all of the headers returned by the filegroup
# are public (e.g., internal headers that are included by public
# headers), but the internal headers need to be packaged in the
# pip_package for the public headers to be properly included.
#
# Public headers are therefore defined by those that are both:
#
# 1) "publicly visible" as defined by bazel
# 2) Have documentation.
#
# This matches the policy of "public" for our python API.
transitive_hdrs(
name = "included_headers",
deps = [
"//tensorflow/compiler/tf2xla:xla_compiled_cpu_function",
"//tensorflow/compiler/mlir:mlir_graph_optimization_pass",
"//tensorflow/core:core_cpu",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/platform:stream_executor",
"//tensorflow/cc/saved_model:loader",
"//tensorflow/cc/saved_model:reader",
"//tensorflow/cc/saved_model:bundle_v2",
"//tensorflow/c/experimental/filesystem:filesystem_interface",
"//tensorflow/c/experimental/next_pluggable_device:c_api_hdrs",
"//tensorflow/c/experimental/pluggable_profiler:pluggable_profiler_hdrs",
"//tensorflow/c/experimental/stream_executor:stream_executor_hdrs",
"//tensorflow/c/experimental/grappler:grappler_hdrs",
"//tensorflow/c:kernels_hdrs",
"//tensorflow/c:ops_hdrs",
"@xla//xla/pjrt/c:pjrt_c_api_hdrs",
"//tensorflow/core/common_runtime/next_pluggable_device/c:plugin_c_api_hdrs",
# TODO(rostam): Revisit these after cc_shared_library everywhere
"//tensorflow/cc:client_session",
"//tensorflow/cc:cc_ops",
# WARNING: None of the C/C++ code under python/ has any API guarantees, and TF team
# reserves the right to change APIs and other header-level interfaces. If your custom
# op uses these headers, it may break when users upgrade their version of tensorflow.
# NOTE(ebrevdo): None of the python symbols are exported by libtensorflow_framework.so.
# Code that relies on these headers should dynamically link to
# _pywrap_tensorflow_internal.so as well.
"//tensorflow/python/grappler:model_analyzer_lib",
"//tensorflow/python/lib/core:py_exception_registry",
"//tensorflow/python/lib/core:pybind11_proto",
"//tensorflow/python/lib/core:pybind11_absl",
"//tensorflow/python/lib/core:pybind11_lib",
"//tensorflow/python/lib/core:pybind11_status",
"//tensorflow/python/lib/core:py_func_lib",
"//tensorflow/python/lib/core:py_seq_tensor",
"//tensorflow/python/lib/core:py_util",
"//tensorflow/python/util:cpp_python_util",
"//tensorflow/python/util:kernel_registry",
"//tensorflow/python/framework:python_op_gen",
"//tensorflow/python/client:tf_session_helper",
] + if_cuda([
"@local_config_cuda//cuda:cuda_headers",
]),
)
filegroup(
name = "xla_cmake",
srcs = ["xla_build/CMakeLists.txt"],
)
filegroup(
name = "licenses",
srcs = [
"//:LICENSE",
"//third_party/icu/data:LICENSE",
"@XNNPACK//:LICENSE",
"@com_google_absl//:LICENSE",
"@com_google_protobuf//:LICENSE",
"@curl//:COPYING",
"@ducc//:LICENSE",
"@flatbuffers//:LICENSE",
"@gemmlowp//:LICENSE",
"@jpegxl//:LICENSE",
"@libjpeg_turbo//:LICENSE.md",
"@llvm-project//llvm:LICENSE.TXT",
"@llvm-project//mlir:LICENSE.TXT",
"@local_config_tensorrt//:LICENSE",
"@ml_dtypes_py//:LICENSE",
"@org_brotli//:LICENSE",
"@png//:LICENSE",
"@ruy//:LICENSE",
"@snappy//:COPYING",
"@stablehlo//:LICENSE",
"@tf_runtime//:LICENSE",
] + select({
"//tensorflow:android": [],
"//tensorflow:ios": [],
"//tensorflow:linux_s390x": [],
"//tensorflow:windows": [],
"//conditions:default": [
"@com_github_googlecloudplatform_google_cloud_cpp//:LICENSE",
],
}) + if_cuda([
"@local_config_nccl//:LICENSE",
]) + if_mkl([
"@xla//third_party/mkl_dnn:LICENSE",
]) + if_enable_mkl(["@xla//xla/tsl/mkl:LICENSE"]) + if_not_system_lib(
"absl_py",
[
"@absl_py//:LICENSE",
],
) + tf_additional_license_deps(),
)
# Add dynamic kernel dso files here.
DYNAMIC_LOADED_KERNELS = [
"//tensorflow/core/kernels:libtfkernel_sobol_op.so",
]
COMMON_PIP_DEPS = [
"//tensorflow/compiler/mlir/stablehlo:stablehlo",
"//tensorflow/compiler/mlir/tensorflow:gen_mlir_passthrough_op_py",
"//tensorflow/core/function/trace_type:serialization_test_proto_py",
"//tensorflow/core/function/trace_type:serialization",
"//tensorflow/dtensor/python/tests:multi_client_test_util",
"//tensorflow/dtensor/python/tests:test_util",
"//tensorflow/lite/python:tflite_convert",
"//tensorflow/lite/tools:visualize",
"//tensorflow/python/autograph/core:test_lib",
"//tensorflow/python/autograph/impl/testing:pybind_for_testing",
"//tensorflow/python/autograph/pyct/testing:basic_definitions",
"//tensorflow/python/autograph/pyct/testing:decorators",
"//tensorflow/python/autograph/pyct/common_transformers:common_transformers",
"//tensorflow/python/autograph/pyct/static_analysis:type_inference",
"//tensorflow/python/autograph/utils:context_managers",
"//tensorflow/python/autograph/utils:tensor_list",
"//tensorflow/python/compiler:compiler",
"//tensorflow/python/ops:cond_v2",
"//tensorflow/python:distributed_framework_test_lib",
"//tensorflow/python/distribute:distribute_test_lib_pip",
"//tensorflow/python/training/experimental:loss_scale",
"//tensorflow/python/training/experimental:loss_scale_optimizer",
"//tensorflow/python/data/benchmarks:benchmark_base",
"//tensorflow/python/data/experimental/kernel_tests/service:multi_process_cluster",
"//tensorflow/python/data/experimental/kernel_tests/service:test_base",
"//tensorflow/python/data/experimental/ops:testing",
"//tensorflow/python/data/experimental/service:server_lib",
"//tensorflow/python/data/kernel_tests:checkpoint_test_base",
"//tensorflow/python/data/kernel_tests:test_base",
"//tensorflow/python/data/kernel_tests:tf_record_test_base",
"//tensorflow/python/ops/ragged:ragged_tensor_test_ops",
"//tensorflow/python/debug:debug_pip",
"//tensorflow/python/distribute:combinations",
"//tensorflow/python/distribute/coordinator:fault_tolerance_test_base",
"//tensorflow/python/distribute/failure_handling:check_preemption_py",
"//tensorflow/python/distribute/failure_handling:failure_handling_lib",
"//tensorflow/python/distribute/failure_handling:failure_handling_util",
"//tensorflow/python/distribute:multi_process_runner",
"//tensorflow/python/eager:eager_pip",
"//tensorflow/python/keras:combinations",
"//tensorflow/python/keras/mixed_precision:test_util",
"//tensorflow/python/keras/utils:dataset_creator",
"//tensorflow/python/kernel_tests/nn_ops:bias_op_base",
"//tensorflow/python/kernel_tests/nn_ops:cudnn_deterministic_base",
"//tensorflow/python/kernel_tests/nn_ops:depthwise_conv_op_base",
"//tensorflow/python/kernel_tests/nn_ops:xent_op_test_base",
"//tensorflow/python/kernel_tests/random:util",
"//tensorflow/python/kernel_tests/signal:test_util",
"//tensorflow/python/kernel_tests/sparse_ops:sparse_xent_op_test_base",
"//tensorflow/python/lib:__init__",
"//tensorflow/python/ops/parallel_for:test_util",
"//tensorflow/python/ops/structured:structured_tensor_dynamic",
"//tensorflow/python/platform:resource_loader",
"//tensorflow/python/profiler:trace",
"//tensorflow/python/saved_model:saved_model",
"//tensorflow/python/summary:__init__",
"//tensorflow/python/summary:plugin_asset",
"//tensorflow/python/summary:summary_iterator",
"//tensorflow/python/summary:summary_py",
"//tensorflow/python/tools:tools_pip",
"//tensorflow/python/tools/api/generator:create_python_api",
"//tensorflow/python/tools/api/generator2/extractor:extractor",
"//tensorflow/python/tools/api/generator2/generator:generator",
"//tensorflow/python/tpu",
"//tensorflow/python/util:deprecated_module",
"//tensorflow/python/util:deprecated_module_new",
"//tensorflow/python/util:example_parser_configuration",
"//tensorflow/python/util:function_utils",
"//tensorflow/python/util:keyword_args",
"//tensorflow/python/util:lock_util",
"//tensorflow/python/util:module_wrapper",
"//tensorflow/python/util:serialization",
"//tensorflow/python/ops:image_grad_test_base",
"//tensorflow/python/framework:memory_checker",
"//tensorflow/python/framework:test_ops",
"//tensorflow/python/ops:while_v2",
"//tensorflow/tools/common:public_api",
"//tensorflow/tools/common:test_module1",
"//tensorflow/tools/common:traverse",
"//tensorflow/tools/docs:tf_doctest_lib",
"@tsl//tsl/profiler/protobuf:trace_events_proto_py",
"//tensorflow/python/distribute:parameter_server_strategy_v2",
"//tensorflow/python/distribute/coordinator:cluster_coordinator",
"//tensorflow/python/distribute/coordinator:remote_eager_lib",
"//tensorflow/python/distribute/coordinator:metric_utils",
"//tensorflow/python/distribute/experimental/rpc:rpc_ops",
"//tensorflow/python/util:pywrap_xla_ops",
"//tensorflow:tensorflow_py",
"//tensorflow/tools/compatibility:tf_upgrade_v2",
] + if_with_tpu_support([
"//tensorflow/python/tools:grpc_tpu_worker",
"//tensorflow/python/tools:grpc_tpu_worker_service",
])
py_binary(
name = "build_pip_package_py",
srcs = [
"build_pip_package.py",
],
main = "build_pip_package.py",
deps = [
"//tensorflow/tools/pip_package/utils:py_utils",
"@pypi//setuptools",
"@pypi//wheel",
],
)
collect_data_files(
name = "cc_deps",
deps = [
#`tensorflow_py` is the main target to collect all required `data` dependencies
"//tensorflow:tensorflow_py",
# utils below are not part of tensorflow_py DAG
"//tensorflow/compiler/mlir/stablehlo:stablehlo_extension",
"//tensorflow/python/autograph/impl/testing:pybind_for_testing",
"//tensorflow/python/framework:memory_checker",
"//tensorflow/python/util:pywrap_xla_ops",
],
)
transitive_py_deps(
name = "py_deps",
deps = COMMON_PIP_DEPS,
)
py_binary(
name = "modify_setup_py_binary",
srcs = [
"modify_setup_py.py",
],
main = "modify_setup_py.py",
deps = ["@xla//third_party/py:setup_py_nvidia_dependencies_util"],
)
_cuda_major_version_split = _cudart_version.split(".")[0] if _cudart_version else ""
cuda_major_version = _cuda_major_version_split if _cuda_major_version_split else "12"
cuda_suffix = "_cu12" if cuda_major_version == "12" else ""
genrule(
name = "setup_py",
srcs = ["setup.py.tpl"],
outs = ["setup.py"],
cmd = """
$(location :modify_setup_py_binary) \\
--template_file $(location setup.py.tpl) \\
--output_file $(OUTS) \\
--nvidia_wheel_versions_data '{nvidia_wheel_versions}' \\
--tf_version "{wheel_version}{wheel_version_suffix}" \\
--cuda_version {cuda_version}
""".format(
cuda_version = cuda_major_version,
nvidia_wheel_versions = NVIDIA_WHEEL_VERSIONS,
wheel_version = TF_VERSION,
wheel_version_suffix = TF_SEMANTIC_VERSION_SUFFIX,
),
tools = [":modify_setup_py_binary"],
)
py_binary(
name = "setup_py_binary",
srcs = [":setup_py"],
data = [
"//:AUTHORS",
"//:LICENSE",
],
main = "setup.py",
)
tf_wheel(
name = "wheel",
# Explicitly list all external repositories that are needed in build_pip_package.py
external_repos = [
"com_google_absl",
"eigen_archive",
"jsoncpp_git",
"com_google_protobuf",
"xla",
"tsl",
"llvm-project",
"cpuinfo",
"FXdiv",
"net_zstd",
"org_brotli",
"pthreadpool",
"riegeli",
"XNNPACK",
"local_config_cuda",
"local_config_tensorrt",
"cuda_cccl",
"cuda_cublas",
"cuda_cudart",
"cuda_cudnn",
"cuda_cufft",
"cuda_cupti",
"cuda_curand",
"cuda_cusolver",
"cuda_cusparse",
"cuda_nvcc",
"cuda_nvml",
"cuda_nvrtc",
"cuda_nvtx",
"cuda_nvjitlink",
],
headers = [
":included_headers",
":licenses",
"//tensorflow/core:protos_all_proto_srcs",
],
platform_name = select({
"@platforms//os:osx": "macosx",
"@platforms//os:macos": "macosx",
"@platforms//os:windows": "win",
"@platforms//os:linux": "linux",
}),
platform_tag = select({
"@platforms//cpu:aarch64": "arm64",
"@platforms//cpu:arm64": "arm64",
"@platforms//cpu:x86_64": "x86_64",
"@platforms//cpu:ppc": "ppc64le",
}),
source_files = [
"MANIFEST.in",
"//tensorflow/tools/pip_package:THIRD_PARTY_NOTICES.txt",
":setup_py",
":cc_deps",
":py_deps",
"//tensorflow:tf_python_api_gen_v2",
] + select({
"//tensorflow:dynamic_loaded_kernels": if_pywrap(
if_false = DYNAMIC_LOADED_KERNELS,
if_true = [],
),
"//conditions:default": [],
}) + select({
"//conditions:default": if_pywrap(
if_false = [
"//tensorflow:tensorflow_cc",
"//tensorflow:tensorflow_framework",
],
if_true = [
"//tensorflow/python:pywrap_tensorflow_binaries",
],
),
"//tensorflow:windows": if_pywrap(
if_false = [
"//tensorflow/python:pywrap_tensorflow_import_lib_file",
],
if_true = [
"//tensorflow/python:pywrap_tensorflow_binaries",
],
),
}) + if_mkl_ml(["@xla//xla/tsl/mkl:intel_binary_blob"]),
xla_aot_compiled = [
"//tensorflow/compiler/tf2xla:xla_compiled_cpu_runtime_srcs",
":xla_cmake",
],
)
genrule(
name = "empty_test",
outs = ["empty_test.py"],
cmd = "echo '' > $@",
)
py_test(
name = "prebuilt_wheel_import_api_packages_test_cpu",
srcs = if_wheel_dependency(
["import_api_packages_test.py"],
[":empty_test"],
),
main = if_wheel_dependency("import_api_packages_test.py", "empty_test.py"),
tags = [
"cpu",
"manual",
],
deps = if_wheel_dependency([":prebuilt_tf_py_import"]) + [
"@pypi//absl_py",
],
)
py_test(
name = "prebuilt_wheel_import_api_packages_test_gpu",
srcs = if_wheel_dependency(
["import_api_packages_test.py"],
[":empty_test"],
),
exec_properties = if_cuda(
tf_exec_properties({"tags": tf_cuda_tests_tags()}),
{},
),
main = if_wheel_dependency("import_api_packages_test.py", "empty_test.py"),
tags = [
"gpu",
"manual",
],
deps = if_wheel_dependency([":prebuilt_tf_py_import"]) + [
"@pypi//absl_py",
],
)
py_test(
name = "import_api_packages_test_cpu",
srcs = ["import_api_packages_test.py"],
main = "import_api_packages_test.py",
tags = [
"cpu",
"manual",
],
deps = [
":tf_py_import",
"@pybind11_abseil//pybind11_abseil:status_casters",
],
)
py_test(
name = "import_api_packages_test_gpu",
srcs = ["import_api_packages_test.py"],
exec_properties = if_cuda(
tf_exec_properties({"tags": tf_cuda_tests_tags()}),
{},
),
main = "import_api_packages_test.py",
tags = [
"gpu",
"manual",
],
deps = [
":tf_py_import",
"@pybind11_abseil//pybind11_abseil:status_casters",
],
)
verify_manylinux_compliance_test(
name = "manylinux_compliance_test",
aarch64_compliance_tag = "manylinux_2_27_aarch64",
ppc64le_compliance_tag = "manylinux_2_17_ppc64le",
test_tags = [
"manual",
],
wheel = ":wheel",
x86_64_compliance_tag = "manylinux_2_27_x86_64",
)
filegroup(
name = "nvidia_wheel_deps",
srcs = if_cuda([
"@pypi_nvidia_cublas{cuda}//:pkg".format(cuda = cuda_suffix),
"@pypi_nvidia_cuda_cupti{cuda}//:pkg".format(cuda = cuda_suffix),
"@pypi_nvidia_cuda_nvcc{cuda}//:pkg".format(cuda = cuda_suffix),
"@pypi_nvidia_cuda_nvrtc{cuda}//:pkg".format(cuda = cuda_suffix),
"@pypi_nvidia_cuda_runtime{cuda}//:pkg".format(cuda = cuda_suffix),
"@pypi_nvidia_cudnn_cu{cuda}//:pkg".format(cuda = cuda_major_version),
"@pypi_nvidia_cufft{cuda}//:pkg".format(cuda = cuda_suffix),
"@pypi_nvidia_curand{cuda}//:pkg".format(cuda = cuda_suffix),
"@pypi_nvidia_cusolver{cuda}//:pkg".format(cuda = cuda_suffix),
"@pypi_nvidia_cusparse{cuda}//:pkg".format(cuda = cuda_suffix),
"@pypi_nvidia_nccl_cu{cuda}//:pkg".format(cuda = cuda_major_version),
"@pypi_nvidia_nvjitlink{cuda}//:pkg".format(cuda = cuda_suffix),
"@pypi_nvidia_nvshmem_cu{cuda}//:pkg".format(cuda = cuda_major_version),
]) + if_cuda(([
"@pypi_nvidia_nvvm{cuda}//:pkg".format(cuda = cuda_suffix),
"@pypi_nvidia_cuda_crt{cuda}//:pkg".format(cuda = cuda_suffix),
] if cuda_major_version == "13" else [])),
)
py_import(
name = "tf_py_import",
wheel = ":wheel",
wheel_deps = if_cuda([":nvidia_wheel_deps"]),
deps = [
"@pypi//absl_py",
"@pypi//astunparse",
"@pypi//flatbuffers",
"@pypi//gast",
"@pypi//ml_dtypes",
"@pypi//numpy",
"@pypi//opt_einsum",
"@pypi//packaging",
"@pypi//protobuf",
"@pypi//requests",
"@pypi//termcolor",
"@pypi//typing_extensions",
"@pypi//wrapt",
],
)
py_import(
name = "prebuilt_tf_py_import",
wheel = if_wheel_dependency(
tf_wheel_dep(),
":wheel",
),
wheel_deps = if_cuda([":nvidia_wheel_deps"]),
deps = [
"@pypi//absl_py",
"@pypi//astunparse",
"@pypi//flatbuffers",
"@pypi//gast",
"@pypi//ml_dtypes",
"@pypi//numpy",
"@pypi//opt_einsum",
"@pypi//packaging",
"@pypi//protobuf",
"@pypi//requests",
"@pypi//termcolor",
"@pypi//typing_extensions",
"@pypi//wrapt",
],
)
tensorflow/tools/pip_package/MANIFEST.in
+15
View File
@@ -0,0 +1,15 @@
include LICENSE
include README
recursive-include * *.py
recursive-include * *.pyd
recursive-include * *.pyi
recursive-include * *.pd
recursive-include * *.so.[0-9]
recursive-include * *.dylib
recursive-include * *.dll
recursive-include * *.lib
recursive-include * *.csv
recursive-include tensorflow *
recursive-exclude tensorflow *.md
recursive-exclude tensorflow/_api/ *
include tensorflow/_api/v2/api_packages.txt
tensorflow/tools/pip_package/README
+1
View File
@@ -0,0 +1 @@
TensorFlow
tensorflow/tools/pip_package/THIRD_PARTY_NOTICES.txt
+9412
View File
File diff suppressed because it is too large Load Diff
tensorflow/tools/pip_package/build_pip_package.py
+507
View File
@@ -0,0 +1,507 @@
# Copyright 2023 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 rearrange files and build the wheel.
In a nutshell this script does:
1) Takes lists of paths to .h/.py/.so/etc files.
2) Creates a temporary directory.
3) Copies files from #1 to #2 with some exceptions and corrections.
4) A wheel is created from the files in the temp directory.
Most of the corrections are related to tsl/xla vendoring:
These files used to be a part of source code but were moved to an external repo.
To not break the TF API, we pretend that it's still part of the it.
"""
import argparse
import json
import os
import shutil
import subprocess
import sys
import tempfile
from tensorflow.tools.pip_package.utils.utils import copy_file
from tensorflow.tools.pip_package.utils.utils import create_init_files
from tensorflow.tools.pip_package.utils.utils import is_macos
from tensorflow.tools.pip_package.utils.utils import is_windows
from tensorflow.tools.pip_package.utils.utils import replace_inplace
def get_repo_path(repo_name: str) -> str:
"""Returns the path to an external repository.
The canonical repository name is provided by the tf_wheel rule in
tensorflow/tools/pip_package/utils/tf_wheel.bzl.
Args:
repo_name: The apparent repository name.
Returns:
The path to the external repository.
"""
env_key = (
repo_name.upper().replace("/", "_").replace("-", "_").replace(":", "_")
)
canonical_name = os.environ.get(f"{env_key}_CANONICAL_REPO_NAME", repo_name)
return f"external/{canonical_name}/"
def parse_args() -> argparse.Namespace:
"""Arguments parser."""
parser = argparse.ArgumentParser(
description="Helper for building pip package", fromfile_prefix_chars="@"
)
parser.add_argument(
"--output-name",
required=True,
help="Output file for the wheel, mandatory",
)
parser.add_argument(
"--project-name",
required=True,
help="Project name to be passed to setup.py",
)
parser.add_argument(
"--platform",
required=True,
help="Platform name to be passed to setup.py",
)
parser.add_argument(
"--headers", help="header files for the wheel", action="append"
)
parser.add_argument(
"--srcs", help="source files for the wheel", action="append"
)
parser.add_argument("--dests", help="", action="append", default=[])
parser.add_argument(
"--xla_aot", help="xla aot compiled sources", action="append"
)
parser.add_argument("--version", help="TF version")
parser.add_argument("--collab", help="True if collaborator build")
return parser.parse_args()
def _copy_cuda_tree(
srcs_dir: str, repo_relative_path: str, destination_relative_path: str
) -> None:
source_path = os.path.join(srcs_dir, repo_relative_path)
if os.path.exists(source_path):
destination_path = os.path.join(srcs_dir, destination_relative_path)
shutil.copytree(
source_path,
destination_path,
dirs_exist_ok=True,
)
def prepare_headers(headers: list[str], srcs_dir: str) -> None:
"""Copy and rearrange header files in the target directory.
Filter out headers by their path and replace paths for some of them.
Args:
headers: a list of paths to header files.
srcs_dir: target directory where headers are copied to.
"""
path_to_exclude = [
# Bazel generates _virtual_includes directories for cc_library targets,
# which are not needed in the pip package.
"cuda_cccl/_virtual_includes",
"cuda_cublas/_virtual_includes",
"cuda_cudart/_virtual_includes",
"cuda_cudnn/_virtual_includes",
"cuda_cufft/_virtual_includes",
"cuda_cupti/_virtual_includes",
"cuda_curand/_virtual_includes",
"cuda_cusolver/_virtual_includes",
"cuda_cusparse/_virtual_includes",
"cuda_nccl/_virtual_includes",
"cuda_nvcc/_virtual_includes",
"cuda_nvjitlink/_virtual_includes",
"cuda_nvml/_virtual_includes",
"cuda_nvrtc/_virtual_includes",
"cuda_nvrtc_builtins/_virtual_includes",
"cuda_nvtx/_virtual_includes",
get_repo_path("pypi"),
get_repo_path("jsoncpp_git") + "src",
get_repo_path("local_config_cuda") + "cuda/_virtual_includes",
get_repo_path("local_config_tensorrt"),
"python_x86_64",
"python_aarch64",
get_repo_path("llvm-project") + "llvm/",
get_repo_path("cpuinfo"),
get_repo_path("FXdiv"),
get_repo_path("net_zstd"),
get_repo_path("org_brotli") + "c",
get_repo_path("org_brotli") + "_virtual_includes",
get_repo_path("pthreadpool"),
get_repo_path("riegeli") + "riegeli",
get_repo_path("XNNPACK") + "src/",
]
path_to_replace = {
get_repo_path("com_google_absl"): "",
get_repo_path("eigen_archive"): "",
get_repo_path("jsoncpp_git"): "",
get_repo_path("com_google_protobuf") + "src/": "",
get_repo_path("xla"): "tensorflow/compiler",
get_repo_path("tsl"): "tensorflow",
}
for file in headers:
if file.endswith("cc.inc"):
continue
if any(i in file for i in path_to_exclude):
continue
for path, val in path_to_replace.items():
if path in file:
copy_file(file, os.path.join(srcs_dir, val), path)
break
else:
copy_file(file, srcs_dir)
cuda_src_path = os.path.join(
srcs_dir, get_repo_path("local_config_cuda") + "cuda"
)
if os.path.exists(cuda_src_path):
shutil.copytree(
cuda_src_path,
os.path.join(srcs_dir, "third_party/gpus"),
)
_copy_cuda_tree(srcs_dir, get_repo_path("cuda_cccl"), "third_party/gpus/cuda")
_copy_cuda_tree(
srcs_dir, get_repo_path("cuda_cublas"), "third_party/gpus/cuda"
)
_copy_cuda_tree(
srcs_dir, get_repo_path("cuda_cudart"), "third_party/gpus/cuda"
)
_copy_cuda_tree(
srcs_dir, get_repo_path("cuda_cudnn"), "third_party/gpus/cudnn"
)
_copy_cuda_tree(
srcs_dir, get_repo_path("cuda_cufft"), "third_party/gpus/cuda"
)
_copy_cuda_tree(
srcs_dir,
get_repo_path("cuda_cupti"),
"third_party/gpus/cuda/extras/CUPTI",
)
_copy_cuda_tree(
srcs_dir, get_repo_path("cuda_curand"), "third_party/gpus/cuda"
)
_copy_cuda_tree(
srcs_dir, get_repo_path("cuda_cusolver"), "third_party/gpus/cuda"
)
_copy_cuda_tree(
srcs_dir, get_repo_path("cuda_cusparse"), "third_party/gpus/cuda"
)
_copy_cuda_tree(srcs_dir, get_repo_path("cuda_nvcc"), "third_party/gpus/cuda")
_copy_cuda_tree(
srcs_dir, get_repo_path("cuda_nvjitlink"), "third_party/gpus/cuda"
)
_copy_cuda_tree(
srcs_dir, get_repo_path("cuda_nvml"), "third_party/gpus/cuda/nvml"
)
_copy_cuda_tree(
srcs_dir, get_repo_path("cuda_nvrtc"), "third_party/gpus/cuda"
)
_copy_cuda_tree(srcs_dir, get_repo_path("cuda_nvtx"), "third_party/gpus/cuda")
shutil.copytree(
os.path.join(srcs_dir, "tensorflow/compiler/xla"),
os.path.join(srcs_dir, "xla"),
)
shutil.copytree(
os.path.join(srcs_dir, "tensorflow/tsl"), os.path.join(srcs_dir, "tsl")
)
def prepare_srcs(
deps: list[str], deps_destinations: list[str], srcs_dir: str
) -> None:
"""Rearrange source files in target the target directory.
Exclude `external` files and move vendored xla/tsl files accordingly.
Args:
deps: a list of paths to files.
deps_destinations: a list of json files with mapping of deps to their
destinations for deps whose original path and path inside the wheel are
different.
srcs_dir: target directory where files are copied to.
"""
path_to_replace = {
get_repo_path("xla"): "tensorflow/compiler",
get_repo_path("tsl"): "tensorflow",
}
deps_mapping_dict = {}
for deps_destination in deps_destinations:
with open(deps_destination, "r") as deps_destination_file:
deps_mapping_dict.update(json.load(deps_destination_file))
for file in deps:
for path, val in path_to_replace.items():
if path in file:
copy_file(file, os.path.join(srcs_dir, val), path)
break
else:
# exclude external py files
if "external" not in file:
if file in deps_mapping_dict:
dest = deps_mapping_dict[file]
if dest:
copy_file(file, srcs_dir, None, dest)
else:
copy_file(file, srcs_dir, None, None)
def prepare_aot(aot: list[str], srcs_dir: str) -> None:
"""Rearrange xla_aot files in target the target directory.
Args:
aot: a list of paths to files that should be in xla_aot directory.
srcs_dir: target directory where files are copied to.
"""
for file in aot:
if get_repo_path("tsl") in file:
copy_file(file, srcs_dir, get_repo_path("tsl"))
elif get_repo_path("xla") in file:
copy_file(file, srcs_dir, get_repo_path("xla"))
else:
copy_file(file, srcs_dir)
shutil.move(
os.path.join(
srcs_dir, "tensorflow/tools/pip_package/xla_build/CMakeLists.txt"
),
os.path.join(srcs_dir, "CMakeLists.txt"),
)
def prepare_wheel_srcs(
headers: list[str],
srcs: list[str],
dests: list[str],
aot: list[str],
srcs_dir: str,
version: str,
) -> None:
"""Rearrange source and header files.
Args:
headers: a list of paths to header files.
srcs: a list of paths to the rest of files.
dests: a list of paths to files with srcs files destinations.
aot: a list of paths to files that should be in xla_aot directory.
srcs_dir: directory to copy files to.
version: tensorflow version.
"""
prepare_headers(headers, os.path.join(srcs_dir, "tensorflow/include"))
prepare_srcs(srcs, dests, srcs_dir)
prepare_aot(aot, os.path.join(srcs_dir, "tensorflow/xla_aot_runtime_src"))
# Every directory that contains a .py file gets an empty __init__.py file.
create_init_files(os.path.join(srcs_dir, "tensorflow"))
# move MANIFEST and THIRD_PARTY_NOTICES to the root
shutil.move(
os.path.join(srcs_dir, "tensorflow/tools/pip_package/MANIFEST.in"),
os.path.join(srcs_dir, "MANIFEST.in"),
)
shutil.move(
os.path.join(
srcs_dir, "tensorflow/tools/pip_package/THIRD_PARTY_NOTICES.txt"
),
os.path.join(srcs_dir, "tensorflow/THIRD_PARTY_NOTICES.txt"),
)
update_xla_tsl_imports(os.path.join(srcs_dir, "tensorflow"))
# Means the wheel is built with pywrap rules
if dests:
return
if not is_windows():
rename_libtensorflow(os.path.join(srcs_dir, "tensorflow"), version)
if not is_macos() and not is_windows():
patch_so(srcs_dir)
def update_xla_tsl_imports(srcs_dir: str) -> None:
"""Workaround for TSL and XLA vendoring."""
replace_inplace(srcs_dir, "from tsl", "from tensorflow.tsl")
replace_inplace(srcs_dir, "from xla", "from tensorflow.compiler.xla")
def patch_so(srcs_dir: str) -> None:
"""Patch .so files.
We must patch some of .so files otherwise auditwheel will fail.
Args:
srcs_dir: target directory with .so files to patch.
"""
to_patch = {
"tensorflow/python/_pywrap_tensorflow_internal.so": (
"$ORIGIN/../../tensorflow/compiler/xla/tsl/python/lib/core"
),
(
"tensorflow/compiler/mlir/quantization/tensorflow/python/"
"pywrap_function_lib.so"
): "$ORIGIN/../../../../../python",
(
"tensorflow/compiler/mlir/quantization/tensorflow/python/"
"pywrap_quantize_model.so"
): "$ORIGIN/../../../../../python",
(
"tensorflow/compiler/mlir/tensorflow_to_stablehlo/python/"
"pywrap_tensorflow_to_stablehlo.so"
): "$ORIGIN/../../../../python",
(
"tensorflow/compiler/mlir/lite/python/_pywrap_converter_api.so"
): "$ORIGIN/../../../../python",
}
for file, path in to_patch.items():
file_path = os.path.join(srcs_dir, file)
if not os.path.exists(file_path):
continue
rpath = (
subprocess.check_output(["patchelf", "--print-rpath", file_path])
.decode()
.strip()
)
new_rpath = rpath + ":" + path
subprocess.run(
["patchelf", "--set-rpath", new_rpath, file_path],
check=True,
)
subprocess.run(
["patchelf", "--shrink-rpath", file_path],
check=True,
)
def rename_libtensorflow(srcs_dir: str, version: str):
"""Update libtensorflow_cc file name.
Bazel sets full TF version in name but libtensorflow_cc must contain only
major. Update accordingly to the platform:
e.g. libtensorflow_cc.so.2.15.0 -> libtensorflow_cc.2
Args:
srcs_dir: target directory with files.
version: Major version to be set.
"""
major_version = version.split(".")[0]
if is_macos():
shutil.move(
os.path.join(srcs_dir, "libtensorflow_cc.{}.dylib".format(version)),
os.path.join(
srcs_dir, "libtensorflow_cc.{}.dylib".format(major_version)
),
)
shutil.move(
os.path.join(
srcs_dir, "libtensorflow_framework.{}.dylib".format(version)
),
os.path.join(
srcs_dir, "libtensorflow_framework.{}.dylib".format(major_version)
),
)
else:
shutil.move(
os.path.join(srcs_dir, "libtensorflow_cc.so.{}".format(version)),
os.path.join(srcs_dir, "libtensorflow_cc.so.{}".format(major_version)),
)
shutil.move(
os.path.join(srcs_dir, "libtensorflow_framework.so.{}".format(version)),
os.path.join(
srcs_dir, "libtensorflow_framework.so.{}".format(major_version)
),
)
def build_wheel(
dir_path: str,
cwd: str,
project_name: str,
platform: str,
collab: str = False,
) -> None:
"""Build the wheel in the target directory.
Args:
dir_path: directory where the wheel will be stored
cwd: path to directory with wheel source files
project_name: name to pass to setup.py.
platform: platform name to pass to setup.py.
collab: defines if this is a collab build
"""
env = os.environ.copy()
if is_windows():
# HOMEPATH is not set by bazel but it's required by setuptools.
env["HOMEPATH"] = "C:"
# project_name is needed by setup.py.
env["project_name"] = project_name
if collab == "True":
env["collaborator_build"] = True
# Note: (Required for rules_python >= 1.7.0)
# Modern rules_python no longer exports PYTHONPATH to subprocesses by default
# for stricter hermeticity (see release 1.7.0). We must explicitly propagate
# the current sys.path so spawned child processes can resolve external
# dependencies.
env["PYTHONPATH"] = os.pathsep.join([os.path.abspath(p) for p in sys.path])
subprocess.run(
[
sys.executable,
"tensorflow/tools/pip_package/setup.py",
"bdist_wheel",
f"--dist-dir={dir_path}",
f"--plat-name={platform}",
],
check=True,
cwd=cwd,
env=env,
)
if __name__ == "__main__":
args = parse_args()
temp_dir = tempfile.TemporaryDirectory(prefix="tensorflow_wheel")
temp_dir_path = temp_dir.name
try:
prepare_wheel_srcs(
args.headers,
args.srcs,
args.dests,
args.xla_aot,
temp_dir_path,
args.version,
)
build_wheel(
os.path.join(os.getcwd(), args.output_name),
temp_dir_path,
args.project_name,
args.platform,
args.collab,
)
finally:
temp_dir.cleanup()
tensorflow/tools/pip_package/check_load_py_test.py
+88
View File
@@ -0,0 +1,88 @@
# Copyright 2017 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
#
# 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.
# ==============================================================================
"""Tests to check that py_test are properly loaded in BUILD files."""
import os
import subprocess
os.chdir(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../..')))
def check_output_despite_error(args):
"""Get output of args from command line, even if there are errors.
Args:
args: a list of command line args.
Returns:
output as string.
"""
try:
output = subprocess.check_output(args, shell=True, stderr=subprocess.STDOUT)
except subprocess.CalledProcessError as e:
output = e.output
return output.strip()
def main():
# Get all py_test target, note bazel query result will also include
# cuda_py_test etc.
try:
targets = subprocess.check_output([
'bazel', 'query',
'kind(py_test, //tensorflow/contrib/... + '
'//tensorflow/python/... - '
'//tensorflow/contrib/tensorboard/...)']).strip()
except subprocess.CalledProcessError as e:
targets = e.output
targets = targets.decode("utf-8") if isinstance(targets, bytes) else targets
# Only keep py_test targets, and filter out targets with 'no_pip' tag.
valid_targets = []
for target in targets.split('\n'):
kind = check_output_despite_error(['buildozer', 'print kind', target])
if kind == 'py_test':
tags = check_output_despite_error(['buildozer', 'print tags', target])
if 'no_pip' not in tags:
valid_targets.append(target)
# Get all BUILD files for all valid targets.
build_files = set()
for target in valid_targets:
build_files.add(os.path.join(target[2:].split(':')[0], 'BUILD'))
# Check if BUILD files load py_test.
files_missing_load = []
for build_file in build_files:
updated_build_file = subprocess.check_output(
['buildozer', '-stdout', 'new_load //tensorflow:tensorflow.bzl py_test',
build_file])
with open(build_file, 'r') as f:
if f.read() != updated_build_file:
files_missing_load.append(build_file)
if files_missing_load:
raise RuntimeError('The following files are missing %s:\n %s' % (
'load("//tensorflow:tensorflow.bzl", "py_test").\nThis load statement'
' is needed because otherwise pip tests will try to use their '
'dependencies, which are not visible to them.',
'\n'.join(files_missing_load)))
else:
print('TEST PASSED.')
if __name__ == '__main__':
main()
tensorflow/tools/pip_package/import_api_packages_test.py
+107
View File
@@ -0,0 +1,107 @@
# Copyright 2023 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
#
# 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.
# ==============================================================================
"""Import API packages test.
This Python test verifies whether the API v2 packages can be imported from the
current build. It utilizes the `_api/v2/api_packages.txt` list of packages from
the local wheel file specified in the `requirements_lock_<python_version>.txt`.
Packages are imported one by one in alphabetical order during runtime.
The test doesn't identify package's order-dependent issues; for instance,
importing "tf.foo" followed by "tf.bar" won't reveal that "tf.bar" depends on
"tf.foo" being imported first.
The `_api/v2/api_packages.txt` file is generated during the TensorFlow API v2
init files creation process and is subsequently stored in the wheel file after
the build. It also contains a few paths that cannot be directly imported. These
paths point to attributes or sub-modules within a module's namespace, but they
don't correspond to an actual file or directory on the filesystem. The list of
such paths is stored in the packages_for_skip variable and will be skipped
during the test.
"""
import logging
import os
import unittest
try:
import importlib.resources as pkg_resources # pylint: disable=g-import-not-at-top
except ImportError:
import importlib_resources as pkg_resources # pylint: disable=g-import-not-at-top
logging.basicConfig(level=logging.INFO)
class ImportApiPackagesTest(unittest.TestCase):
"""ImportApiPackagesTest class. See description at the top of the file."""
def setUp(self):
def _get_api_packages_v2():
api_packages_path = os.path.join("_api", "v2", "api_packages.txt")
logging.info("Load api packages file: %s", api_packages_path)
return set(
pkg_resources.files("tensorflow")
.joinpath(api_packages_path)
.read_text()
.splitlines()
)
super().setUp()
self.api_packages_v2 = _get_api_packages_v2()
# Some paths in the api_packages_path file cannot be directly imported.
# These paths may point to attributes or sub-modules within a module's
# namespace, but they don't correspond to an actual file
# or directory on the filesystem.
self.packages_for_skip = [
"tensorflow.distribute.cluster_resolver",
"tensorflow.lite.experimental.authoring",
"tensorflow.distribute.experimental.coordinator",
"tensorflow.summary.experimental",
"tensorflow.distribute.coordinator",
"tensorflow.distribute.experimental.partitioners",
]
def test_import_runtime(self):
"""Try to import all packages from api packages file one by one."""
version = "v2"
failed_packages = []
logging.info("Try to import packages at runtime...")
for package_name in self.api_packages_v2:
# Convert package name to the short version:
# tensorflow._api.v2.distribute.experimental to
# tensorflow.distribute.experimental
short_package_name = package_name.replace(f"_api.{version}.", "")
# skip non-importable paths
if short_package_name not in self.packages_for_skip:
try:
__import__(short_package_name)
except ImportError:
logging.exception("error importing %s", short_package_name)
failed_packages.append(package_name)
if failed_packages:
self.fail(
"Failed to import"
f" {len(failed_packages)}/{len(self.api_packages_v2)} packages"
f" {version}:\n{failed_packages}"
)
logging.info("Import of packages was successful.")
if __name__ == "__main__":
unittest.main()
tensorflow/tools/pip_package/modify_setup_py.py
+97
View File
@@ -0,0 +1,97 @@
# Copyright 2025 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
#
# 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..
# ==============================================================================
"""Modify setup.py with TensorFlow and NVIDIA wheel versions."""
import argparse
import pathlib
import third_party.py.setup_py_nvidia_dependencies_util as util
def _update_setup_with_tf_and_nvidia_wheel_versions(
template_path: pathlib.Path,
output_path: pathlib.Path,
nvidia_wheel_versions_data: str,
tf_version: str,
cuda_version: str,
):
"""Updates a setup.py template with TensorFlow and NVIDIA wheel versions.
This function reads a setup.py template file, replaces placeholder versions
for TensorFlow and various NVIDIA-related wheels based on the provided
data, and writes the result to an output file.
Args:
template_path: Path to the input setup.py.tpl template file.
output_path: Path where the modified setup.py file will be written.
nvidia_wheel_versions_data: A string containing NVIDIA wheel version data,
with each line in the format "wheel_name version_spec".
tf_version: The version string for the TensorFlow package.
cuda_version: The CUDA version string.
"""
with open(template_path) as f:
content = f.read()
content = content.replace("_VERSION = '0.0.0'", f"_VERSION = '{tf_version}'")
content = util.get_setup_py_content_with_nvidia_wheel_versions(
content, cuda_version, nvidia_wheel_versions_data
)
with open(output_path, "w") as f:
f.write(content)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--template_file",
type=pathlib.Path,
required=True,
help="Path to the setup.py.tpl template file",
)
parser.add_argument(
"--output_file",
type=pathlib.Path,
required=True,
help="Path to write the generated setup.py file",
)
parser.add_argument(
"--nvidia_wheel_versions_data",
default=None,
required=True,
help="NVIDIA wheel versions data",
)
parser.add_argument(
"--cuda_version",
type=str,
required=True,
help="The CUDA version string",
default="12",
)
parser.add_argument(
"--tf_version",
type=str,
required=True,
help="The TensorFlow package version string",
)
args = parser.parse_args()
_update_setup_with_tf_and_nvidia_wheel_versions(
args.template_file,
args.output_file,
args.nvidia_wheel_versions_data,
args.tf_version,
args.cuda_version,
)
tensorflow/tools/pip_package/pip_smoke_test.py
+191
View File
@@ -0,0 +1,191 @@
# Copyright 2015 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
#
# 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.
# ==============================================================================
"""This pip smoke test verifies dependency files exist in the pip package.
This script runs bazel queries to see what python files are required by the
tests and ensures they are in the pip package superset.
"""
import os
import subprocess
os.chdir(os.path.abspath(os.path.join(os.path.dirname(__file__), "../../..")))
PIP_PACKAGE_QUERY_EXPRESSION = (
"deps(//tensorflow/tools/pip_package:build_pip_package)")
# List of file paths containing BUILD files that should not be included for the
# pip smoke test.
BUILD_DENYLIST = [
"tensorflow/lite",
"tensorflow/compiler/mlir/lite",
"tensorflow/compiler/mlir/tfrt",
"tensorflow/core/runtime_fallback",
"tensorflow/core/tfrt",
"tensorflow/python/kernel_tests/signal",
"tensorflow/examples",
"tensorflow/tools/android",
"tensorflow/tools/toolchains",
"tensorflow/python/autograph/tests",
"tensorflow/python/eager/benchmarks",
]
def GetBuild(dir_base):
"""Get the list of BUILD file all targets recursively startind at dir_base."""
items = []
for root, _, files in os.walk(dir_base):
for name in files:
if (name == "BUILD" and not any(x in root for x in BUILD_DENYLIST)):
items.append("//" + root + ":all")
return items
def BuildPyTestDependencies():
python_targets = GetBuild("tensorflow/python")
tensorflow_targets = GetBuild("tensorflow")
# Build list of test targets,
# python - attr(manual|pno_pip)
targets = " + ".join(python_targets)
targets += ' - attr(tags, "manual|no_pip", %s)' % " + ".join(
tensorflow_targets)
query_kind = "kind(py_test, %s)" % targets
# Skip benchmarks etc.
query_filter = 'filter("^((?!benchmark).)*$", %s)' % query_kind
# Get the dependencies
query_deps = "deps(%s, 1)" % query_filter
return python_targets, query_deps
PYTHON_TARGETS, PY_TEST_QUERY_EXPRESSION = BuildPyTestDependencies()
# TODO(amitpatankar): Clean up denylist.
# List of dependencies that should not included in the pip package.
DEPENDENCY_DENYLIST = [
"//tensorflow/cc/saved_model:saved_model_test_files",
"//tensorflow/cc/saved_model:saved_model_half_plus_two",
"//tensorflow:no_tensorflow_py_deps",
"//tensorflow/tools/pip_package:win_pip_package_marker",
"//tensorflow/core:image_testdata",
"//tensorflow/core/kernels/cloud:bigquery_reader_ops",
"//tensorflow/python:extra_py_tests_deps",
"//tensorflow/python:mixed_precision",
"//tensorflow/python:tf_optimizer",
"//tensorflow/python/framework:test_file_system.so",
"//tensorflow/python/debug:grpc_tensorflow_server.par",
"//tensorflow/python/feature_column:vocabulary_testdata",
"//tensorflow/python/util:nest_test_main_lib",
# lite
"//tensorflow/lite/experimental/examples/lstm:rnn_cell",
"//tensorflow/lite/experimental/examples/lstm:rnn_cell.py",
"//tensorflow/lite/experimental/examples/lstm:unidirectional_sequence_lstm_test", # pylint:disable=line-too-long
"//tensorflow/lite/experimental/examples/lstm:unidirectional_sequence_lstm_test.py", # pylint:disable=line-too-long
"//tensorflow/lite/python:interpreter",
"//tensorflow/lite/python:interpreter_test",
"//tensorflow/lite/python:interpreter.py",
"//tensorflow/lite/python:interpreter_test.py",
]
def main():
"""This script runs the pip smoke test.
Raises:
RuntimeError: If any dependencies for py_tests exist in subSet
Prerequisites:
1. Bazel is installed.
2. Running in github repo of tensorflow.
3. Configure has been run.
"""
# pip_package_dependencies_list is the list of included files in pip packages
pip_package_dependencies = subprocess.check_output([
"bazel", "cquery", "--experimental_cc_shared_library",
PIP_PACKAGE_QUERY_EXPRESSION
])
if isinstance(pip_package_dependencies, bytes):
pip_package_dependencies = pip_package_dependencies.decode("utf-8")
pip_package_dependencies_list = pip_package_dependencies.strip().split("\n")
pip_package_dependencies_list = [
x.split()[0] for x in pip_package_dependencies_list
]
print("Pip package superset size: %d" % len(pip_package_dependencies_list))
# tf_py_test_dependencies is the list of dependencies for all python
# tests in tensorflow
tf_py_test_dependencies = subprocess.check_output([
"bazel", "cquery", "--experimental_cc_shared_library",
PY_TEST_QUERY_EXPRESSION
])
if isinstance(tf_py_test_dependencies, bytes):
tf_py_test_dependencies = tf_py_test_dependencies.decode("utf-8")
tf_py_test_dependencies_list = tf_py_test_dependencies.strip().split("\n")
tf_py_test_dependencies_list = [
x.split()[0] for x in tf_py_test_dependencies.strip().split("\n")
]
print("Pytest dependency subset size: %d" % len(tf_py_test_dependencies_list))
missing_dependencies = []
# File extensions and endings to ignore
ignore_extensions = [
"_test", "_test.py", "_test_cpu", "_test_cpu.py", "_test_gpu",
"_test_gpu.py", "_test_lib"
]
ignored_files_count = 0
denylisted_dependencies_count = len(DEPENDENCY_DENYLIST)
# Compare dependencies
for dependency in tf_py_test_dependencies_list:
if dependency and dependency.startswith("//tensorflow"):
ignore = False
# Ignore extensions
if any(dependency.endswith(ext) for ext in ignore_extensions):
ignore = True
ignored_files_count += 1
# Check if the dependency is in the pip package, the dependency denylist,
# or should be ignored because of its file extension.
if not (ignore or dependency in pip_package_dependencies_list or
dependency in DEPENDENCY_DENYLIST):
missing_dependencies.append(dependency)
print("Ignored files count: %d" % ignored_files_count)
print("Denylisted dependencies count: %d" % denylisted_dependencies_count)
if missing_dependencies:
print("Missing the following dependencies from pip_packages:")
for missing_dependency in missing_dependencies:
print("\nMissing dependency: %s " % missing_dependency)
print("Affected Tests:")
rdep_query = ("rdeps(kind(py_test, %s), %s)" %
(" + ".join(PYTHON_TARGETS), missing_dependency))
affected_tests = subprocess.check_output(
["bazel", "cquery", "--experimental_cc_shared_library", rdep_query])
affected_tests_list = affected_tests.split("\n")[:-2]
print("\n".join(affected_tests_list))
raise RuntimeError("""
One or more added test dependencies are not in the pip package.
If these test dependencies need to be in TensorFlow pip package, please add them to //tensorflow/tools/pip_package/BUILD.
Else add no_pip tag to the test.""")
else:
print("TEST PASSED")
if __name__ == "__main__":
main()
tensorflow/tools/pip_package/redundant_tensorflow_gpu/README.md
+30
View File
@@ -0,0 +1,30 @@
tensorflow-gpu has been removed. Please install tensorflow instead.
The tensorflow package supports GPU accelerated operations via Nvidia CUDA.
## Removal Information
tensorflow and tensorflow-gpu have been the same package since TensorFlow
2.1, released in September 2019. Although the checksums differ due to metadata,
they were built in the same way and both provide GPU support via Nvidia CUDA.
As of January 2023, tensorflow-gpu has been removed and has been replaced with
this new, empty package that generates an error upon installation.
All existing versions of tensorflow-gpu are still available, but the
TensorFlow team has stopped releasing any new tensorflow-gpu packages, and
will not release any patches for existing tensorflow-gpu versions.
## About this package
This simple package raises a warning if setup.py is executed as part of a
package installation. This intentionally prevents users from installing
the package.
To build and upload this package's source distribution (sdist) to testpypi:
```
$ vim setup.cfg # update the version number and package name
$ python3 -m pip install --user twine
$ python3 setup.py sdist
$ twine upload --repository testpypi dist/*
$ pip3 install the_name_of_your_test_package -i https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple
```
tensorflow/tools/pip_package/redundant_tensorflow_gpu/setup.cfg
+49
View File
@@ -0,0 +1,49 @@
# Copyright 2022 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
#
# 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.
# ==============================================================================
[metadata]
name = tensorflow-gpu
version = 2.12.0
description = Removed: please install "tensorflow" instead.
long_description = file:README.md
long_description_content_type=text/markdown
python_requires = >=3.7
license = Apache 2.0
keywords = tensorflow, tensor, machine learning
classifiers =
Development Status :: 5 - Production/Stable
Intended Audience :: Developers
Intended Audience :: Education
Intended Audience :: Science/Research
License :: OSI Approved :: Apache Software License
License :: OSI Approved :: Apache Software License
Operating System :: OS Independent
Programming Language :: Python :: 3
Programming Language :: Python :: 3 :: Only
Topic :: Scientific/Engineering
Topic :: Scientific/Engineering :: Artificial Intelligence
Topic :: Scientific/Engineering :: Mathematics
Topic :: Software Development
Topic :: Software Development :: Libraries
Topic :: Software Development :: Libraries :: Python Modules
project_urls =
Homepage = https://www.tensorflow.org
Issue Tracker = https://github.com/tensorflow/tensorflow/issues
Download = https://github.com/tensorflow/tensorflow/tags
author = Google, Inc.
author_email = packages@tensorflow.org
[options]
install_requires =
python_version>"3.7"
tensorflow/tools/pip_package/redundant_tensorflow_gpu/setup.py
+44
View File
@@ -0,0 +1,44 @@
# Copyright 2022 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
#
# 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.
# ==============================================================================
"""Minimal setup.py that can build an sdist, but warns on install."""
import sys
import setuptools
TF_REMOVAL_WARNING = """
=========================================================
The "tensorflow-gpu" package has been removed!
Please install "tensorflow" instead.
Other than the name, the two packages have been identical
since TensorFlow 2.1, or roughly since Sep 2019. For more
information, see: pypi.org/project/tensorflow-gpu
=========================================================
"""
# Cover all "pip install" situations
if (
"bdist_wheel" in sys.argv
or "install" in sys.argv
or "bdist_egg" in sys.argv
):
sys.exit(TF_REMOVAL_WARNING)
if __name__ == "__main__":
setuptools.setup()
tensorflow/tools/pip_package/redundant_tf_nightly_gpu/README.md
+29
View File
@@ -0,0 +1,29 @@
tf-nightly-gpu has been removed. Please install tf-nightly instead.
The tf-nightly package supports GPU accelerated operations via Nvidia CUDA.
## Removal Information
tf-nightly and tf-nightly-gpu have been the same package since TensorFlow
2.1, released in September 2019. Although the checksums differ due to metadata,
they were built in the same way and both provide GPU support via Nvidia CUDA.
As of January 2023, tf-nightly-gpu has been removed and has been replaced with
this new, empty package that generates an error upon installation.
The TensorFlow team has stopped releasing any new tf-nightly-gpu packages, and
tf-nightly-gpu packages may disappear at any time. Please switch to tf-nightly.
## About this package
This simple package raises a warning if setup.py is executed as part of a
package installation. This intentionally prevents users from installing
the package.
To build and upload this package's source distribution (sdist) to testpypi:
```
$ vim setup.cfg # update the version number and package name
$ python3 -m pip install --user twine
$ python3 setup.py sdist
$ twine upload --repository testpypi dist/*
$ pip3 install the_name_of_your_test_package -i https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple
```
tensorflow/tools/pip_package/redundant_tf_nightly_gpu/setup.cfg
+49
View File
@@ -0,0 +1,49 @@
# Copyright 2022 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
#
# 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.
# ==============================================================================
[metadata]
name = tf-nightly-gpu
version = 2.12.0
description = Removed: please install "tf-nightly" instead.
long_description = file:README.md
long_description_content_type=text/markdown
python_requires = >=3.7
license = Apache 2.0
keywords = tensorflow, tensor, machine learning
classifiers =
Development Status :: 5 - Production/Stable
Intended Audience :: Developers
Intended Audience :: Education
Intended Audience :: Science/Research
License :: OSI Approved :: Apache Software License
License :: OSI Approved :: Apache Software License
Operating System :: OS Independent
Programming Language :: Python :: 3
Programming Language :: Python :: 3 :: Only
Topic :: Scientific/Engineering
Topic :: Scientific/Engineering :: Artificial Intelligence
Topic :: Scientific/Engineering :: Mathematics
Topic :: Software Development
Topic :: Software Development :: Libraries
Topic :: Software Development :: Libraries :: Python Modules
project_urls =
Homepage = https://www.tensorflow.org
Issue Tracker = https://github.com/tensorflow/tensorflow/issues
Download = https://github.com/tensorflow/tensorflow/tags
author = Google, Inc.
author_email = packages@tensorflow.org
[options]
install_requires =
python_version>"3.7"
tensorflow/tools/pip_package/redundant_tf_nightly_gpu/setup.py
+44
View File
@@ -0,0 +1,44 @@
# Copyright 2022 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
#
# 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.
# ==============================================================================
"""Minimal setup.py that can build an sdist, but warns on install."""
import sys
import setuptools
TF_REMOVAL_WARNING = """
=========================================================
The "tf-nightly-gpu" package has been removed!
Please install "tf-nightly" instead.
Other than the name, the two packages have been identical
since tf-nightly 2.1, or roughly since Sep 2019. For more
information, see: pypi.org/project/tf-nightly-gpu
=========================================================
"""
# Cover all "pip install" situations
if (
"bdist_wheel" in sys.argv
or "install" in sys.argv
or "bdist_egg" in sys.argv
):
sys.exit(TF_REMOVAL_WARNING)
if __name__ == "__main__":
setuptools.setup()
tensorflow/tools/pip_package/setup.py.tpl
+450
View File
@@ -0,0 +1,450 @@
# Copyright 2015 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
#
# 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..
# ==============================================================================
"""TensorFlow is an open source machine learning framework for everyone.
[![Python](https://img.shields.io/pypi/pyversions/tensorflow.svg?style=plastic)](https://badge.fury.io/py/tensorflow)
[![PyPI](https://badge.fury.io/py/tensorflow.svg)](https://badge.fury.io/py/tensorflow)
TensorFlow is an open source software library for high performance numerical
computation. Its flexible architecture allows easy deployment of computation
across a variety of platforms (CPUs, GPUs, TPUs), and from desktops to clusters
of servers to mobile and edge devices.
Originally developed by researchers and engineers from the Google Brain team
within Google's AI organization, it comes with strong support for machine
learning and deep learning and the flexible numerical computation core is used
across many other scientific domains. TensorFlow is licensed under [Apache
2.0](https://github.com/tensorflow/tensorflow/blob/master/LICENSE).
"""
import datetime
import fnmatch
import os
import platform
import re
import sys
from setuptools import Command
from setuptools import find_namespace_packages
from setuptools import setup
from setuptools.command.install import install as InstallCommandBase
from setuptools.dist import Distribution
# A genrule //tensorflow/tools/pip_package:setup_py replaces dummy string by
# by the data provided in //tensorflow/tf_version.bzl.
# The version suffix can be set by passing the build parameters
# --repo_env=ML_WHEEL_BUILD_DATE=<date> and
# --repo_env=ML_WHEEL_VERSION_SUFFIX=<suffix>.
# To update the project version, update tf_version.bzl.
# This version string is semver compatible, but incompatible with pip.
# For pip, we will remove all '-' characters from this string, and use the
# result for pip.
_VERSION = '0.0.0'
cuda_version = 0 # placeholder
cuda_major_version = '12' # placeholder
cuda_wheel_suffix = '' # placeholder
nvidia_cublas_version = '' # placeholder
nvidia_cuda_cupti_version = '' # placeholder
nvidia_cuda_nvcc_version = '' # placeholder
nvidia_cuda_runtime_version = '' # placeholder
nvidia_cudnn_version = '' # placeholder
nvidia_cufft_version = '' # placeholder
nvidia_cusolver_version = '' # placeholder
nvidia_cusparse_version = '' # placeholder
nvidia_nccl_version = '' # placeholder
nvidia_nvjitlink_version = '' # placeholder
nvidia_cuda_nvrtc_version = '' # placeholder
nvidia_curand_version = '' # placeholder
nvidia_nvshmem_version = '' # placeholder
# We use the same setup.py for all tensorflow_* packages and for the nightly
# equivalents (tf_nightly_*). The package is controlled from the argument line
# when building the pip package.
project_name = 'tensorflow'
if os.environ.get('project_name', None):
project_name = os.environ['project_name']
collaborator_build = os.environ.get('collaborator_build', False)
# Returns standard if a tensorflow-* package is being built, and nightly if a
# tf_nightly-* package is being built.
def standard_or_nightly(standard, nightly):
return nightly if 'tf_nightly' in project_name else standard
# All versions of TF need these packages. We indicate the widest possible range
# of package releases possible to be as up-to-date as possible as well as to
# accomodate as many pre-installed packages as possible.
# For packages that don't have yet a stable release, we pin using `~= 0.x` which
# means we accept any `0.y` version (y >= x) but not the first major release. We
# will need additional testing for that.
# NOTE: This assumes that all packages follow SemVer. If a package follows a
# different versioning scheme (e.g., PVP), we use different bound specifier and
# comment the versioning scheme.
REQUIRED_PACKAGES = [
'absl-py >= 1.0.0',
'astunparse >= 1.6.0',
'flatbuffers >= 25.9.23',
'gast >=0.2.1,!=0.5.0,!=0.5.1,!=0.5.2',
'google_pasta >= 0.1.1',
'libclang >= 13.0.0',
'opt_einsum >= 2.3.2',
'packaging',
'protobuf >= 6.31.1, < 8.0.0',
'requests >= 2.21.0, < 3',
'setuptools',
'six >= 1.12.0',
'termcolor >= 1.1.0',
'typing_extensions >= 3.6.6',
'wrapt >= 1.11.0',
# grpcio does not build correctly on big-endian machines due to lack of
# BoringSSL support.
# See https://github.com/tensorflow/tensorflow/issues/17882.
'grpcio >= 1.24.3, < 2.0' if sys.byteorder == 'little' else None,
# TensorFlow exposes the TF API for certain TF ecosystem packages like
# keras. When TF depends on those packages, the package version needs to
# match the current TF version. For tf_nightly, we install the nightly
# variant of each package instead, which must be one version ahead of the
# current release version. These also usually have "alpha" or "dev" in their
# version name. During the TF release process the version of these
# dependencies on the release branch is updated to the stable releases (RC
# or final). For example, 'keras-nightly ~= 2.14.0.dev' will be replaced by
# 'keras >= 2.14.0rc0, < 2.15' on the release branch after the branch cut.
'keras-nightly >= 3.12.0.dev',
'numpy >= 1.26.0',
# Starting with 3.15, only MacOS 14 and 15 are supported.
'h5py >= 3.11.0, < 3.15.0' if sys.version_info.minor <= 13 else 'h5py ~= 3.15.1',
'ml_dtypes >= 0.5.1, < 1.0.0',
]
REQUIRED_PACKAGES = [p for p in REQUIRED_PACKAGES if p is not None]
FAKE_REQUIRED_PACKAGES = [
# The depedencies here below are not actually used but are needed for
# package managers like poetry to parse as they are confused by the
# different architectures having different requirements.
# The entries here should be a simple duplicate of those in the collaborator
# build section.
]
if platform.system() == 'Linux' and platform.machine() == 'x86_64':
REQUIRED_PACKAGES.append(FAKE_REQUIRED_PACKAGES)
if collaborator_build:
# If this is a collaborator build, then build an "installer" wheel and
# add the collaborator packages as the only dependencies.
REQUIRED_PACKAGES = [
# Install the TensorFlow package built by Intel if the user is on a
# Windows machine.
standard_or_nightly('tensorflow-intel', 'tf-nightly-intel')
+ '=='
+ _VERSION
+ ';platform_system=="Windows"',
]
# Set up extra packages, which are optional sets of other Python package deps.
# E.g. "pip install tensorflow[and-cuda]" below installs the normal TF deps,
# plus the CUDA libraries listed.
EXTRA_PACKAGES = {
'and-cuda': [
# TODO(nluehr): set nvidia-* versions based on build components.
f'nvidia-cublas{cuda_wheel_suffix}{nvidia_cublas_version}',
f'nvidia-cuda-cupti{cuda_wheel_suffix}{nvidia_cuda_cupti_version}',
f'nvidia-cuda-nvcc{cuda_wheel_suffix}{nvidia_cuda_nvcc_version}',
f'nvidia-cuda-nvrtc{cuda_wheel_suffix}{nvidia_cuda_nvrtc_version}',
f'nvidia-cuda-runtime{cuda_wheel_suffix}{nvidia_cuda_runtime_version}',
f'nvidia-cudnn-cu13{nvidia_cudnn_version}' if cuda_major_version == '13' else f'nvidia-cudnn-cu12{nvidia_cudnn_version}',
f'nvidia-cufft{cuda_wheel_suffix}{nvidia_cufft_version}',
f'nvidia-curand{cuda_wheel_suffix}{nvidia_curand_version}',
f'nvidia-cusolver{cuda_wheel_suffix}{nvidia_cusolver_version}',
f'nvidia-cusparse{cuda_wheel_suffix}{nvidia_cusparse_version}',
f'nvidia-nccl-cu13{nvidia_nccl_version}' if cuda_major_version == '13' else f'nvidia-nccl-cu12{nvidia_nccl_version}',
f'nvidia-nvjitlink{cuda_wheel_suffix}{nvidia_nvjitlink_version}',
f'nvidia-nvshmem-cu13{nvidia_nvshmem_version}' if cuda_major_version == '13' else f'nvidia-nvshmem-cu12{nvidia_nvshmem_version}',
],
'gcs-filesystem': [
('tensorflow-io-gcs-filesystem>=0.23.1; '
'sys_platform!="win32" and python_version<"3.13"'),
('tensorflow-io-gcs-filesystem>=0.23.1; '
'sys_platform=="win32" and python_version<"3.12"'),
]
}
DOCLINES = __doc__.split('\n')
# pylint: disable=line-too-long
CONSOLE_SCRIPTS = [
'tflite_convert = tensorflow.lite.python.tflite_convert:main',
'toco = tensorflow.lite.python.tflite_convert:main',
'saved_model_cli = tensorflow.python.tools.saved_model_cli:main',
(
'import_pb_to_tensorboard ='
' tensorflow.python.tools.import_pb_to_tensorboard:main'
),
# We need to keep the TensorBoard command, even though the console script
# is now declared by the tensorboard pip package. If we remove the
# TensorBoard command, pip will inappropriately remove it during install,
# even though the command is not removed, just moved to a different wheel.
# We exclude it anyway if building tf_nightly.
standard_or_nightly('tensorboard = tensorboard.main:run_main', None),
'tf_upgrade_v2 = tensorflow.tools.compatibility.tf_upgrade_v2_main:main',
]
CONSOLE_SCRIPTS = [s for s in CONSOLE_SCRIPTS if s is not None]
# pylint: enable=line-too-long
class BinaryDistribution(Distribution):
def has_ext_modules(self):
return True
class InstallCommand(InstallCommandBase):
"""Override the dir where the headers go."""
def finalize_options(self):
ret = InstallCommandBase.finalize_options(self) # pylint: disable=assignment-from-no-return
self.install_headers = os.path.join(self.install_platlib, 'tensorflow',
'include')
self.install_lib = self.install_platlib
return ret
class InstallHeaders(Command):
"""Override how headers are copied.
The install_headers that comes with setuptools copies all files to
the same directory. But we need the files to be in a specific directory
hierarchy for -I <include_dir> to work correctly.
"""
description = 'install C/C++ header files'
user_options = [
('install-dir=', 'd', 'directory to install header files to'),
('force', 'f', 'force installation (overwrite existing files)'),
]
boolean_options = ['force']
def initialize_options(self):
self.install_dir = None
self.force = 0
self.outfiles = []
def finalize_options(self):
self.set_undefined_options('install', ('install_headers', 'install_dir'),
('force', 'force'))
def mkdir_and_copy_file(self, header):
install_dir = os.path.join(self.install_dir, os.path.dirname(header))
# Windows platform uses "\" in path strings, the external header location
# expects "/" in paths. Hence, we replaced "\" with "/" for this reason
install_dir = install_dir.replace('\\', '/')
# Get rid of some extra intervening directories so we can have fewer
# directories for -I
install_dir = re.sub('/google/protobuf_archive/src', '', install_dir)
# Copy external code headers into tensorflow/include.
# A symlink would do, but the wheel file that gets created ignores
# symlink within the directory hierarchy.
# NOTE(keveman): Figure out how to customize bdist_wheel package so
# we can do the symlink.
# pylint: disable=line-too-long
external_header_locations = {
'/tensorflow/include/external/com_google_absl': '',
'/tensorflow/include/external/ducc': '/ducc',
'/tensorflow/include/external/eigen_archive': '',
'/tensorflow/include/external/ml_dtypes_py': '',
'/tensorflow/include/tensorflow/compiler/xla': (
'/tensorflow/include/xla'
),
'/tensorflow/include/tensorflow/tsl': '/tensorflow/include/tsl',
}
# pylint: enable=line-too-long
for location in external_header_locations:
if location in install_dir:
extra_dir = install_dir.replace(location,
external_header_locations[location])
if not os.path.exists(extra_dir):
self.mkpath(extra_dir)
self.copy_file(header, extra_dir)
if not os.path.exists(install_dir):
self.mkpath(install_dir)
return self.copy_file(header, install_dir)
def run(self):
hdrs = self.distribution.headers
if not hdrs:
return
self.mkpath(self.install_dir)
for header in hdrs:
(out, _) = self.mkdir_and_copy_file(header)
self.outfiles.append(out)
def get_inputs(self):
return self.distribution.headers or []
def get_outputs(self):
return self.outfiles
def find_files(pattern, root):
"""Return all the files matching pattern below root dir."""
for dirpath, _, files in os.walk(root):
for filename in fnmatch.filter(files, pattern):
yield os.path.join(dirpath, filename)
so_lib_paths = [
i for i in os.listdir('.')
if os.path.isdir(i) and fnmatch.fnmatch(i, '_solib_*')
]
matches = []
for path in so_lib_paths:
matches.extend(['../' + x for x in find_files('*', path) if '.py' not in x])
if '_tpu' in project_name:
REQUIRED_PACKAGES.append([f'libtpu~=0.0.14'])
CONSOLE_SCRIPTS.extend([
'start_grpc_tpu_worker = tensorflow.python.tools.grpc_tpu_worker:run',
('start_grpc_tpu_service = '
'tensorflow.python.tools.grpc_tpu_worker_service:run'),
])
if os.name == 'nt':
EXTENSION_NAME = 'python/_pywrap_tensorflow_internal.pyd'
else:
EXTENSION_NAME = 'python/_pywrap_tensorflow_internal.so'
headers = (
list(find_files('*.proto', 'tensorflow/compiler'))
+ list(find_files('*.proto', 'tensorflow/core'))
+ list(find_files('*.proto', 'tensorflow/python'))
+ list(find_files('*.proto', 'tensorflow/python/framework'))
+ list(find_files('*.proto', 'tensorflow/tsl'))
+ list(find_files('*.def', 'tensorflow/compiler'))
+ list(find_files('*.h', 'tensorflow/c'))
+ list(find_files('*.h', 'tensorflow/cc'))
+ list(find_files('*.h', 'tensorflow/compiler'))
+ list(find_files('*.h.inc', 'tensorflow/compiler'))
+ list(find_files('*.h', 'tensorflow/core'))
+ list(find_files('*.h', 'tensorflow/lite/kernels/shim'))
+ list(find_files('*.h', 'tensorflow/python'))
+ list(find_files('*.h', 'tensorflow/python/client'))
+ list(find_files('*.h', 'tensorflow/python/framework'))
+ list(find_files('*.h', 'tensorflow/stream_executor'))
+ list(find_files('*.h', 'tensorflow/compiler/xla/stream_executor'))
+ list(find_files('*.h', 'tensorflow/tsl'))
+ list(find_files('*.h', 'google/com_google_protobuf/src'))
+ list(find_files('*.inc', 'google/com_google_protobuf/src'))
+ list(find_files('*', 'third_party/gpus'))
+ list(find_files('*.h', 'tensorflow/include/external/com_google_absl'))
+ list(find_files('*.inc', 'tensorflow/include/external/com_google_absl'))
+ list(find_files('*.h', 'tensorflow/include/external/ducc/google'))
+ list(find_files('*', 'tensorflow/include/external/eigen_archive'))
+ list(find_files('*.h', 'tensorflow/include/external/ml_dtypes_py'))
)
# Quite a lot of setup() options are different if this is a collaborator package
# build. We explicitly list the differences here, then unpack the dict as
# options at the end of the call to setup() below. For what each keyword does,
# see https://setuptools.pypa.io/en/latest/references/keywords.html.
if collaborator_build:
collaborator_build_dependent_options = {
'cmdclass': {},
'distclass': None,
'entry_points': {},
'headers': [],
'include_package_data': None,
'packages': [],
'package_data': {},
}
else:
collaborator_build_dependent_options = {
'cmdclass': {
'install_headers': InstallHeaders,
'install': InstallCommand,
},
'distclass': BinaryDistribution,
'entry_points': {
'console_scripts': CONSOLE_SCRIPTS,
},
'headers': headers,
'include_package_data': True,
'packages': find_namespace_packages(),
'package_data': {
'tensorflow': [EXTENSION_NAME] + matches,
},
}
CLASSIFIERS = [
'Development Status :: 5 - Production/Stable',
# TODO(angerson) Add IFTTT when possible
'Intended Audience :: Developers',
'Intended Audience :: Education',
'Intended Audience :: Science/Research',
'License :: OSI Approved :: Apache Software License',
'Programming Language :: Python :: 3',
'Programming Language :: Python :: 3.10',
'Programming Language :: Python :: 3.11',
'Programming Language :: Python :: 3.12',
'Programming Language :: Python :: 3.13',
'Programming Language :: Python :: 3 :: Only',
'Topic :: Scientific/Engineering',
'Topic :: Scientific/Engineering :: Mathematics',
'Topic :: Scientific/Engineering :: Artificial Intelligence',
'Topic :: Software Development',
'Topic :: Software Development :: Libraries',
'Topic :: Software Development :: Libraries :: Python Modules',
]
if cuda_major_version:
CLASSIFIERS.extend([
f'Environment :: GPU :: NVIDIA CUDA :: {cuda_major_version}',
f'Environment :: GPU :: NVIDIA CUDA :: {cuda_major_version} :: {cuda_major_version}.0',
])
setup(
name=project_name,
version=_VERSION.replace('-', ''),
description=DOCLINES[0],
long_description='\n'.join(DOCLINES[2:]),
long_description_content_type='text/markdown',
url='https://www.tensorflow.org/',
download_url='https://github.com/tensorflow/tensorflow/tags',
author='Google Inc.',
author_email='packages@tensorflow.org',
install_requires=REQUIRED_PACKAGES,
extras_require=EXTRA_PACKAGES,
# Add in any packaged data.
zip_safe=False,
# Supported Python versions
# Python 3.9 support was dropped in TensorFlow 2.17.
python_requires='>=3.10',
# PyPI package information.
classifiers=sorted(CLASSIFIERS),
license='Apache 2.0',
keywords='tensorflow tensor machine learning',
nvidia_cuda_nvcc_version=nvidia_cuda_nvcc_version,
nvidia_cuda_nvrtc_version=nvidia_cuda_nvrtc_version,
nvidia_cuda_runtime_version=nvidia_cuda_runtime_version,
**collaborator_build_dependent_options
)
tensorflow/tools/pip_package/simple_console.py
+29
View File
@@ -0,0 +1,29 @@
# Copyright 2015 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
#
# 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.
# ==============================================================================
"""Start a simple interactive console with TensorFlow available."""
import code
import sys
def main(_):
"""Run an interactive console."""
code.interact()
return 0
if __name__ == '__main__':
sys.exit(main(sys.argv))
tensorflow/tools/pip_package/simple_console_for_windows.py
+29
View File
@@ -0,0 +1,29 @@
# Copyright 2015 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
#
# 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.
# ==============================================================================
"""Start a simple interactive console with TensorFlow available."""
import code
import sys
def main(_):
"""Run an interactive console."""
code.interact()
return 0
if __name__ == '__main__':
sys.exit(main(sys.argv))
tensorflow/tools/pip_package/utils/BUILD
+8
View File
@@ -0,0 +1,8 @@
package(default_visibility = ["//visibility:public"])
py_library(
name = "py_utils",
srcs = [
"utils.py",
],
)
tensorflow/tools/pip_package/utils/tf_wheel.bzl
+172
View File
@@ -0,0 +1,172 @@
# Copyright 2023 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
#
# 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.
"""Rule to build custom wheel.
It parses prvoided inputs and then calls `build_pip_package_py` binary with following args:
1) `--project-name` - name to be passed to setup.py file. It will define name of the wheel.
Should be set via --repo_env=WHEEL_NAME=tensorflow_cpu.
2) `--collab` - whether this is a collaborator build.
3) `--version` - tensorflow version.
4) `--headers` - paths to header file.
5) `--srcs` - paths to source files
6) `--dests` - json file with source to destination mappings for files whose original
location does not match its destination in packaged wheel; if the destination is an
empty string the source file will be ignored.
7) `--xla_aot` - paths to files that should be in xla_aot directory.
"""
load("@bazel_skylib//rules:common_settings.bzl", "BuildSettingInfo")
load(
"@python_version_repo//:py_version.bzl",
"HERMETIC_PYTHON_VERSION",
"MACOSX_DEPLOYMENT_TARGET",
"WHEEL_COLLAB",
"WHEEL_NAME",
)
load(
"//tensorflow:tf_version.bzl",
"TF_VERSION",
"TF_WHEEL_VERSION_SUFFIX",
)
def get_canonical_repo_name(apparent_repo_name):
"""Returns the canonical repo name for the given apparent repo name seen by the module this bzl file belongs to."""
if not apparent_repo_name.startswith("@"):
apparent_repo_name = "@" + apparent_repo_name
return Label(apparent_repo_name).workspace_name
def _get_wheel_platform_name(platform_name, platform_tag):
macos_platform_version = "{}_".format(MACOSX_DEPLOYMENT_TARGET.replace(".", "_")) if MACOSX_DEPLOYMENT_TARGET else ""
tag = platform_tag
if platform_tag == "x86_64" and platform_name == "win":
tag = "amd64"
if platform_tag == "arm64" and platform_name == "linux":
tag = "aarch64"
return "{platform_name}_{platform_version}{platform_tag}".format(
platform_name = platform_name,
platform_tag = tag,
platform_version = macos_platform_version,
)
def _get_full_wheel_name(
platform_name,
platform_tag,
wheel_version):
python_version = HERMETIC_PYTHON_VERSION.replace(".", "")
return "{wheel_name}-{wheel_version}-cp{python_version}-cp{python_version}-{wheel_platform_tag}.whl".format(
wheel_name = WHEEL_NAME,
wheel_version = wheel_version,
python_version = python_version,
wheel_platform_tag = _get_wheel_platform_name(
platform_name,
platform_tag,
),
)
def _is_dest_file(basename, dest_files_suffixes):
for suffix in dest_files_suffixes:
if basename.endswith(suffix):
return True
return False
def _tf_wheel_impl(ctx):
include_cuda_libs = ctx.attr.include_cuda_libs[BuildSettingInfo].value
override_include_cuda_libs = ctx.attr.override_include_cuda_libs[BuildSettingInfo].value
if include_cuda_libs and not override_include_cuda_libs:
fail("TF wheel shouldn't be built with CUDA dependencies." +
" Please provide `--config=cuda_wheel` for bazel build command." +
" If you absolutely need to add CUDA dependencies, provide" +
" `--@local_config_cuda//cuda:override_include_cuda_libs=true`.")
executable = ctx.executable.wheel_binary
full_wheel_version = (TF_VERSION + TF_WHEEL_VERSION_SUFFIX)
full_wheel_name = _get_full_wheel_name(
platform_name = ctx.attr.platform_name,
platform_tag = ctx.attr.platform_tag,
wheel_version = full_wheel_version,
)
wheel_dir_name = "wheel_house"
output_file = ctx.actions.declare_file("{wheel_dir}/{wheel_name}".format(
wheel_dir = wheel_dir_name,
wheel_name = full_wheel_name,
))
wheel_dir = output_file.path[:output_file.path.rfind("/")]
args = ctx.actions.args()
args.add("--project-name", WHEEL_NAME)
args.add("--platform", _get_wheel_platform_name(
ctx.attr.platform_name,
ctx.attr.platform_tag,
))
args.add("--collab", str(WHEEL_COLLAB))
args.add("--output-name", wheel_dir)
args.add("--version", TF_VERSION)
headers = ctx.files.headers[:]
for f in headers:
args.add("--headers=%s" % (f.path))
xla_aot = ctx.files.xla_aot_compiled[:]
for f in xla_aot:
args.add("--xla_aot=%s" % (f.path))
srcs = []
for src in ctx.attr.source_files:
for f in src.files.to_list():
srcs.append(f)
if _is_dest_file(f.basename, ctx.attr.dest_files_suffixes):
args.add("--dests=%s" % (f.path))
else:
args.add("--srcs=%s" % (f.path))
args.set_param_file_format("flag_per_line")
args.use_param_file("@%s", use_always = False)
env = {}
for repo in ctx.attr.external_repos:
env[repo.upper().replace("-", "_").replace("/", "_") + "_CANONICAL_REPO_NAME"] = get_canonical_repo_name(repo)
ctx.actions.run(
mnemonic = "TFWheel",
arguments = [args],
inputs = srcs + headers + xla_aot,
outputs = [output_file],
executable = executable,
use_default_shell_env = True,
env = env,
)
return [DefaultInfo(files = depset(direct = [output_file]))]
tf_wheel = rule(
attrs = {
"source_files": attr.label_list(allow_files = True),
"dest_files_suffixes": attr.string_list(default = ["_wheel_locations.json"]),
"headers": attr.label_list(allow_files = True),
"xla_aot_compiled": attr.label_list(allow_files = True),
"wheel_binary": attr.label(
default = Label("//tensorflow/tools/pip_package:build_pip_package_py"),
executable = True,
cfg = "exec",
),
"include_cuda_libs": attr.label(default = Label("@local_config_cuda//cuda:include_cuda_libs")),
"override_include_cuda_libs": attr.label(default = Label("@local_config_cuda//cuda:override_include_cuda_libs")),
"platform_tag": attr.string(mandatory = True),
"platform_name": attr.string(mandatory = True),
"external_repos": attr.string_list(),
},
implementation = _tf_wheel_impl,
)
def tf_wheel_dep():
return "@pypi_{}//:whl".format(WHEEL_NAME)
tensorflow/tools/pip_package/utils/utils.py
+79
View File
@@ -0,0 +1,79 @@
# Copyright 2023 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.
"""Utils for pip package builder."""
import os
import shutil
import sys
def is_windows() -> bool:
return sys.platform.startswith("win32")
def is_macos() -> bool:
return sys.platform.startswith("darwin")
def create_init_files(dst_dir: str) -> None:
"""Create __init__.py files."""
for root, _, files in os.walk(dst_dir):
if any(file.endswith(".py") or file.endswith(".so") for file in files):
curr_dir = root
while curr_dir != dst_dir:
init_path = os.path.join(curr_dir, "__init__.py")
if not os.path.exists(init_path):
open(init_path, "w").close()
curr_dir = os.path.dirname(curr_dir)
def replace_inplace(directory, search, to_replace) -> None:
"""Traverse the directory and replace search phrase in each file."""
for root, _, files in os.walk(directory):
for file_name in files:
if file_name.endswith(".py"):
file_path = os.path.join(root, file_name)
with open(file_path, "r", encoding="utf-8") as file:
filedata = file.read()
if search in filedata:
filedata = filedata.replace(search, to_replace)
with open(file_path, "w") as file:
file.write(filedata)
def copy_file(
src_file: str,
dst_dir: str,
strip: str = None,
dest_file: str = None,
) -> None:
"""Copy a file to the destination directory.
Args:
src_file: file to be copied
dst_dir: destination directory
strip: prefix to strip before copying to destination
dest_file: destanation file location if different from src_file
"""
# strip `bazel-out/.../bin/` for generated files.
dest = dest_file if dest_file else src_file
if dest.startswith("bazel-out"):
dest = dest[dest.index("bin") + 4 :]
if strip:
dest = dest.removeprefix(strip)
dest_dir_path = os.path.join(dst_dir, os.path.dirname(dest))
os.makedirs(dest_dir_path, exist_ok=True)
shutil.copy(src_file, dest_dir_path)
os.chmod(os.path.join(dst_dir, dest), 0o644)
tensorflow/tools/pip_package/xla_build/CMakeLists.txt
+34
View File
@@ -0,0 +1,34 @@
# Copyright 2020 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
#
# 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.
# ============================================================================
cmake_minimum_required(VERSION 3.5)
file(GLOB_RECURSE TF_RUNTIME_SRC "*.cc")
add_library(tf_xla_runtime_objects OBJECT
${TF_RUNTIME_SRC}
)
target_include_directories(tf_xla_runtime_objects PRIVATE ../include)
target_compile_options(tf_xla_runtime_objects PRIVATE
-ftemplate-backtrace-limit=0
-Wno-ignored-attributes
-Wno-deprecated-copy
-Wno-cast-qual
-Wno-sign-compare
)
add_library(tf_xla_runtime STATIC
$<TARGET_OBJECTS:tf_xla_runtime_objects>
)
tensorflow/tools/pip_package/xla_build/pip_test/CMakeLists.txt
+61
View File
@@ -0,0 +1,61 @@
# Copyright 2022 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
#
# 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.
# ============================================================================
cmake_minimum_required(VERSION 3.14)
project(TestAOT)
set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CXX_STANDARD_REQUIRED ON)
include(FetchContent)
FetchContent_Declare(
googletest
GIT_REPOSITORY https://github.com/google/googletest.git
GIT_TAG release-1.12.1
)
FetchContent_MakeAvailable(googletest)
enable_testing()
add_definitions(-DTF_PIP_INTEGRATION_TEST)
set(TENSORFLOW_PACKAGE_PATH "/usr/local/lib/python3.9/dist-packages/tensorflow"
CACHE STRING "Full path of tensorflow python package")
set(CMAKE_ARCHIVE_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/lib)
add_executable(aot_compiled_test
aot_compiled_test.cc
aot_compiled_vars_and_arithmetic.o
aot_compiled_vars_and_arithmetic_frozen.o
aot_compiled_x_matmul_y_small.o
aot_compiled_x_matmul_y_large.o
aot_compiled_x_matmul_y_large_multithreaded.o
aot_compiled_x_plus_y.o
)
# The paths assume the setup performed by run_integration_test.sh.
target_include_directories(aot_compiled_test PUBLIC
"/tmp/generated_models"
"${TENSORFLOW_PACKAGE_PATH}/include"
"/tmp/tf_third_party")
add_subdirectory("${TENSORFLOW_PACKAGE_PATH}/xla_aot_runtime_src"
${CMAKE_ARCHIVE_OUTPUT_DIRECTORY}/tf_runtime)
target_link_libraries(aot_compiled_test GTest::gtest_main tf_xla_runtime)
include(GoogleTest)
gtest_discover_tests(aot_compiled_test)
tensorflow/tools/pip_package/xla_build/pip_test/run_xla_aot_test.sh
+159
View File
@@ -0,0 +1,159 @@
#!/bin/bash
# Copyright 2015 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
#
# 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.
# ==============================================================================
#=============================================================================
# this script must be run from the root of the tensorflow git repository.
#
# it *can* be run outside docker/kokoro, on a dev machine, as long as cmake and
# ninja-build packages are installed, and the tensorflow PIP package (the one
# under test, presumably) is installed.
#
# the script creates a few directories under /tmp, see a few lines below.
#
# under kokoro, this is run by learning/brain/testing/kokoro/rel/docker/aot_compile.sh
#=============================================================================
set -eo pipefail -o history
# If the TENSORFLOW_PACKAGE_PATH variable is set, add an argument to the cmake
# command-line to explicitly set the package path.
# This must be done before disallowing unset variables.
CMAKE_TENSORFLOW_PACKAGE_ARG=
if [ ! -z "$TENSORFLOW_PACKAGE_PATH" ]; then
CMAKE_TENSORFLOW_PACKAGE_ARG="-DTENSORFLOW_PACKAGE_PATH=${TENSORFLOW_PACKAGE_PATH}"
fi
set -u
echo "Building x_matmul_y models"
python3 \
tensorflow/python/tools/make_aot_compile_models.py \
--out_dir=/tmp/saved_models
# LINT.IfChange
GEN_ROOT=/tmp/generated_models
GEN_PREFIX="${GEN_ROOT}/tensorflow/python/tools"
PROJECT=/tmp/project
TF_THIRD_PARTY=/tmp/tf_third_party
# LINT.ThenChange(//tensorflow/tools/pip_package/xla_build/pip_test/CMakeLists.txt)
export PATH=$PATH:${HOME}/.local/bin
rm -rf "${GEN_ROOT}" "${PROJECT}" "${TF_THIRD_PARTY}"
# We don't want to -Itensorflow, to avoid unwanted dependencies.
echo "Copying third_party stuff (for eigen)"
mkdir -p "${TF_THIRD_PARTY}"
cp -rf third_party "${TF_THIRD_PARTY}/"
echo "AOT models"
saved_model_cli aot_compile_cpu \
--dir tensorflow/cc/saved_model/testdata/VarsAndArithmeticObjectGraph \
--output_prefix "${GEN_PREFIX}/aot_compiled_vars_and_arithmetic" \
--variables_to_feed variable_x \
--cpp_class VarsAndArithmetic --signature_def_key serving_default \
--tag_set serve
saved_model_cli aot_compile_cpu \
--dir tensorflow/cc/saved_model/testdata/VarsAndArithmeticObjectGraph \
--output_prefix "${GEN_PREFIX}/aot_compiled_vars_and_arithmetic_frozen" \
--cpp_class VarsAndArithmeticFrozen --signature_def_key serving_default \
--tag_set serve
saved_model_cli aot_compile_cpu \
--dir /tmp/saved_models/x_matmul_y_small \
--output_prefix "${GEN_PREFIX}/aot_compiled_x_matmul_y_small" \
--cpp_class XMatmulYSmall --signature_def_key serving_default --tag_set serve
saved_model_cli aot_compile_cpu \
--dir /tmp/saved_models/x_matmul_y_large \
--output_prefix "${GEN_PREFIX}/aot_compiled_x_matmul_y_large" \
--cpp_class XMatmulYLarge --signature_def_key serving_default --tag_set serve
saved_model_cli aot_compile_cpu \
--dir /tmp/saved_models/x_matmul_y_large \
--output_prefix "${GEN_PREFIX}/aot_compiled_x_matmul_y_large_multithreaded" \
--multithreading True \
--cpp_class XMatmulYLargeMultithreaded --signature_def_key serving_default \
--tag_set serve
saved_model_cli aot_compile_cpu \
--dir tensorflow/cc/saved_model/testdata/x_plus_y_v2_debuginfo \
--output_prefix "${GEN_PREFIX}/aot_compiled_x_plus_y" \
--cpp_class XPlusY --signature_def_key serving_default --tag_set serve
echo "Creating project and copying object files"
mkdir -p "${PROJECT}"
cp -f \
"${GEN_PREFIX}/aot_compiled_vars_and_arithmetic.o" \
"${GEN_PREFIX}/aot_compiled_vars_and_arithmetic_frozen.o" \
"${GEN_PREFIX}/aot_compiled_x_matmul_y_small.o" \
"${GEN_PREFIX}/aot_compiled_x_matmul_y_large.o" \
"${GEN_PREFIX}/aot_compiled_x_matmul_y_large_multithreaded.o" \
"${GEN_PREFIX}/aot_compiled_x_plus_y.o" \
"${PROJECT}"
echo "Copying build and source files"
cp tensorflow/tools/pip_package/xla_build/pip_test/CMakeLists.txt \
tensorflow/python/tools/aot_compiled_test.cc \
"${PROJECT}"
echo "Building"
mkdir "${PROJECT}/build"
cmake -GNinja -S "${PROJECT}" -B "${PROJECT}/build" -DCMAKE_BUILD_TYPE=Release $CMAKE_TENSORFLOW_PACKAGE_ARG
ninja -C "${PROJECT}/build"
echo "Running test"
"${PROJECT}/build/aot_compiled_test"
echo "Cross-compile AOT models"
failed=0
function check_crosscompile() {
local triple="${1}"
local check_string="${2}"
echo "============================"
echo "Cross-compiling to ${triple}"
echo "============================"
saved_model_cli aot_compile_cpu \
--dir tensorflow/cc/saved_model/testdata/VarsAndArithmeticObjectGraph \
--output_prefix "${GEN_PREFIX}/${triple}/aot_compiled_vars_and_arithmetic" \
--variables_to_feed variable_x \
--target_triple "${triple}" \
--cpp_class VarsAndArithmetic \
--signature_def_key serving_default \
--tag_set serve
file "${GEN_PREFIX}/${triple}/aot_compiled_vars_and_arithmetic.o" \
| grep "${check_string}"
if [ 0 -ne $? ]
then
echo "${triple}: FAILED"
failed=1
else
echo "${triple}: SUCCESS"
fi
}
check_crosscompile aarch64-unknown-linux-gnu \
"ELF 64-bit LSB relocatable, ARM aarch64"
check_crosscompile x86_64-unknown-linux-gnu "ELF 64-bit LSB relocatable, x86-64"
check_crosscompile arm64-apple-darwin "Mach-O 64-bit arm64 object"
check_crosscompile x86_64-apple-darwin "Mach-O 64-bit x86_64 object"
# check_crosscompile aarch64-pc-windows-msvc "<tbd>"
check_crosscompile x86_64-pc-windows-msvc "Intel amd64 COFF object file"
# check_crosscompile riscv64gc-unknown-linux-gnu "<tbd>"
exit "${failed}"