wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 3 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/python/platform/BUILD
+402
View File
@@ -0,0 +1,402 @@
# platform package
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load(
"//tensorflow:tensorflow.bzl",
"if_oss",
)
load("//tensorflow:tensorflow.default.bzl", "pybind_extension", "tf_py_build_info_genrule", "tf_py_strict_test", "tf_python_pybind_extension")
load("//tensorflow/core/platform:build_config.bzl", "pyx_library", "tf_additional_all_protos", "tf_additional_lib_deps", "tf_proto_library", "tf_protos_grappler") # @unused
visibility = [
"//tensorflow:__subpackages__",
"//tensorflow/dtensor:dtensor-internal",
# copybara:uncomment "//learning/brain/python/platform:__subpackages__",
# copybara:uncomment "//learning/brain/contrib/eager/numlib/benchmarks/kumamon:__subpackages__",
# copybara:uncomment "//learning/brain/mobile/lite/tooling/model_analyzer:__subpackages__",
# copybara:uncomment "//tensorflow_serving/model_servers:__subpackages__",
# copybara:uncomment "//third_party/odml/model_customization/quantization:__subpackages__",
]
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = visibility,
licenses = ["notice"],
)
tf_py_build_info_genrule(
name = "py_build_info_gen",
out = "build_info.py",
)
py_library(
name = "build_info",
srcs = ["build_info.py"],
)
py_library(
name = "windows_lib_diagnostics",
srcs = ["windows_lib_diagnostics.py"],
)
py_library(
name = "self_check",
srcs = ["self_check.py"],
deps = if_oss(
[
":build_info",
":_pywrap_cpu_feature_guard",
],
),
)
py_library(
name = "benchmark",
srcs = ["benchmark.py"],
visibility = visibility + ["//tensorflow:internal"],
deps = [
":gfile",
":tf_logging",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/client:timeline",
"//tensorflow/python/framework:ops",
"//tensorflow/python/util:tf_export",
"//tensorflow/python/util:tf_inspect",
"@absl_py//absl:app",
],
)
py_library(
name = "analytics",
srcs = ["analytics.py"],
)
py_library(
name = "device_context",
srcs = ["device_context.py"],
deps = [
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:control_flow_ops",
],
)
py_library(
name = "test",
srcs = ["googletest.py"],
# copybara:uncomment_begin(google-only)
# visibility = visibility + ["//tensorflow:internal"],
# copybara:uncomment_end_and_comment_begin
visibility = [
"//visibility:public",
],
# copybara:comment_end
deps = [
":benchmark",
":flags",
"//tensorflow/python/util:tf_export",
] + if_oss([
":tf_logging",
"@absl_py//absl:app",
"@absl_py//absl/testing:absltest",
"//tensorflow/python/framework:errors",
"//tensorflow/python/util:tf_inspect",
]),
)
tf_py_strict_test(
name = "resource_loader_test",
size = "small",
srcs = ["resource_loader_test.py"],
data = [
"resource_loader.py",
],
tags = [
"no_pip",
"no_windows",
],
deps = [
":resource_loader",
":test",
],
)
tf_py_strict_test(
name = "sysconfig_test",
size = "small",
srcs = ["sysconfig_test.py"],
data = [
"sysconfig.py",
],
tags = [
"no_mac", # TODO(b/259295275) re-enable after fixing sysconfig.get_path breakage
"no_oss", # TODO(b/259295275) re-enable after fixing sysconfig.get_path breakage
"no_pip",
"no_windows",
],
deps = [
":client_testlib",
":sysconfig",
":test",
"//tensorflow:tensorflow_py",
],
)
tf_py_strict_test(
name = "flags_test",
size = "small",
srcs = ["flags_test.py"],
tags = [
"no_mac", # TODO(b/259295275) re-enable after fixing sysconfig.get_path breakage
"no_oss", # TODO(b/263966250) re-enable after fixing sysconfig.get_path breakage
],
deps = [
":client_testlib",
":flags",
"@absl_py//absl/flags",
],
)
tf_py_strict_test(
name = "stacktrace_handler_test",
size = "small",
srcs = ["stacktrace_handler_test.py"],
tags = [
"no_oss", # TODO(b/263966250) re-enable after fixing sysconfig.get_path breakage
"no_windows",
"nomac",
"nozapfhahn",
],
deps = [
":client_testlib",
":tf_logging",
],
)
tf_py_strict_test(
name = "app_test",
size = "small",
srcs = ["app_test.py"],
tags = [
"manual",
"notap",
],
deps = [
":app",
":flags",
],
)
tf_python_pybind_extension(
name = "_pywrap_stacktrace_handler",
srcs = ["stacktrace_handler_wrapper.cc"],
hdrs = [
"//tensorflow/core/platform:stacktrace_handler_hdrs",
"@tsl//tsl/platform:stacktrace_handler_hdrs",
],
enable_stub_generation = True,
pytype_srcs = [
"_pywrap_stacktrace_handler.pyi",
],
deps = [
"@pybind11",
"@tsl//tsl/platform:stacktrace_handler_hdrs_lib",
"@xla//third_party/python_runtime:headers",
],
)
tf_python_pybind_extension(
name = "_pywrap_cpu_feature_guard",
srcs = ["cpu_feature_guard_wrapper.cc"],
enable_stub_generation = True,
pytype_srcs = [
"_pywrap_cpu_feature_guard.pyi",
],
deps = [
"//tensorflow/core/platform:cpu_feature_guard_hdr", # Only depend on header to avoid ODR issues.
"@pybind11",
"@xla//third_party/python_runtime:headers",
],
)
py_library(
name = "client_testlib",
srcs = ["test.py"],
# copybara:uncomment_begin(google-only)
# visibility = visibility + [
# "//third_party/cloud_tpu/convergence_tools:__subpackages__",
# "//third_party/py/tf_slim:__subpackages__",
# "//tensorflow:internal",
# "//tensorflow_models:__subpackages__",
# ],
# copybara:uncomment_end_and_comment_begin
visibility = [
"//visibility:public",
],
# copybara:comment_end
deps = [
":test",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:gradient_checker",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "app",
srcs = ["app.py"],
deps = [
":flags",
"//tensorflow/python/util:_pywrap_util_port",
"//tensorflow/python/util:tf_export",
"@absl_py//absl:app",
],
)
py_library(
name = "sysconfig",
srcs = ["sysconfig.py"],
deps = [
":build_info",
"//tensorflow/python/client:pywrap_tf_session",
"//tensorflow/python/framework:versions",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "__init__",
srcs = ["__init__.py"],
deps = [
],
)
py_library(
name = "control_imports",
srcs = ["control_imports.py"],
deps = [
],
)
py_library(
name = "parameterized",
srcs = ["parameterized.py"],
)
py_library(
name = "remote_utils",
srcs = ["remote_utils.py"],
deps = [
],
)
py_library(
name = "gfile",
srcs = ["gfile.py"],
# copybara:uncomment_begin(google-only)
# visibility = visibility,
# copybara:uncomment_end_and_comment_begin
visibility = [
"//visibility:public",
],
# copybara:comment_end
deps = [
"//tensorflow/python/lib/io:file_io",
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "tf_logging",
srcs = ["tf_logging.py"],
# copybara:uncomment_begin(google-only)
# visibility = visibility + [
# "//learning/brain/contrib/learn:__pkg__",
# "//learning/brain/mobile/lite/lstm:__pkg__",
# "//third_party/py/tf_keras/distribute:__pkg__",
# "//third_party/py/tf_keras/layers/preprocessing:__pkg__",
# "//third_party/py/tf_keras/models:__pkg__",
# "//third_party/py/tf_keras/optimizers:__pkg__",
# "//third_party/py/tf_keras/tests:__pkg__",
# "//third_party/py/tf_slim/training:__pkg__",
# ],
# copybara:uncomment_end_and_comment_begin
visibility = [
"//visibility:public",
],
# copybara:comment_end
deps = [
"//tensorflow/python/util:tf_export",
"@absl_py//absl/logging",
],
)
py_library(
name = "flags",
srcs = ["flags.py"],
visibility = ["//visibility:public"],
deps = [
"//tensorflow/python/util:tf_decorator_py",
"@absl_py//absl/flags",
],
)
py_library(
name = "resource_loader",
srcs = ["resource_loader.py"],
visibility = ["//visibility:public"],
deps = [
":tf_logging",
"//tensorflow/python/util:tf_export",
"//tensorflow/python/util:tf_inspect",
] + if_oss([
"@rules_python//python/runfiles",
]),
)
tf_py_strict_test(
name = "build_info_test",
size = "small",
srcs = ["build_info_test.py"],
main = "build_info_test.py",
tags = [
"no_pip",
"notap",
],
deps = [
":build_info",
":client_testlib",
"//tensorflow/compiler/tf2tensorrt:_pywrap_py_utils",
],
)
tf_py_strict_test(
name = "benchmark_test",
size = "small",
srcs = ["benchmark_test.py"],
main = "benchmark_test.py",
tags = [
"no_pip",
],
deps = [
":benchmark",
":client_testlib",
"//tensorflow/core:protos_all_py",
],
)
pybind_extension(
name = "_pywrap_tf2",
srcs = ["enable_tf2.cc"],
hdrs = ["//tensorflow/core/platform:enable_tf2_hdr"],
enable_stub_generation = True,
pytype_srcs = [
"_pywrap_tf2.pyi",
],
deps = [
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/platform:enable_tf2_utils",
"@pybind11",
],
)
tensorflow/python/platform/__init__.py
View File
tensorflow/python/platform/_pywrap_cpu_feature_guard.pyi
+16
View File
@@ -0,0 +1,16 @@
# 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.
# ==============================================================================
def InfoAboutUnusedCPUFeatures() -> None: ...
tensorflow/python/platform/_pywrap_stacktrace_handler.pyi
+16
View File
@@ -0,0 +1,16 @@
# 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.
# ==============================================================================
def InstallStacktraceHandler() -> None: ...
tensorflow/python/platform/_pywrap_tf2.pyi
+17
View File
@@ -0,0 +1,17 @@
# 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.
# ==============================================================================
def enable(arg0: bool) -> None: ...
def is_enabled() -> bool: ...
tensorflow/python/platform/analytics.py
+24
View File
@@ -0,0 +1,24 @@
# Copyright 2019 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.
# ==============================================================================
"""Analytics helpers library."""
def track_usage(tool_id, tags):
"""No usage tracking for external library.
Args:
tool_id: A string identifier for tool to be tracked.
tags: list of string tags that will be added to the tracking.
"""
del tool_id, tags # Unused externally.
tensorflow/python/platform/app.py
+36
View File
@@ -0,0 +1,36 @@
# 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.
# ==============================================================================
"""Generic entry point script."""
import sys as _sys
from absl.app import run as _run
from tensorflow.python.platform import flags
from tensorflow.python.util.tf_export import tf_export
def _parse_flags_tolerate_undef(argv):
"""Parse args, returning any unknown flags (ABSL defaults to crashing)."""
return flags.FLAGS(_sys.argv if argv is None else argv, known_only=True)
@tf_export(v1=['app.run'])
def run(main=None, argv=None):
"""Runs the program with an optional 'main' function and 'argv' list."""
main = main or _sys.modules['__main__'].main
_run(main=main, argv=argv, flags_parser=_parse_flags_tolerate_undef)
tensorflow/python/platform/app_test.py
+41
View File
@@ -0,0 +1,41 @@
# 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.
# ==============================================================================
"""Tests for our flags implementation."""
import sys
from tensorflow.python.platform import app
from tensorflow.python.platform import flags
FLAGS = flags.FLAGS
flags.DEFINE_boolean('myflag', False, '')
def main(argv):
if (len(argv) != 3):
print("Length of argv was not 3: ", argv)
sys.exit(-1)
if argv[1] != "--passthrough":
print("--passthrough argument not in argv")
sys.exit(-1)
if argv[2] != "extra":
print("'extra' argument not in argv")
sys.exit(-1)
if __name__ == '__main__':
sys.argv.extend(["--myflag", "--passthrough", "extra"])
app.run()
tensorflow/python/platform/benchmark.py
+489
View File
@@ -0,0 +1,489 @@
# Copyright 2016 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.
# ==============================================================================
"""Utilities to run benchmarks."""
import math
import numbers
import os
import re
import sys
import time
import types
from absl import app
from tensorflow.core.protobuf import config_pb2
from tensorflow.core.protobuf import rewriter_config_pb2
from tensorflow.core.util import test_log_pb2
from tensorflow.python.client import timeline
from tensorflow.python.framework import ops
from tensorflow.python.platform import gfile
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.util import tf_inspect
from tensorflow.python.util.tf_export import tf_export
# When a subclass of the Benchmark class is created, it is added to
# the registry automatically
GLOBAL_BENCHMARK_REGISTRY = set()
# Environment variable that determines whether benchmarks are written.
# See also tensorflow/core/util/reporter.h TestReporter::kTestReporterEnv.
TEST_REPORTER_TEST_ENV = "TEST_REPORT_FILE_PREFIX"
# Environment variable that lets the TensorFlow runtime allocate a new
# threadpool for each benchmark.
OVERRIDE_GLOBAL_THREADPOOL = "TF_OVERRIDE_GLOBAL_THREADPOOL"
def _rename_function(f, arg_num, name):
"""Rename the given function's name appears in the stack trace."""
func_code = f.__code__
new_code = func_code.replace(co_argcount=arg_num, co_name=name)
return types.FunctionType(new_code, f.__globals__, name, f.__defaults__,
f.__closure__)
def _global_report_benchmark(
name, iters=None, cpu_time=None, wall_time=None,
throughput=None, extras=None, metrics=None):
"""Method for recording a benchmark directly.
Args:
name: The BenchmarkEntry name.
iters: (optional) How many iterations were run
cpu_time: (optional) Total cpu time in seconds
wall_time: (optional) Total wall time in seconds
throughput: (optional) Throughput (in MB/s)
extras: (optional) Dict mapping string keys to additional benchmark info.
metrics: (optional) A list of dict representing metrics generated by the
benchmark. Each dict should contain keys 'name' and'value'. A dict
can optionally contain keys 'min_value' and 'max_value'.
Raises:
TypeError: if extras is not a dict.
IOError: if the benchmark output file already exists.
"""
logging.info("Benchmark [%s] iters: %d, wall_time: %g, cpu_time: %g,"
"throughput: %g, extras: %s, metrics: %s", name,
iters if iters is not None else -1,
wall_time if wall_time is not None else -1,
cpu_time if cpu_time is not None else -1,
throughput if throughput is not None else -1,
str(extras) if extras else "None",
str(metrics) if metrics else "None")
entries = test_log_pb2.BenchmarkEntries()
entry = entries.entry.add()
entry.name = name
if iters is not None:
entry.iters = iters
if cpu_time is not None:
entry.cpu_time = cpu_time
if wall_time is not None:
entry.wall_time = wall_time
if throughput is not None:
entry.throughput = throughput
if extras is not None:
if not isinstance(extras, dict):
raise TypeError("extras must be a dict")
for (k, v) in extras.items():
if isinstance(v, numbers.Number):
entry.extras[k].double_value = v
else:
entry.extras[k].string_value = str(v)
if metrics is not None:
if not isinstance(metrics, list):
raise TypeError("metrics must be a list")
for metric in metrics:
if "name" not in metric:
raise TypeError("metric must has a 'name' field")
if "value" not in metric:
raise TypeError("metric must has a 'value' field")
metric_entry = entry.metrics.add()
metric_entry.name = metric["name"]
metric_entry.value = metric["value"]
if "min_value" in metric:
metric_entry.min_value.value = metric["min_value"]
if "max_value" in metric:
metric_entry.max_value.value = metric["max_value"]
test_env = os.environ.get(TEST_REPORTER_TEST_ENV, None)
if test_env is None:
# Reporting was not requested, just print the proto
print(str(entries))
return
serialized_entry = entries.SerializeToString()
mangled_name = name.replace("/", "__")
output_path = "%s%s" % (test_env, mangled_name)
if gfile.Exists(output_path):
raise IOError("File already exists: %s" % output_path)
with gfile.GFile(output_path, "wb") as out:
out.write(serialized_entry)
class _BenchmarkRegistrar(type):
"""The Benchmark class registrar. Used by abstract Benchmark class."""
def __new__(mcs, clsname, base, attrs):
newclass = type.__new__(mcs, clsname, base, attrs)
if not newclass.is_abstract():
GLOBAL_BENCHMARK_REGISTRY.add(newclass)
return newclass
@tf_export("__internal__.test.ParameterizedBenchmark", v1=[])
class ParameterizedBenchmark(_BenchmarkRegistrar):
"""Metaclass to generate parameterized benchmarks.
Use this class as a metaclass and override the `_benchmark_parameters` to
generate multiple benchmark test cases. For example:
class FooBenchmark(metaclass=tf.test.ParameterizedBenchmark,
tf.test.Benchmark):
# The `_benchmark_parameters` is expected to be a list with test cases.
# Each of the test case is a tuple, with the first time to be test case
# name, followed by any number of the parameters needed for the test case.
_benchmark_parameters = [
('case_1', Foo, 1, 'one'),
('case_2', Bar, 2, 'two'),
]
def benchmark_test(self, target_class, int_param, string_param):
# benchmark test body
The example above will generate two benchmark test cases:
"benchmark_test__case_1" and "benchmark_test__case_2".
"""
def __new__(mcs, clsname, base, attrs):
param_config_list = attrs["_benchmark_parameters"]
def create_benchmark_function(original_benchmark, params):
return lambda self: original_benchmark(self, *params)
for name in attrs.copy().keys():
if not name.startswith("benchmark"):
continue
original_benchmark = attrs[name]
del attrs[name]
for param_config in param_config_list:
test_name_suffix = param_config[0]
params = param_config[1:]
benchmark_name = name + "__" + test_name_suffix
if benchmark_name in attrs:
raise Exception(
"Benchmark named {} already defined.".format(benchmark_name))
benchmark = create_benchmark_function(original_benchmark, params)
# Renaming is important because `report_benchmark` function looks up the
# function name in the stack trace.
attrs[benchmark_name] = _rename_function(benchmark, 1, benchmark_name)
return super().__new__(mcs, clsname, base, attrs)
class Benchmark(metaclass=_BenchmarkRegistrar):
"""Abstract class that provides helper functions for running benchmarks.
Any class subclassing this one is immediately registered in the global
benchmark registry.
Only methods whose names start with the word "benchmark" will be run during
benchmarking.
"""
@classmethod
def is_abstract(cls):
# mro: (_BenchmarkRegistrar, Benchmark) means this is Benchmark
return len(cls.mro()) <= 2
def _get_name(self, overwrite_name=None):
"""Returns full name of class and method calling report_benchmark."""
# Find the caller method (outermost Benchmark class)
stack = tf_inspect.stack()
calling_class = None
name = None
for frame in stack[::-1]:
f_locals = frame[0].f_locals
f_self = f_locals.get("self", None)
if isinstance(f_self, Benchmark):
calling_class = f_self # Get the outermost stack Benchmark call
name = frame[3] # Get the method name
break
if calling_class is None:
raise ValueError("Unable to determine calling Benchmark class.")
# Use the method name, or overwrite_name is provided.
name = overwrite_name or name
# Prefix the name with the class name.
class_name = type(calling_class).__name__
name = "%s.%s" % (class_name, name)
return name
def report_benchmark(
self,
iters=None,
cpu_time=None,
wall_time=None,
throughput=None,
extras=None,
name=None,
metrics=None):
"""Report a benchmark.
Args:
iters: (optional) How many iterations were run
cpu_time: (optional) Median or mean cpu time in seconds.
wall_time: (optional) Median or mean wall time in seconds.
throughput: (optional) Throughput (in MB/s)
extras: (optional) Dict mapping string keys to additional benchmark info.
Values may be either floats or values that are convertible to strings.
name: (optional) Override the BenchmarkEntry name with `name`.
Otherwise it is inferred from the top-level method name.
metrics: (optional) A list of dict, where each dict has the keys below
name (required), string, metric name
value (required), double, metric value
min_value (optional), double, minimum acceptable metric value
max_value (optional), double, maximum acceptable metric value
"""
name = self._get_name(overwrite_name=name)
_global_report_benchmark(
name=name, iters=iters, cpu_time=cpu_time, wall_time=wall_time,
throughput=throughput, extras=extras, metrics=metrics)
@tf_export("test.benchmark_config")
def benchmark_config():
"""Returns a tf.compat.v1.ConfigProto for disabling the dependency optimizer.
Returns:
A TensorFlow ConfigProto object.
"""
config = config_pb2.ConfigProto()
config.graph_options.rewrite_options.dependency_optimization = (
rewriter_config_pb2.RewriterConfig.OFF)
return config
@tf_export("test.Benchmark")
class TensorFlowBenchmark(Benchmark):
"""Abstract class that provides helpers for TensorFlow benchmarks."""
def __init__(self):
# Allow TensorFlow runtime to allocate a new threadpool with different
# number of threads for each new benchmark.
os.environ[OVERRIDE_GLOBAL_THREADPOOL] = "1"
super().__init__()
@classmethod
def is_abstract(cls):
# mro: (_BenchmarkRegistrar, Benchmark, TensorFlowBenchmark) means
# this is TensorFlowBenchmark.
return len(cls.mro()) <= 3
def run_op_benchmark(self,
sess,
op_or_tensor,
feed_dict=None,
burn_iters=2,
min_iters=10,
store_trace=False,
store_memory_usage=True,
name=None,
extras=None,
mbs=0):
"""Run an op or tensor in the given session. Report the results.
Args:
sess: `Session` object to use for timing.
op_or_tensor: `Operation` or `Tensor` to benchmark.
feed_dict: A `dict` of values to feed for each op iteration (see the
`feed_dict` parameter of `Session.run`).
burn_iters: Number of burn-in iterations to run.
min_iters: Minimum number of iterations to use for timing.
store_trace: Boolean, whether to run an extra untimed iteration and
store the trace of iteration in returned extras.
The trace will be stored as a string in Google Chrome trace format
in the extras field "full_trace_chrome_format". Note that trace
will not be stored in test_log_pb2.TestResults proto.
store_memory_usage: Boolean, whether to run an extra untimed iteration,
calculate memory usage, and store that in extras fields.
name: (optional) Override the BenchmarkEntry name with `name`.
Otherwise it is inferred from the top-level method name.
extras: (optional) Dict mapping string keys to additional benchmark info.
Values may be either floats or values that are convertible to strings.
mbs: (optional) The number of megabytes moved by this op, used to
calculate the ops throughput.
Returns:
A `dict` containing the key-value pairs that were passed to
`report_benchmark`. If `store_trace` option is used, then
`full_chrome_trace_format` will be included in return dictionary even
though it is not passed to `report_benchmark` with `extras`.
"""
for _ in range(burn_iters):
sess.run(op_or_tensor, feed_dict=feed_dict)
deltas = [None] * min_iters
for i in range(min_iters):
start_time = time.time()
sess.run(op_or_tensor, feed_dict=feed_dict)
end_time = time.time()
delta = end_time - start_time
deltas[i] = delta
extras = extras if extras is not None else {}
unreported_extras = {}
if store_trace or store_memory_usage:
run_options = config_pb2.RunOptions(
trace_level=config_pb2.RunOptions.FULL_TRACE)
run_metadata = config_pb2.RunMetadata()
sess.run(op_or_tensor, feed_dict=feed_dict,
options=run_options, run_metadata=run_metadata)
tl = timeline.Timeline(run_metadata.step_stats)
if store_trace:
unreported_extras["full_trace_chrome_format"] = (
tl.generate_chrome_trace_format())
if store_memory_usage:
step_stats_analysis = tl.analyze_step_stats(show_memory=True)
allocator_maximums = step_stats_analysis.allocator_maximums
for k, v in allocator_maximums.items():
extras["allocator_maximum_num_bytes_%s" % k] = v.num_bytes
def _median(x):
if not x:
return -1
s = sorted(x)
l = len(x)
lm1 = l - 1
return (s[l//2] + s[lm1//2]) / 2.0
def _mean_and_stdev(x):
if not x:
return -1, -1
l = len(x)
mean = sum(x) / l
if l == 1:
return mean, -1
variance = sum([(e - mean) * (e - mean) for e in x]) / (l - 1)
return mean, math.sqrt(variance)
median_delta = _median(deltas)
benchmark_values = {
"iters": min_iters,
"wall_time": median_delta,
"extras": extras,
"name": name,
"throughput": mbs / median_delta
}
self.report_benchmark(**benchmark_values)
mean_delta, stdev_delta = _mean_and_stdev(deltas)
unreported_extras["wall_time_mean"] = mean_delta
unreported_extras["wall_time_stdev"] = stdev_delta
benchmark_values["extras"].update(unreported_extras)
return benchmark_values
def evaluate(self, tensors):
"""Evaluates tensors and returns numpy values.
Args:
tensors: A Tensor or a nested list/tuple of Tensors.
Returns:
tensors numpy values.
"""
sess = ops.get_default_session() or self.cached_session()
return sess.run(tensors)
def _run_benchmarks(regex):
"""Run benchmarks that match regex `regex`.
This function goes through the global benchmark registry, and matches
benchmark class and method names of the form
`module.name.BenchmarkClass.benchmarkMethod` to the given regex.
If a method matches, it is run.
Args:
regex: The string regular expression to match Benchmark classes against.
Raises:
ValueError: If no benchmarks were selected by the input regex.
"""
registry = list(GLOBAL_BENCHMARK_REGISTRY)
selected_benchmarks = []
# Match benchmarks in registry against regex
for benchmark in registry:
benchmark_name = "%s.%s" % (benchmark.__module__, benchmark.__name__)
attrs = dir(benchmark)
# Don't instantiate the benchmark class unless necessary
benchmark_instance = None
for attr in attrs:
if not attr.startswith("benchmark"):
continue
candidate_benchmark_fn = getattr(benchmark, attr)
if not callable(candidate_benchmark_fn):
continue
full_benchmark_name = "%s.%s" % (benchmark_name, attr)
if regex == "all" or re.search(regex, full_benchmark_name):
selected_benchmarks.append(full_benchmark_name)
# Instantiate the class if it hasn't been instantiated
benchmark_instance = benchmark_instance or benchmark()
# Get the method tied to the class
instance_benchmark_fn = getattr(benchmark_instance, attr)
# Call the instance method
instance_benchmark_fn()
if not selected_benchmarks:
raise ValueError("No benchmarks matched the pattern: '{}'".format(regex))
def benchmarks_main(true_main, argv=None):
"""Run benchmarks as declared in argv.
Args:
true_main: True main function to run if benchmarks are not requested.
argv: the command line arguments (if None, uses sys.argv).
"""
if argv is None:
argv = sys.argv
found_arg = [
arg
for arg in argv
if arg.startswith("--benchmark_filter=")
or arg.startswith("-benchmark_filter=")
]
if found_arg:
# Remove --benchmark_filter arg from sys.argv
argv.remove(found_arg[0])
regex = found_arg[0].split("=")[1]
app.run(lambda _: _run_benchmarks(regex), argv=argv)
else:
true_main()
tensorflow/python/platform/benchmark_test.py
+78
View File
@@ -0,0 +1,78 @@
# Copyright 2019 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.
# ==============================================================================
"""Test for the tf.test.benchmark."""
import os
from google.protobuf import json_format
from tensorflow.core.util import test_log_pb2
from tensorflow.python.platform import benchmark
from tensorflow.python.platform import test
class BenchmarkTest(test.TestCase, benchmark.TensorFlowBenchmark):
def testReportBenchmark(self):
output_dir = self.get_temp_dir() + os.path.sep
os.environ['TEST_REPORT_FILE_PREFIX'] = output_dir
proto_file_path = os.path.join(output_dir,
'BenchmarkTest.testReportBenchmark')
if os.path.exists(proto_file_path):
os.remove(proto_file_path)
self.report_benchmark(
iters=2000,
wall_time=1000,
name='testReportBenchmark',
metrics=[{'name': 'metric_name_1', 'value': 0, 'min_value': 1},
{'name': 'metric_name_2', 'value': 90, 'min_value': 0,
'max_value': 95}])
with open(proto_file_path, 'rb') as f:
benchmark_entries = test_log_pb2.BenchmarkEntries()
benchmark_entries.ParseFromString(f.read())
actual_result = json_format.MessageToDict(
benchmark_entries, preserving_proto_field_name=True,
always_print_fields_with_no_presence=True)['entry'][0]
os.remove(proto_file_path)
expected_result = {
'name': 'BenchmarkTest.testReportBenchmark',
# google.protobuf.json_format.MessageToDict() will convert
# int64 field to string.
'iters': '2000',
'wall_time': 1000,
'cpu_time': 0,
'throughput': 0,
'extras': {},
'metrics': [
{
'name': 'metric_name_1',
'value': 0,
'min_value': 1
},
{
'name': 'metric_name_2',
'value': 90,
'min_value': 0,
'max_value': 95
}
]
}
self.assertEqual(2000, benchmark_entries.entry[0].iters)
self.assertDictEqual(expected_result, actual_result)
if __name__ == '__main__':
test.main()
tensorflow/python/platform/build_info_test.py
+47
View File
@@ -0,0 +1,47 @@
# 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.
# ==============================================================================
"""Test for the generated build_info script."""
import platform
from tensorflow.python.platform import build_info
from tensorflow.python.platform import test
class BuildInfoTest(test.TestCase):
def testBuildInfo(self):
self.assertEqual(build_info.build_info['is_rocm_build'],
test.is_built_with_rocm())
self.assertEqual(build_info.build_info['is_cuda_build'],
test.is_built_with_cuda())
# TODO(b/173044576): make the test work for Windows.
if platform.system() != 'Windows':
# pylint: disable=g-import-not-at-top
from tensorflow.compiler.tf2tensorrt._pywrap_py_utils import is_tensorrt_enabled
self.assertEqual(build_info.build_info['is_tensorrt_build'],
is_tensorrt_enabled())
def testDeterministicOrder(self):
# The dict may contain other keys depending on the platform, but the ones
# it always contains should be in order.
self.assertContainsSubsequence(
build_info.build_info.keys(),
('is_cuda_build', 'is_rocm_build', 'is_tensorrt_build'))
if __name__ == '__main__':
test.main()
tensorflow/python/platform/control_imports.py
+27
View File
@@ -0,0 +1,27 @@
# 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.
# ==============================================================================
"""Switch between Google or open source dependencies."""
# Switch between Google and OSS dependencies
USE_OSS = True
# Per-dependency switches determining whether each dependency is ready
# to be replaced by its OSS equivalence.
# TODO(danmane,mrry,opensource): Flip these switches, then remove them
OSS_APP = True
OSS_FLAGS = True
OSS_GFILE = True
OSS_GOOGLETEST = True
OSS_LOGGING = True
OSS_PARAMETERIZED = True
tensorflow/python/platform/cpu_feature_guard_wrapper.cc
+22
View File
@@ -0,0 +1,22 @@
/* 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.
==============================================================================*/
#include "pybind11/pybind11.h" // from @pybind11
#include "tensorflow/core/platform/cpu_feature_guard.h"
PYBIND11_MODULE(_pywrap_cpu_feature_guard, m) {
m.def("InfoAboutUnusedCPUFeatures",
&tensorflow::port::InfoAboutUnusedCPUFeatures);
};
tensorflow/python/platform/device_context.py
+18
View File
@@ -0,0 +1,18 @@
# Copyright 2019 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.
# ==============================================================================
"""Helpers to get device context."""
def enclosing_tpu_context():
pass
tensorflow/python/platform/enable_tf2.cc
+22
View File
@@ -0,0 +1,22 @@
/* 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.
==============================================================================*/
#include "pybind11/pybind11.h" // from @pybind11
#include "tensorflow/core/platform/enable_tf2_utils.h"
PYBIND11_MODULE(_pywrap_tf2, m) {
m.def("enable", &tensorflow::set_tf2_execution);
m.def("is_enabled", &tensorflow::tf2_execution_enabled);
}
tensorflow/python/platform/flags.py
+121
View File
@@ -0,0 +1,121 @@
# 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.
# ==============================================================================
"""Import router for absl.flags. See https://github.com/abseil/abseil-py."""
import logging as _logging
import sys as _sys
# go/tf-wildcard-import
from absl.flags import * # pylint: disable=wildcard-import
from tensorflow.python.util import tf_decorator
# Since we wrap absl.flags DEFINE functions, we need to declare this module
# does not affect key flags.
disclaim_key_flags() # pylint: disable=undefined-variable
_RENAMED_ARGUMENTS = {
'flag_name': 'name',
'default_value': 'default',
'docstring': 'help',
}
def _wrap_define_function(original_function):
"""Wraps absl.flags's define functions so tf.flags accepts old names."""
def wrapper(*args, **kwargs):
"""Wrapper function that turns old keyword names to new ones."""
has_old_names = False
for old_name, new_name in _RENAMED_ARGUMENTS.items():
if old_name in kwargs:
has_old_names = True
value = kwargs.pop(old_name)
kwargs[new_name] = value
if has_old_names:
_logging.warning(
'Use of the keyword argument names (flag_name, default_value, '
'docstring) is deprecated, please use (name, default, help) instead.')
return original_function(*args, **kwargs)
return tf_decorator.make_decorator(original_function, wrapper)
class _FlagValuesWrapper:
"""Wrapper class for absl.flags.FLAGS.
The difference is that tf.flags.FLAGS implicitly parses flags with sys.argv
when accessing the FLAGS values before it's explicitly parsed,
while absl.flags.FLAGS raises an exception.
"""
def __init__(self, flags_object):
self.__dict__['__wrapped'] = flags_object
def __getattribute__(self, name):
if name == '__dict__':
return super().__getattribute__(name)
return self.__dict__['__wrapped'].__getattribute__(name)
def __getattr__(self, name):
wrapped = self.__dict__['__wrapped']
# To maintain backwards compatibility, implicitly parse flags when reading
# a flag.
if not wrapped.is_parsed():
wrapped(_sys.argv)
return wrapped.__getattr__(name)
def __setattr__(self, name, value):
return self.__dict__['__wrapped'].__setattr__(name, value)
def __delattr__(self, name):
return self.__dict__['__wrapped'].__delattr__(name)
def __dir__(self):
return self.__dict__['__wrapped'].__dir__()
def __getitem__(self, name):
return self.__dict__['__wrapped'].__getitem__(name)
def __setitem__(self, name, flag):
return self.__dict__['__wrapped'].__setitem__(name, flag)
def __len__(self):
return self.__dict__['__wrapped'].__len__()
def __iter__(self):
return self.__dict__['__wrapped'].__iter__()
def __str__(self):
return self.__dict__['__wrapped'].__str__()
def __call__(self, *args, **kwargs):
return self.__dict__['__wrapped'].__call__(*args, **kwargs)
# pylint: disable=invalid-name,used-before-assignment
# absl.flags APIs use `default` as the name of the default value argument.
# Allow the following functions continue to accept `default_value`.
DEFINE_string = _wrap_define_function(DEFINE_string)
DEFINE_boolean = _wrap_define_function(DEFINE_boolean)
DEFINE_bool = DEFINE_boolean
DEFINE_float = _wrap_define_function(DEFINE_float)
DEFINE_integer = _wrap_define_function(DEFINE_integer)
# pylint: enable=invalid-name,used-before-assignment
FLAGS = _FlagValuesWrapper(FLAGS) # pylint: disable=used-before-assignment
tensorflow/python/platform/flags_test.py
+132
View File
@@ -0,0 +1,132 @@
# 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.
# ==============================================================================
"""Sanity tests for tf.flags."""
import sys
import unittest
from absl import flags as absl_flags
from tensorflow.python.platform import flags
from tensorflow.python.platform import test
flags.DEFINE_string( # pylint: disable=no-value-for-parameter
flag_name='old_string',
default_value='default',
docstring='docstring')
flags.DEFINE_string(
name='new_string', default='default', help='docstring')
flags.DEFINE_integer( # pylint: disable=no-value-for-parameter
flag_name='old_integer',
default_value=1,
docstring='docstring')
flags.DEFINE_integer(
name='new_integer', default=1, help='docstring')
flags.DEFINE_float( # pylint: disable=no-value-for-parameter
flag_name='old_float',
default_value=1.5,
docstring='docstring')
flags.DEFINE_float(
name='new_float', default=1.5, help='docstring')
flags.DEFINE_bool( # pylint: disable=no-value-for-parameter
flag_name='old_bool',
default_value=True,
docstring='docstring')
flags.DEFINE_bool(
name='new_bool', default=True, help='docstring')
flags.DEFINE_boolean( # pylint: disable=no-value-for-parameter
flag_name='old_boolean',
default_value=False,
docstring='docstring')
flags.DEFINE_boolean(
name='new_boolean', default=False, help='docstring')
class FlagsTest(unittest.TestCase):
def setUp(self):
self.original_flags = flags.FlagValues()
self.wrapped_flags = flags._FlagValuesWrapper(self.original_flags)
flags.DEFINE_string(
'test', 'default', 'test flag', flag_values=self.wrapped_flags)
def test_attribute_overrides(self):
# Test that methods defined in absl.flags.FlagValues are the same as the
# wrapped ones.
self.assertEqual(flags.FLAGS.is_parsed, absl_flags.FLAGS.is_parsed)
def test_getattr(self):
self.assertFalse(self.wrapped_flags.is_parsed())
with test.mock.patch.object(sys, 'argv', new=['program', '--test=new']):
self.assertEqual('new', self.wrapped_flags.test)
self.assertTrue(self.wrapped_flags.is_parsed())
def test_setattr(self):
self.assertEqual('default', self.wrapped_flags.test)
self.wrapped_flags.test = 'new'
self.assertEqual('new', self.wrapped_flags.test)
def test_delattr(self):
del self.wrapped_flags.test
self.assertNotIn('test', self.wrapped_flags)
with self.assertRaises(AttributeError):
_ = self.wrapped_flags.test
def test_dir(self):
self.assertEqual(['test'], dir(self.wrapped_flags))
def test_getitem(self):
self.assertIs(self.original_flags['test'], self.wrapped_flags['test'])
def test_setitem(self):
flag = flags.Flag(flags.ArgumentParser(), flags.ArgumentSerializer(),
'fruit', 'apple', 'the fruit type')
self.wrapped_flags['fruit'] = flag
self.assertIs(self.original_flags['fruit'], self.wrapped_flags['fruit'])
self.assertEqual('apple', self.wrapped_flags.fruit)
def test_len(self):
self.assertEqual(1, len(self.wrapped_flags))
def test_iter(self):
self.assertEqual(['test'], list(self.wrapped_flags))
def test_str(self):
self.assertEqual(str(self.wrapped_flags), str(self.original_flags))
def test_call(self):
self.wrapped_flags(['program', '--test=new'])
self.assertEqual('new', self.wrapped_flags.test)
def test_keyword_arguments(self):
test_cases = (
('old_string', 'default'),
('new_string', 'default'),
('old_integer', 1),
('new_integer', 1),
('old_float', 1.5),
('new_float', 1.5),
('old_bool', True),
('new_bool', True),
('old_boolean', False),
('new_boolean', False),
)
for flag_name, default_value in test_cases:
self.assertEqual(default_value, absl_flags.FLAGS[flag_name].default)
self.assertEqual('docstring', absl_flags.FLAGS[flag_name].help)
if __name__ == '__main__':
unittest.main()
tensorflow/python/platform/gfile.py
+135
View File
@@ -0,0 +1,135 @@
# 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.
# ==============================================================================
"""Import router for file_io."""
# pylint: disable=unused-import
from tensorflow.python.lib.io.file_io import copy as Copy
from tensorflow.python.lib.io.file_io import create_dir as MkDir
from tensorflow.python.lib.io.file_io import delete_file as Remove
from tensorflow.python.lib.io.file_io import delete_recursively as DeleteRecursively
from tensorflow.python.lib.io.file_io import file_exists as Exists
from tensorflow.python.lib.io.file_io import FileIO as _FileIO
from tensorflow.python.lib.io.file_io import get_matching_files as Glob
from tensorflow.python.lib.io.file_io import is_directory as IsDirectory
from tensorflow.python.lib.io.file_io import list_directory as ListDirectory
from tensorflow.python.lib.io.file_io import recursive_create_dir as MakeDirs
from tensorflow.python.lib.io.file_io import rename as Rename
from tensorflow.python.lib.io.file_io import stat as Stat
from tensorflow.python.lib.io.file_io import walk as Walk
# pylint: enable=unused-import
from tensorflow.python.util.deprecation import deprecated
from tensorflow.python.util.tf_export import tf_export
@tf_export('io.gfile.GFile', v1=['gfile.GFile', 'gfile.Open', 'io.gfile.GFile'])
class GFile(_FileIO):
r"""File I/O wrappers without thread locking.
The main roles of the `tf.io.gfile` module are:
1. To provide an API that is close to Python's file I/O objects, and
2. To provide an implementation based on TensorFlow's C++ FileSystem API.
The C++ FileSystem API supports multiple file system implementations,
including local files, Google Cloud Storage (using a `gs://` prefix, and
HDFS (using an `hdfs://` prefix). TensorFlow exports these as `tf.io.gfile`,
so that you can use these implementations for saving and loading checkpoints,
writing to TensorBoard logs, and accessing training data (among other uses).
However, if all your files are local, you can use the regular Python file
API without any problem.
*Note*: though similar to Python's I/O implementation, there are semantic
differences to make `tf.io.gfile` more efficient for backing filesystems. For
example, a write mode file will not be opened until the first write call to
minimize RPC invocations in network filesystems.
Once you obtain a `GFile` object, you can use it in most ways as you would any
Python's file object:
>>> with open("/tmp/x", "w") as f:
... f.write("asdf")
4
>>> with tf.io.gfile.GFile("/tmp/x") as f:
... f.read()
'asdf'
The difference is that you can specify URI schemes to use other filesystems
(e.g., `gs://` for GCS, `s3://` for S3, etc.), if they are supported. Using
`file://` as an example, we have:
>>> with tf.io.gfile.GFile("file:///tmp/x", "w") as f:
... f.write("qwert")
... f.write("asdf")
>>> tf.io.gfile.GFile("file:///tmp/x").read()
'qwertasdf'
You can also read all lines of a file directly:
>>> with tf.io.gfile.GFile("file:///tmp/x", "w") as f:
... f.write("asdf\n")
... f.write("qwer\n")
>>> tf.io.gfile.GFile("/tmp/x").readlines()
['asdf\n', 'qwer\n']
You can iterate over the lines:
>>> with tf.io.gfile.GFile("file:///tmp/x", "w") as f:
... f.write("asdf\n")
... f.write("qwer\n")
>>> for line in tf.io.gfile.GFile("/tmp/x"):
... print(line[:-1]) # removes the end of line character
asdf
qwer
Random access read is possible if the underlying filesystem supports it:
>>> with open("/tmp/x", "w") as f:
... f.write("asdfqwer")
>>> f = tf.io.gfile.GFile("/tmp/x")
>>> f.read(3)
'asd'
>>> f.seek(4)
>>> f.tell()
4
>>> f.read(3)
'qwe'
>>> f.tell()
7
>>> f.close()
"""
def __init__(self, name, mode='r'):
super(GFile, self).__init__(name=name, mode=mode)
@tf_export(v1=['gfile.FastGFile'])
class FastGFile(_FileIO):
"""File I/O wrappers without thread locking.
Note, that this is somewhat like builtin Python file I/O, but
there are semantic differences to make it more efficient for
some backing filesystems. For example, a write mode file will
not be opened until the first write call (to minimize RPC
invocations in network filesystems).
"""
@deprecated(None, 'Use tf.gfile.GFile.')
def __init__(self, name, mode='r'):
super(FastGFile, self).__init__(name=name, mode=mode)
# Does not alias to Open so that we use our version of GFile to strip
# 'b' mode.
Open = GFile
tensorflow/python/platform/googletest.py
+273
View File
@@ -0,0 +1,273 @@
# 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.
# ==============================================================================
"""Imports absltest as a replacement for testing.pybase.googletest."""
import atexit
import os
import sys
import tempfile
# go/tf-wildcard-import
# pylint: disable=wildcard-import,redefined-builtin
from absl import app
from absl.testing.absltest import *
# pylint: enable=wildcard-import,redefined-builtin
from tensorflow.python.framework import errors
from tensorflow.python.lib.io import file_io
from tensorflow.python.platform import benchmark
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.util import tf_decorator
from tensorflow.python.util import tf_inspect
from tensorflow.python.util.tf_export import tf_export
Benchmark = benchmark.TensorFlowBenchmark # pylint: disable=invalid-name
absltest_main = main
# We keep a global variable in this module to make sure we create the temporary
# directory only once per test binary invocation.
_googletest_temp_dir = ''
# pylint: disable=invalid-name
# pylint: disable=undefined-variable
def g_main(argv):
"""Delegate to absltest.main."""
absltest_main(argv=argv)
# Redefine main to allow running benchmarks
def main(argv=None): # pylint: disable=function-redefined
def main_wrapper():
args = argv
if args is None:
args = sys.argv
return app.run(main=g_main, argv=args)
benchmark.benchmarks_main(true_main=main_wrapper, argv=argv)
def GetTempDir():
"""Return a temporary directory for tests to use."""
global _googletest_temp_dir
if not _googletest_temp_dir:
if os.environ.get('TEST_TMPDIR'):
temp_dir = tempfile.mkdtemp(prefix=os.environ['TEST_TMPDIR'])
else:
first_frame = tf_inspect.stack()[-1][0]
temp_dir = os.path.join(tempfile.gettempdir(),
os.path.basename(tf_inspect.getfile(first_frame)))
temp_dir = tempfile.mkdtemp(prefix=temp_dir.rstrip('.py'))
# Make sure we have the correct path separators.
temp_dir = temp_dir.replace('/', os.sep)
def delete_temp_dir(dirname=temp_dir):
try:
file_io.delete_recursively(dirname)
except errors.OpError as e:
logging.error('Error removing %s: %s', dirname, e)
atexit.register(delete_temp_dir)
_googletest_temp_dir = temp_dir
return _googletest_temp_dir
def test_src_dir_path(relative_path):
"""Creates an absolute test srcdir path given a relative path.
Args:
relative_path: a path relative to tensorflow root.
e.g. "contrib/session_bundle/example".
Returns:
An absolute path to the linked in runfiles.
"""
return os.path.join(
os.environ['TEST_SRCDIR'],
os.environ['TEST_WORKSPACE'],
'tensorflow',
relative_path,
)
def StatefulSessionAvailable():
return False
@tf_export(v1=['test.StubOutForTesting'])
class StubOutForTesting(object):
"""Support class for stubbing methods out for unit testing.
Sample Usage:
You want os.path.exists() to always return true during testing.
stubs = StubOutForTesting()
stubs.Set(os.path, 'exists', lambda x: 1)
...
stubs.CleanUp()
The above changes os.path.exists into a lambda that returns 1. Once
the ... part of the code finishes, the CleanUp() looks up the old
value of os.path.exists and restores it.
"""
def __init__(self):
self.cache = []
self.stubs = []
def __del__(self):
"""Do not rely on the destructor to undo your stubs.
You cannot guarantee exactly when the destructor will get called without
relying on implementation details of a Python VM that may change.
"""
self.CleanUp()
# __enter__ and __exit__ allow use as a context manager.
def __enter__(self):
return self
def __exit__(self, unused_exc_type, unused_exc_value, unused_tb):
self.CleanUp()
def CleanUp(self):
"""Undoes all SmartSet() & Set() calls, restoring original definitions."""
self.SmartUnsetAll()
self.UnsetAll()
def SmartSet(self, obj, attr_name, new_attr):
"""Replace obj.attr_name with new_attr.
This method is smart and works at the module, class, and instance level
while preserving proper inheritance. It will not stub out C types however
unless that has been explicitly allowed by the type.
This method supports the case where attr_name is a staticmethod or a
classmethod of obj.
Notes:
- If obj is an instance, then it is its class that will actually be
stubbed. Note that the method Set() does not do that: if obj is
an instance, it (and not its class) will be stubbed.
- The stubbing is using the builtin getattr and setattr. So, the __get__
and __set__ will be called when stubbing (TODO: A better idea would
probably be to manipulate obj.__dict__ instead of getattr() and
setattr()).
Args:
obj: The object whose attributes we want to modify.
attr_name: The name of the attribute to modify.
new_attr: The new value for the attribute.
Raises:
AttributeError: If the attribute cannot be found.
"""
_, obj = tf_decorator.unwrap(obj)
if (tf_inspect.ismodule(obj) or
(not tf_inspect.isclass(obj) and attr_name in obj.__dict__)):
orig_obj = obj
orig_attr = getattr(obj, attr_name)
else:
if not tf_inspect.isclass(obj):
mro = list(tf_inspect.getmro(obj.__class__))
else:
mro = list(tf_inspect.getmro(obj))
mro.reverse()
orig_attr = None
found_attr = False
for cls in mro:
try:
orig_obj = cls
orig_attr = getattr(obj, attr_name)
found_attr = True
except AttributeError:
continue
if not found_attr:
raise AttributeError('Attribute not found.')
# Calling getattr() on a staticmethod transforms it to a 'normal' function.
# We need to ensure that we put it back as a staticmethod.
old_attribute = obj.__dict__.get(attr_name)
if old_attribute is not None and isinstance(old_attribute, staticmethod):
orig_attr = staticmethod(orig_attr)
self.stubs.append((orig_obj, attr_name, orig_attr))
setattr(orig_obj, attr_name, new_attr)
def SmartUnsetAll(self):
"""Reverses SmartSet() calls, restoring things to original definitions.
This method is automatically called when the StubOutForTesting()
object is deleted; there is no need to call it explicitly.
It is okay to call SmartUnsetAll() repeatedly, as later calls have
no effect if no SmartSet() calls have been made.
"""
for args in reversed(self.stubs):
setattr(*args)
self.stubs = []
def Set(self, parent, child_name, new_child):
"""In parent, replace child_name's old definition with new_child.
The parent could be a module when the child is a function at
module scope. Or the parent could be a class when a class' method
is being replaced. The named child is set to new_child, while the
prior definition is saved away for later, when UnsetAll() is
called.
This method supports the case where child_name is a staticmethod or a
classmethod of parent.
Args:
parent: The context in which the attribute child_name is to be changed.
child_name: The name of the attribute to change.
new_child: The new value of the attribute.
"""
old_child = getattr(parent, child_name)
old_attribute = parent.__dict__.get(child_name)
if old_attribute is not None and isinstance(old_attribute, staticmethod):
old_child = staticmethod(old_child)
self.cache.append((parent, old_child, child_name))
setattr(parent, child_name, new_child)
def UnsetAll(self):
"""Reverses Set() calls, restoring things to their original definitions.
This method is automatically called when the StubOutForTesting()
object is deleted; there is no need to call it explicitly.
It is okay to call UnsetAll() repeatedly, as later calls have no
effect if no Set() calls have been made.
"""
# Undo calls to Set() in reverse order, in case Set() was called on the
# same arguments repeatedly (want the original call to be last one undone)
for (parent, old_child, child_name) in reversed(self.cache):
setattr(parent, child_name, old_child)
self.cache = []
tensorflow/python/platform/logging_test.py
+28
View File
@@ -0,0 +1,28 @@
# 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.
# ==============================================================================
from tensorflow.python.platform import googletest
from tensorflow.python.platform import tf_logging as logging
class EventLoaderTest(googletest.TestCase):
def test_log(self):
# Just check that logging works without raising an exception.
logging.error("test log message")
if __name__ == "__main__":
googletest.main()
tensorflow/python/platform/parameterized.py
+17
View File
@@ -0,0 +1,17 @@
# 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.
# ==============================================================================
"""Extension to unittest to run parameterized tests."""
raise ImportError("Not implemented yet.")
tensorflow/python/platform/remote_utils.py
+32
View File
@@ -0,0 +1,32 @@
# Copyright 2019 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.
# ==============================================================================
"""Platform-specific helpers for connecting to remote servers."""
def get_default_communication_protocol():
return 'grpc'
def is_remote_path(_):
return False
def get_appendable_file_encoding():
return ''
def coordination_service_type(*args, **kwargs):
del args, kwargs
return None
tensorflow/python/platform/resource_loader.py
+132
View File
@@ -0,0 +1,132 @@
# 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.
# ==============================================================================
"""Resource management library."""
import os as _os
import sys as _sys
from tensorflow.python.util import tf_inspect as _inspect
from tensorflow.python.util.tf_export import tf_export
# pylint: disable=g-import-not-at-top
try:
from python.runfiles import runfiles
except ImportError:
runfiles = None
# pylint: enable=g-import-not-at-top
@tf_export(v1=['resource_loader.load_resource'])
def load_resource(path):
"""Load the resource at given path, where path is relative to tensorflow/.
Args:
path: a string resource path relative to tensorflow/.
Returns:
The contents of that resource.
Raises:
IOError: If the path is not found, or the resource can't be opened.
"""
with open(get_path_to_datafile(path), 'rb') as f:
return f.read()
# pylint: disable=protected-access
@tf_export(v1=['resource_loader.get_data_files_path'])
def get_data_files_path():
"""Get a direct path to the data files colocated with the script.
Returns:
The directory where files specified in data attribute of py_test
and py_binary are stored.
"""
return _os.path.dirname(_inspect.getfile(_sys._getframe(1)))
@tf_export(v1=['resource_loader.get_root_dir_with_all_resources'])
def get_root_dir_with_all_resources():
"""Get a root directory containing all the data attributes in the build rule.
Returns:
The path to the specified file present in the data attribute of py_test
or py_binary. Falls back to returning the same as get_data_files_path if it
fails to detect a bazel runfiles directory.
"""
script_dir = get_data_files_path()
# Create a history of the paths, because the data files are located relative
# to the repository root directory, which is directly under runfiles
# directory.
directories = [script_dir]
data_files_dir = ''
while True:
candidate_dir = directories[-1]
current_directory = _os.path.basename(candidate_dir)
if '.runfiles' in current_directory:
# Our file should never be directly under runfiles.
# If the history has only one item, it means we are directly inside the
# runfiles directory, something is wrong, fall back to the default return
# value, script directory.
if len(directories) > 1:
data_files_dir = directories[-2]
break
else:
new_candidate_dir = _os.path.dirname(candidate_dir)
# If we are at the root directory these two will be the same.
if new_candidate_dir == candidate_dir:
break
else:
directories.append(new_candidate_dir)
return data_files_dir or script_dir
@tf_export(v1=['resource_loader.get_path_to_datafile'])
def get_path_to_datafile(path):
"""Get the path to the specified file in the data dependencies.
The path is relative to tensorflow/
Args:
path: a string resource path relative to tensorflow/
Returns:
The path to the specified file present in the data attribute of py_test
or py_binary.
Raises:
IOError: If the path is not found, or the resource can't be opened.
"""
# First, try finding in the new path.
if runfiles:
r = runfiles.Create()
new_fpath = r.Rlocation(
_os.path.abspath(_os.path.join('tensorflow', path)))
if new_fpath is not None and _os.path.exists(new_fpath):
return new_fpath
# Then, the old style path, as people became dependent on this buggy call.
old_filepath = _os.path.join(
_os.path.dirname(_inspect.getfile(_sys._getframe(1))), path)
return old_filepath
@tf_export(v1=['resource_loader.readahead_file_path'])
def readahead_file_path(path, readahead='128M'): # pylint: disable=unused-argument
"""Readahead files not implemented; simply returns given path."""
return path
tensorflow/python/platform/resource_loader_test.py
+33
View File
@@ -0,0 +1,33 @@
# 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.
# ==============================================================================
from tensorflow.python.platform import googletest
from tensorflow.python.platform import resource_loader
class ResourceLoaderTest(googletest.TestCase):
def test_exception(self):
with self.assertRaises(IOError):
resource_loader.load_resource("/fake/file/path/dne")
def test_exists(self):
contents = resource_loader.load_resource(
"python/platform/resource_loader.py")
self.assertIn(b"tensorflow", contents)
if __name__ == "__main__":
googletest.main()
tensorflow/python/platform/self_check.py
+71
View File
@@ -0,0 +1,71 @@
# 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.
# ==============================================================================
"""Platform-specific code for checking the integrity of the TensorFlow build."""
import os
MSVCP_DLL_NAMES = "msvcp_dll_names"
try:
from tensorflow.python.platform import build_info
except ImportError as exc:
raise ImportError(
"Could not import tensorflow. Do not import"
" tensorflow from its source directory; change directory to outside the"
" TensorFlow source tree, and relaunch your Python interpreter from"
" there."
) from exc
def preload_check():
"""Raises an exception if the environment is not correctly configured.
Raises:
ImportError: If the check detects that the environment is not correctly
configured, and attempting to load the TensorFlow runtime will fail.
"""
if os.name == "nt":
# Attempt to load any DLLs that the Python extension depends on before
# we load the Python extension, so that we can raise an actionable error
# message if they are not found.
if MSVCP_DLL_NAMES in build_info.build_info:
missing = []
for dll_name in build_info.build_info[MSVCP_DLL_NAMES].split(","):
found = False
for path_dir in os.environ.get("PATH", "").split(";"):
if os.path.isfile(os.path.join(path_dir, dll_name)):
found = True
break
if not found:
missing.append(dll_name)
if missing:
raise ImportError(
"Could not find the DLL(s) %r. TensorFlow requires that these DLLs "
"be installed in a directory that is named in your %%PATH%% "
"environment variable. You may install these DLLs by downloading "
'"Microsoft C++ Redistributable for Visual Studio 2015, 2017 and '
'2019" for your platform from this URL: '
"https://support.microsoft.com/help/2977003/the-latest-supported-visual-c-downloads"
% " or ".join(missing)
)
else:
# Load a library that performs CPU feature guard checking. Doing this here
# as a preload check makes it more likely that we detect any CPU feature
# incompatibilities before we trigger them (which would typically result in
# SIGILL).
from tensorflow.python.platform import _pywrap_cpu_feature_guard
_pywrap_cpu_feature_guard.InfoAboutUnusedCPUFeatures()
tensorflow/python/platform/stacktrace_handler_test.py
+77
View File
@@ -0,0 +1,77 @@
# Copyright 2018 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.
# =============================================================================
"""Test to make sure stack trace is generated in case of test failures."""
import argparse
import os
import signal
import subprocess
import sys
from tensorflow.python.platform import test
from tensorflow.python.platform import tf_logging as logging
# FLAGS defined at the bottom:
# child (bool) set to true if we are running in the child process.
FLAGS = None
_CHILD_FLAG_HELP = 'Boolean. Set to true if this is the child process.'
class StacktraceHandlerTest(test.TestCase):
def testChildProcessKillsItself(self):
if FLAGS.child:
os.kill(os.getpid(), signal.SIGABRT)
def testGeneratesStacktrace(self):
if FLAGS.child:
return
# Subprocess sys.argv[0] with --child=True
if sys.executable:
child_process = subprocess.Popen(
[sys.executable, sys.argv[0], '--child=True'], cwd=os.getcwd(),
stdout=subprocess.PIPE, stderr=subprocess.PIPE)
else:
child_process = subprocess.Popen(
[sys.argv[0], '--child=True'], cwd=os.getcwd(),
stdout=subprocess.PIPE, stderr=subprocess.PIPE)
# Capture its output. capture both stdout and stderr and append them.
# We are not worried about timing or order of messages in this test.
child_stdout, child_stderr = child_process.communicate()
child_output = child_stdout + child_stderr
# Make sure the child process is dead before we proceed.
child_process.wait()
logging.info('Output from the child process:')
logging.info(child_output)
# Verify a stack trace is printed.
self.assertIn(b'PyEval_EvalFrame', child_output)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument(
'--child', type=bool, default=False, help=_CHILD_FLAG_HELP)
FLAGS, unparsed = parser.parse_known_args()
# Now update argv, so that unittest library does not get confused.
sys.argv = [sys.argv[0]] + unparsed
test.main()
tensorflow/python/platform/stacktrace_handler_wrapper.cc
+22
View File
@@ -0,0 +1,22 @@
/* Copyright 2019 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.
==============================================================================*/
#include "pybind11/pybind11.h" // from @pybind11
#include "tensorflow/core/platform/stacktrace_handler.h"
PYBIND11_MODULE(_pywrap_stacktrace_handler, m) {
m.def("InstallStacktraceHandler",
&tensorflow::testing::InstallStacktraceHandler);
};
tensorflow/python/platform/sysconfig.py
+141
View File
@@ -0,0 +1,141 @@
# 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.
# ==============================================================================
"""System configuration library.
API docstring: tensorflow.sysconfig
"""
import os.path as _os_path
import platform as _platform
from tensorflow.python.client import pywrap_tf_session
# pylint: disable=g-importing-member
from tensorflow.python.framework.versions import CXX11_ABI_FLAG as _CXX11_ABI_FLAG
from tensorflow.python.framework.versions import CXX_VERSION as _CXX_VERSION
from tensorflow.python.framework.versions import MONOLITHIC_BUILD as _MONOLITHIC_BUILD
from tensorflow.python.framework.versions import VERSION as _VERSION
from tensorflow.python.platform import build_info
from tensorflow.python.util.tf_export import tf_export
# pylint: disable=g-import-not-at-top
@tf_export('sysconfig.get_include')
def get_include():
"""Get the directory containing the TensorFlow C++ header files.
Returns:
The directory as string.
"""
return _os_path.join(get_lib(), 'include')
@tf_export('sysconfig.get_lib')
def get_lib():
"""Get the directory containing the TensorFlow framework library.
Returns:
The directory as string.
"""
return _os_path.normpath(
_os_path.join(_os_path.dirname(__file__), '..', '..')
)
@tf_export('sysconfig.get_compile_flags')
def get_compile_flags():
"""Returns the compilation flags for compiling with TensorFlow.
The returned list of arguments can be passed to the compiler for compiling
against TensorFlow headers. The result is platform dependent.
For example, on a typical Linux system with Python 3.7 the following command
prints `['-I/usr/local/lib/python3.7/dist-packages/tensorflow/include',
'-D_GLIBCXX_USE_CXX11_ABI=1', '-DEIGEN_MAX_ALIGN_BYTES=64']`
>>> print(tf.sysconfig.get_compile_flags())
Returns:
A list of strings for the compiler flags.
"""
flags = []
flags.append('-I%s' % get_include())
flags.append('-D_GLIBCXX_USE_CXX11_ABI=%d' % _CXX11_ABI_FLAG)
cxx_version_flag = None
if _CXX_VERSION == 201103:
cxx_version_flag = '--std=c++11'
elif _CXX_VERSION == 201402:
cxx_version_flag = '--std=c++14'
elif _CXX_VERSION == 201703:
cxx_version_flag = '--std=c++17'
elif _CXX_VERSION == 202002:
cxx_version_flag = '--std=c++20'
if cxx_version_flag:
flags.append(cxx_version_flag)
flags.append('-DEIGEN_MAX_ALIGN_BYTES=%d' %
pywrap_tf_session.get_eigen_max_align_bytes())
return flags
@tf_export('sysconfig.get_link_flags')
def get_link_flags():
"""Returns the linker flags for linking with TensorFlow.
The returned list of arguments can be passed to the linker for linking against
TensorFlow. The result is platform dependent.
For example, on a typical Linux system with Python 3.7 the following command
prints `['-L/usr/local/lib/python3.7/dist-packages/tensorflow',
'-l:libtensorflow_framework.so.2']`
>>> print(tf.sysconfig.get_link_flags())
Returns:
A list of strings for the linker flags.
"""
is_mac = _platform.system() == 'Darwin'
ver = _VERSION.split('.')[0]
flags = []
if not _MONOLITHIC_BUILD:
flags.append('-L%s' % get_lib())
if is_mac:
flags.append('-ltensorflow_framework.%s' % ver)
else:
flags.append('-l:libtensorflow_framework.so.%s' % ver)
return flags
@tf_export('sysconfig.get_build_info')
def get_build_info():
"""Get a dictionary describing TensorFlow's build environment.
Values are generated when TensorFlow is compiled, and are static for each
TensorFlow package. The return value is a dictionary with string keys such as:
- cuda_version
- cudnn_version
- is_cuda_build
- is_rocm_build
- msvcp_dll_names
- nvcuda_dll_name
- cudart_dll_name
- cudnn_dll_name
Note that the actual keys and values returned by this function is subject to
change across different versions of TensorFlow or across platforms.
Returns:
A Dictionary describing TensorFlow's build environment.
"""
return build_info.build_info
tensorflow/python/platform/sysconfig_test.py
+49
View File
@@ -0,0 +1,49 @@
# Copyright 2019 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 re
from tensorflow.python.platform import googletest
from tensorflow.python.platform import sysconfig as sysconfig_lib
from tensorflow.python.platform import test
class SysconfigTest(googletest.TestCase):
def test_get_build_info_works(self):
build_info = sysconfig_lib.get_build_info()
self.assertIsInstance(build_info, dict)
def test_rocm_cuda_info_matches(self):
build_info = sysconfig_lib.get_build_info()
self.assertEqual(build_info["is_rocm_build"], test.is_built_with_rocm())
self.assertEqual(build_info["is_cuda_build"], test.is_built_with_cuda())
def test_compile_flags(self):
# Must contain an include directory, and define _GLIBCXX_USE_CXX11_ABI,
# EIGEN_MAX_ALIGN_BYTES
compile_flags = sysconfig_lib.get_compile_flags()
def list_contains(items, regex_str):
regex = re.compile(regex_str)
return any(regex.match(item) for item in items)
self.assertTrue(list_contains(compile_flags, ".*/include"))
self.assertTrue(list_contains(compile_flags, ".*_GLIBCXX_USE_CXX11_ABI.*"))
self.assertTrue(list_contains(compile_flags, ".*EIGEN_MAX_ALIGN_BYTES.*"))
self.assertTrue(list_contains(compile_flags, ".*std.*"))
if __name__ == "__main__":
googletest.main()
tensorflow/python/platform/test.py
+220
View File
@@ -0,0 +1,220 @@
# 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.
# ==============================================================================
"""Testing."""
import functools
from unittest import mock
# pylint: disable=g-bad-import-order
from tensorflow.python.framework import test_util as _test_util
from tensorflow.python.platform import googletest as _googletest
# pylint: disable=unused-import
from tensorflow.python.framework.test_util import assert_equal_graph_def
from tensorflow.python.framework.test_util import create_local_cluster
from tensorflow.python.framework.test_util import TensorFlowTestCase as TestCase
from tensorflow.python.framework.test_util import gpu_device_name
from tensorflow.python.framework.test_util import is_gpu_available
from tensorflow.python.ops.gradient_checker import compute_gradient_error
from tensorflow.python.ops.gradient_checker import compute_gradient
# pylint: enable=unused-import,g-bad-import-order
from tensorflow.python.util.tf_export import tf_export
tf_export(v1=['test.mock'])(mock)
# Import Benchmark class
Benchmark = _googletest.Benchmark # pylint: disable=invalid-name
# Import StubOutForTesting class
StubOutForTesting = _googletest.StubOutForTesting # pylint: disable=invalid-name
@tf_export('test.main')
def main(argv=None):
"""Runs all unit tests."""
_test_util.InstallStackTraceHandler()
return _googletest.main(argv)
@tf_export(v1=['test.get_temp_dir'])
def get_temp_dir():
"""Returns a temporary directory for use during tests.
There is no need to delete the directory after the test.
@compatibility(TF2)
This function is removed in TF2. Please use `TestCase.get_temp_dir` instead
in a test case.
Outside of a unit test, obtain a temporary directory through Python's
`tempfile` module.
@end_compatibility
Returns:
The temporary directory.
"""
return _googletest.GetTempDir()
@tf_export(v1=['test.test_src_dir_path'])
def test_src_dir_path(relative_path):
"""Creates an absolute test srcdir path given a relative path.
Args:
relative_path: a path relative to tensorflow root.
e.g. "core/platform".
Returns:
An absolute path to the linked in runfiles.
"""
return _googletest.test_src_dir_path(relative_path)
@tf_export('test.is_built_with_cuda')
def is_built_with_cuda():
"""Returns whether TensorFlow was built with CUDA (GPU) support.
This method should only be used in tests written with `tf.test.TestCase`. A
typical usage is to skip tests that should only run with CUDA (GPU).
>>> class MyTest(tf.test.TestCase):
...
... def test_add_on_gpu(self):
... if not tf.test.is_built_with_cuda():
... self.skipTest("test is only applicable on GPU")
...
... with tf.device("GPU:0"):
... self.assertEqual(tf.math.add(1.0, 2.0), 3.0)
TensorFlow official binary is built with CUDA.
"""
return _test_util.IsGoogleCudaEnabled()
@tf_export('test.is_built_with_rocm')
def is_built_with_rocm():
"""Returns whether TensorFlow was built with ROCm (GPU) support.
This method should only be used in tests written with `tf.test.TestCase`. A
typical usage is to skip tests that should only run with ROCm (GPU).
>>> class MyTest(tf.test.TestCase):
...
... def test_add_on_gpu(self):
... if not tf.test.is_built_with_rocm():
... self.skipTest("test is only applicable on GPU")
...
... with tf.device("GPU:0"):
... self.assertEqual(tf.math.add(1.0, 2.0), 3.0)
TensorFlow official binary is NOT built with ROCm.
"""
return _test_util.IsBuiltWithROCm()
@tf_export('test.disable_with_predicate')
def disable_with_predicate(pred, skip_message):
"""Disables the test if pred is true."""
def decorator_disable_with_predicate(func):
@functools.wraps(func)
def wrapper_disable_with_predicate(self, *args, **kwargs):
if pred():
self.skipTest(skip_message)
else:
return func(self, *args, **kwargs)
return wrapper_disable_with_predicate
return decorator_disable_with_predicate
@tf_export('test.is_built_with_gpu_support')
def is_built_with_gpu_support():
"""Returns whether TensorFlow was built with GPU (CUDA or ROCm) support.
This method should only be used in tests written with `tf.test.TestCase`. A
typical usage is to skip tests that should only run with GPU.
>>> class MyTest(tf.test.TestCase):
...
... def test_add_on_gpu(self):
... if not tf.test.is_built_with_gpu_support():
... self.skipTest("test is only applicable on GPU")
...
... with tf.device("GPU:0"):
... self.assertEqual(tf.math.add(1.0, 2.0), 3.0)
TensorFlow official binary is built with CUDA GPU support.
"""
return is_built_with_cuda() or is_built_with_rocm()
@tf_export('test.is_built_with_xla')
def is_built_with_xla():
"""Returns whether TensorFlow was built with XLA support.
This method should only be used in tests written with `tf.test.TestCase`. A
typical usage is to skip tests that should only run with XLA.
>>> class MyTest(tf.test.TestCase):
...
... def test_add_on_xla(self):
... if not tf.test.is_built_with_xla():
... self.skipTest("test is only applicable on XLA")
... @tf.function(jit_compile=True)
... def add(x, y):
... return tf.math.add(x, y)
...
... self.assertEqual(add(tf.ones(()), tf.ones(())), 2.0)
TensorFlow official binary is built with XLA.
"""
return _test_util.IsBuiltWithXLA()
@tf_export('test.is_cpu_target_available')
def is_cpu_target_available(target):
"""Indicates whether TensorFlow was built with support for a given CPU target.
Args:
target: The name of the CPU target whose support to check for.
Returns:
A boolean indicating whether TensorFlow was built with support for the
given CPU target.
This method should only be used in tests written with `tf.test.TestCase`. A
typical usage is to skip tests that should only run with a given target.
>>> class MyTest(tf.test.TestCase):
...
... def test_add_on_aarch64(self):
... if not tf.test.is_cpu_target_available('aarch64'):
... self.skipTest("test is only applicable on AArch64")
... @tf.function(jit_compile=True)
... def add(x, y):
... return tf.math.add(x, y)
...
... self.assertEqual(add(tf.ones(()), tf.ones(())), 2.0)
"""
return _test_util.IsCPUTargetAvailable(target)
tensorflow/python/platform/tf_logging.py
+372
View File
@@ -0,0 +1,372 @@
# 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.
# ==============================================================================
"""Logging utilities."""
# pylint: disable=unused-import
# pylint: disable=g-bad-import-order
# pylint: disable=invalid-name
import logging as _logging
import os as _os
import sys as _sys
import _thread
import time as _time
import traceback as _traceback
from logging import DEBUG
from logging import ERROR
from logging import FATAL
from logging import INFO
from logging import WARN
import threading
from tensorflow.python.util.tf_export import tf_export
# Don't use this directly. Use get_logger() instead.
_logger = None
_logger_lock = threading.Lock()
def error_log(error_msg, level=ERROR):
"""Empty helper method."""
del error_msg, level
def _get_caller(offset=3):
"""Returns a code and frame object for the lowest non-logging stack frame."""
# Use sys._getframe(). This avoids creating a traceback object.
# pylint: disable=protected-access
f = _sys._getframe(offset)
# pylint: enable=protected-access
our_file = f.f_code.co_filename
f = f.f_back
while f:
code = f.f_code
if code.co_filename != our_file:
return code, f
f = f.f_back
return None, None
# The definition of `findCaller` changed in Python 3.2,
# and further changed in Python 3.8
if _sys.version_info.major >= 3 and _sys.version_info.minor >= 8:
def _logger_find_caller(stack_info=False, stacklevel=1): # pylint: disable=g-wrong-blank-lines
code, frame = _get_caller(4)
sinfo = None
if stack_info:
sinfo = '\n'.join(_traceback.format_stack())
if code:
return (code.co_filename, frame.f_lineno, code.co_name, sinfo)
else:
return '(unknown file)', 0, '(unknown function)', sinfo
elif _sys.version_info.major >= 3 and _sys.version_info.minor >= 2:
def _logger_find_caller(stack_info=False): # pylint: disable=g-wrong-blank-lines
code, frame = _get_caller(4)
sinfo = None
if stack_info:
sinfo = '\n'.join(_traceback.format_stack())
if code:
return (code.co_filename, frame.f_lineno, code.co_name, sinfo)
else:
return '(unknown file)', 0, '(unknown function)', sinfo
else:
def _logger_find_caller(): # pylint: disable=g-wrong-blank-lines
code, frame = _get_caller(4)
if code:
return (code.co_filename, frame.f_lineno, code.co_name)
else:
return '(unknown file)', 0, '(unknown function)'
@tf_export('get_logger')
def get_logger():
"""Return TF logger instance.
Returns:
An instance of the Python logging library Logger.
See Python documentation (https://docs.python.org/3/library/logging.html)
for detailed API. Below is only a summary.
The logger has 5 levels of logging from the most serious to the least:
1. FATAL
2. ERROR
3. WARN
4. INFO
5. DEBUG
The logger has the following methods, based on these logging levels:
1. fatal(msg, *args, **kwargs)
2. error(msg, *args, **kwargs)
3. warn(msg, *args, **kwargs)
4. info(msg, *args, **kwargs)
5. debug(msg, *args, **kwargs)
The `msg` can contain string formatting. An example of logging at the `ERROR`
level
using string formatting is:
>>> tf.get_logger().error("The value %d is invalid.", 3)
You can also specify the logging verbosity. In this case, the
WARN level log will not be emitted:
>>> tf.get_logger().setLevel(ERROR)
>>> tf.get_logger().warn("This is a warning.")
"""
global _logger
# Use double-checked locking to avoid taking lock unnecessarily.
if _logger:
return _logger
_logger_lock.acquire()
try:
if _logger:
return _logger
# Scope the TensorFlow logger to not conflict with users' loggers.
logger = _logging.getLogger('tensorflow')
# Override findCaller on the logger to skip internal helper functions
logger.findCaller = _logger_find_caller
# Don't further configure the TensorFlow logger if the root logger is
# already configured. This prevents double logging in those cases.
if not _logging.getLogger().handlers:
# Determine whether we are in an interactive environment
_interactive = False
try:
# This is only defined in interactive shells.
if _sys.ps1:
_interactive = True
except AttributeError:
# Even now, we may be in an interactive shell with `python -i`.
_interactive = _sys.flags.interactive
# If we are in an interactive environment (like Jupyter), set loglevel
# to INFO and pipe the output to stdout.
if _interactive:
logger.setLevel(INFO)
_logging_target = _sys.stdout
else:
_logging_target = _sys.stderr
# Add the output handler.
_handler = _logging.StreamHandler(_logging_target)
_handler.setFormatter(_logging.Formatter(_logging.BASIC_FORMAT, None))
logger.addHandler(_handler)
_logger = logger
return _logger
finally:
_logger_lock.release()
@tf_export(v1=['logging.log'])
def log(level, msg, *args, **kwargs):
get_logger().log(level, msg, *args, **kwargs)
@tf_export(v1=['logging.debug'])
def debug(msg, *args, **kwargs):
get_logger().debug(msg, *args, **kwargs)
@tf_export(v1=['logging.error'])
def error(msg, *args, **kwargs):
get_logger().error(msg, *args, **kwargs)
@tf_export(v1=['logging.fatal'])
def fatal(msg, *args, **kwargs):
get_logger().fatal(msg, *args, **kwargs)
@tf_export(v1=['logging.info'])
def info(msg, *args, **kwargs):
get_logger().info(msg, *args, **kwargs)
@tf_export(v1=['logging.warn'])
def warn(msg, *args, **kwargs):
get_logger().warning(msg, *args, **kwargs)
@tf_export(v1=['logging.warning'])
def warning(msg, *args, **kwargs):
get_logger().warning(msg, *args, **kwargs)
_level_names = {
FATAL: 'FATAL',
ERROR: 'ERROR',
WARN: 'WARN',
INFO: 'INFO',
DEBUG: 'DEBUG',
}
# Mask to convert integer thread ids to unsigned quantities for logging
# purposes
_THREAD_ID_MASK = 2 * _sys.maxsize + 1
_log_prefix = None # later set to google2_log_prefix
# Counter to keep track of number of log entries per token.
_log_counter_per_token = {}
@tf_export(v1=['logging.TaskLevelStatusMessage'])
def TaskLevelStatusMessage(msg):
error(msg)
@tf_export(v1=['logging.flush'])
def flush():
raise NotImplementedError()
# Code below is taken from pyglib/logging
@tf_export(v1=['logging.vlog'])
def vlog(level, msg, *args, **kwargs):
get_logger().log(level, msg, *args, **kwargs)
def _GetNextLogCountPerToken(token):
"""Wrapper for _log_counter_per_token.
Args:
token: The token for which to look up the count.
Returns:
The number of times this function has been called with
*token* as an argument (starting at 0)
"""
global _log_counter_per_token # pylint: disable=global-variable-not-assigned
_log_counter_per_token[token] = 1 + _log_counter_per_token.get(token, -1)
return _log_counter_per_token[token]
@tf_export(v1=['logging.log_every_n'])
def log_every_n(level, msg, n, *args):
"""Log 'msg % args' at level 'level' once per 'n' times.
Logs the 1st call, (N+1)st call, (2N+1)st call, etc.
Not threadsafe.
Args:
level: The level at which to log.
msg: The message to be logged.
n: The number of times this should be called before it is logged.
*args: The args to be substituted into the msg.
"""
count = _GetNextLogCountPerToken(_GetFileAndLine())
log_if(level, msg, not (count % n), *args)
@tf_export(v1=['logging.log_first_n'])
def log_first_n(level, msg, n, *args): # pylint: disable=g-bad-name
"""Log 'msg % args' at level 'level' only first 'n' times.
Not threadsafe.
Args:
level: The level at which to log.
msg: The message to be logged.
n: The number of times this should be called before it is logged.
*args: The args to be substituted into the msg.
"""
count = _GetNextLogCountPerToken(_GetFileAndLine())
log_if(level, msg, count < n, *args)
@tf_export(v1=['logging.log_if'])
def log_if(level, msg, condition, *args):
"""Log 'msg % args' at level 'level' only if condition is fulfilled."""
if condition:
vlog(level, msg, *args)
def _GetFileAndLine():
"""Returns (filename, linenumber) for the stack frame."""
code, f = _get_caller()
if not code:
return ('<unknown>', 0)
return (code.co_filename, f.f_lineno)
def google2_log_prefix(level, timestamp=None, file_and_line=None):
"""Assemble a logline prefix using the google2 format."""
# pylint: disable=global-variable-not-assigned
global _level_names
# pylint: enable=global-variable-not-assigned
# Record current time
now = timestamp or _time.time()
now_tuple = _time.localtime(now)
now_microsecond = int(1e6 * (now % 1.0))
(filename, line) = file_and_line or _GetFileAndLine()
basename = _os.path.basename(filename)
# Severity string
severity = 'I'
if level in _level_names:
severity = _level_names[level][0]
s = '%c%02d%02d %02d:%02d:%02d.%06d %5d %s:%d] ' % (
severity,
now_tuple[1], # month
now_tuple[2], # day
now_tuple[3], # hour
now_tuple[4], # min
now_tuple[5], # sec
now_microsecond,
_get_thread_id(),
basename,
line)
return s
@tf_export(v1=['logging.get_verbosity'])
def get_verbosity():
"""Return how much logging output will be produced."""
return get_logger().getEffectiveLevel()
@tf_export(v1=['logging.set_verbosity'])
def set_verbosity(v):
"""Sets the threshold for what messages will be logged."""
get_logger().setLevel(v)
def _get_thread_id():
"""Get id of current thread, suitable for logging as an unsigned quantity."""
thread_id = _thread.get_ident()
return thread_id & _THREAD_ID_MASK
_log_prefix = google2_log_prefix
tf_export(v1=['logging.DEBUG']).export_constant(__name__, 'DEBUG')
tf_export(v1=['logging.ERROR']).export_constant(__name__, 'ERROR')
tf_export(v1=['logging.FATAL']).export_constant(__name__, 'FATAL')
tf_export(v1=['logging.INFO']).export_constant(__name__, 'INFO')
tf_export(v1=['logging.WARN']).export_constant(__name__, 'WARN')
tensorflow/python/platform/windows_lib_diagnostics.py
+217
View File
@@ -0,0 +1,217 @@
# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# 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.
# ==============================================================================
"""Utility to diagnose Windows DLL loading issues."""
import ctypes
import os
import struct
import sys
# Windows Error Codes
ERROR_ACCESS_DENIED = 0x05
ERROR_PROC_NOT_FOUND = 0x7F
ERROR_MOD_NOT_FOUND = 0x7E
ERROR_BAD_EXE_FORMAT = 0xC1
ERROR_DLL_INIT_FAILED = 0x45A
def get_dll_dependencies(path):
"""Parses PE header to find direct DLL dependencies."""
deps = []
try:
with open(path, "rb") as f:
if f.read(2) != b"MZ":
return []
f.seek(0x3C)
header_offset_raw = f.read(4)
if not header_offset_raw:
return []
pe_offset = struct.unpack("<I", header_offset_raw)[0]
f.seek(pe_offset)
if f.read(4) != b"PE\0\0":
return []
f.seek(2, 1) # machine
num_sections = struct.unpack("<H", f.read(2))[0]
f.seek(12, 1) # stamp etc.
size_of_optional_header = struct.unpack("<H", f.read(2))[0]
f.seek(2, 1) # characteristics
optional_header_start = f.tell()
magic_raw = f.read(2)
if not magic_raw:
return []
magic = struct.unpack("<H", magic_raw)[0]
if magic == 0x20B: # PE32+
import_dir_entry_offset = optional_header_start + 120
elif magic == 0x10B: # PE32
import_dir_entry_offset = optional_header_start + 104
else:
return []
f.seek(import_dir_entry_offset)
import_rva = struct.unpack("<I", f.read(4))[0]
if import_rva == 0:
return []
f.seek(optional_header_start + size_of_optional_header)
sections = []
for _ in range(num_sections):
_ = f.read(8)
vs = struct.unpack("<I", f.read(4))[0]
va = struct.unpack("<I", f.read(4))[0]
_ = struct.unpack("<I", f.read(4))[0]
rp = struct.unpack("<I", f.read(4))[0]
f.seek(16, 1)
sections.append((va, vs, rp))
def rva_to_offset(rva):
for va, vs, ptr in sections:
if va <= rva < va + vs:
return ptr + (rva - va)
return None
import_offset = rva_to_offset(import_rva)
if import_offset is None:
return []
f.seek(import_offset)
while True:
entry = f.read(20)
if not entry or len(entry) < 20:
break
name_rva = struct.unpack("<I", entry[12:16])[0]
if name_rva == 0:
break
current_pos = f.tell()
name_offset = rva_to_offset(name_rva)
if name_offset:
f.seek(name_offset)
name = b""
while True:
char = f.read(1)
if char == b"\0" or not char:
break
name += char
deps.append(name.decode("ascii"))
f.seek(current_pos)
except (IOError, struct.error, UnicodeDecodeError):
# Catch specific exceptions that can occur during file reading and parsing.
pass
return deps
def diagnose_dll_load(path, depth=0, visited=None):
"""Recursively checks DLL dependencies and attempts loading them."""
if visited is None:
visited = set()
if path.lower() in visited:
return
visited.add(path.lower())
indent = " " * depth
if depth == 0:
print(f"\n[TensorFlow DLL Diagnostic] Analyzing: {path}")
# If file doesn't exist, it might be a system DLL
if not os.path.exists(path):
return
deps = get_dll_dependencies(path)
if not deps:
return
target_dir = os.path.dirname(path)
# os.add_dll_directory is available on Windows in Python 3.8+
add_dll_dir = getattr(os, "add_dll_directory", None)
if add_dll_dir is not None and callable(add_dll_dir):
try:
add_dll_dir(target_dir)
except OSError:
pass
for dll_name in deps:
if dll_name.lower() in visited:
continue
# 1. Recurse first (Leaf-first)
if not dll_name.lower().startswith(
("kernel32", "user32", "api-ms-win-", "msvcrt", "ucrtbase")
):
for p in [target_dir] + os.environ["PATH"].split(os.pathsep):
full_p = os.path.join(p, dll_name)
if os.path.exists(full_p):
diagnose_dll_load(full_p, depth + 1, visited)
break
# 2. After recursion, try to load at this level
try:
_ = ctypes.WinDLL(dll_name)
except OSError as e:
err = getattr(e, "winerror", None)
print(f"{indent}[Error] Failed to load {dll_name}: ", end="")
if err == ERROR_MOD_NOT_FOUND:
print(
"MODULE NOT FOUND (0x7E) - Check if the file or its dependencies"
" exist."
)
elif err == ERROR_BAD_EXE_FORMAT:
print("BAD EXE FORMAT (0xC1/193) - Likely a 32-bit vs 64-bit mismatch.")
elif err == ERROR_PROC_NOT_FOUND:
print(
"PROC NOT FOUND (0x7F/127) - Version mismatch in a dependent DLL."
)
elif err == ERROR_ACCESS_DENIED:
print("ACCESS DENIED (0x05) - Check file permissions.")
elif err == ERROR_DLL_INIT_FAILED:
print(
"INITIALIZATION FAILED (0x45A) - The DLL's DllMain returned false."
)
print(
f"{indent} Hint: This often happens if your CPU lacks required"
" instructions (like AVX/AVX2)"
)
print(
f"{indent} or if the Microsoft Visual C++ Redistributable is"
" outdated/missing."
)
else:
print(f"UNKNOWN ERROR ({err}): {e}")
visited.add(dll_name.lower())
def run_diagnosis(path=None):
"""Runs the Windows DLL load diagnosis.
Args:
path: The path to the primary DLL/PYD file to diagnose. If None, it defaults
to '_pywrap_tensorflow_internal.pyd' located relative to this script.
"""
try:
if path is None:
# Determine path relative to this file to avoid importing tensorflow
# This file is typically at tensorflow/python/platform/
platform_dir = os.path.dirname(os.path.abspath(__file__))
python_dir = os.path.dirname(platform_dir)
path = os.path.join(python_dir, "_pywrap_tensorflow_internal.pyd")
if os.name == "nt" and os.path.exists(path):
diagnose_dll_load(path)
else:
print(f"Error: Path does not exist or not on Windows: {path}")
except OSError as e:
print(f"Diagnostic failed: {e}")
if __name__ == "__main__":
if len(sys.argv) > 1:
run_diagnosis(sys.argv[1])
else:
run_diagnosis()