57 lines
1.6 KiB
Python
57 lines
1.6 KiB
Python
# -*- Python -*-
|
|
|
|
load("//tensorflow:tensorflow.bzl", "tf_py_test")
|
|
|
|
# Create a benchmark test target of a TensorFlow C++ test (tf_cc_*_test)
|
|
def tf_cc_logged_benchmark(
|
|
name=None,
|
|
target=None,
|
|
benchmarks="..",
|
|
tags=[],
|
|
test_log_output_prefix=""):
|
|
if not name:
|
|
fail("Must provide a name")
|
|
if not target:
|
|
fail("Must provide a target")
|
|
if (not ":" in target
|
|
or not target.startswith("//")
|
|
or target.endswith(":all")
|
|
or target.endswith(".")):
|
|
fail(" ".join(("Target must be a single well-defined test, e.g.,",
|
|
"//path/to:test. Received: %s" % target)))
|
|
|
|
all_tags = list(set(tags) + \
|
|
set(["benchmark-test", "local", "regression-test"]))
|
|
|
|
tf_py_test(
|
|
name = name,
|
|
tags = all_tags,
|
|
srcs = ["//tensorflow/tools/test:run_and_gather_logs.py"],
|
|
args = [
|
|
"--test_name=" + target
|
|
],
|
|
data = [
|
|
target,
|
|
],
|
|
main = "run_and_gather_logs.py",
|
|
additional_deps = [
|
|
"//tensorflow/tools/test:run_and_gather_logs"
|
|
])
|
|
|
|
# Create a benchmark test target of a TensorFlow python test (*py_tests)
|
|
def tf_py_logged_benchmark(
|
|
name=None,
|
|
target=None,
|
|
benchmarks="..",
|
|
tags=[],
|
|
test_log_output_prefix=""):
|
|
# For now generating a py benchmark is the same as generating a C++
|
|
# benchmark target. In the future this may change, so we have
|
|
# two macros just in case
|
|
tf_cc_logged_benchmark(
|
|
name=name,
|
|
target=target,
|
|
benchmarks=benchmarks,
|
|
tags=tags,
|
|
test_log_output_prefix=test_log_output_prefix)
|