1 # -*- Python -*- 2 3 load("//tensorflow:tensorflow.bzl", "tf_py_test") 4 5 # Create a benchmark test target of a TensorFlow C++ test (tf_cc_*_test) 6 def tf_cc_logged_benchmark( 7 name=None, 8 target=None, 9 benchmarks="..", 10 tags=[], 11 test_log_output_prefix="", 12 benchmark_type="cpp_microbenchmark"): 13 if not name: 14 fail("Must provide a name") 15 if not target: 16 fail("Must provide a target") 17 if (not ":" in target 18 or not target.startswith("//") 19 or target.endswith(":all") 20 or target.endswith(".")): 21 fail(" ".join(("Target must be a single well-defined test, e.g.,", 22 "//path/to:test. Received: %s" % target))) 23 24 all_tags = ( 25 depset(tags) + depset( 26 ["benchmark-test", "local", "manual", "regression-test"])).to_list() 27 28 tf_py_test( 29 name = name, 30 tags = all_tags, 31 size = "large", 32 srcs = ["//tensorflow/tools/test:run_and_gather_logs"], 33 args = [ 34 "--name=//%s:%s" % (PACKAGE_NAME, name), 35 "--test_name=" + target, 36 "--test_args=--benchmarks=%s" % benchmarks, 37 "--benchmark_type=%s" % benchmark_type, 38 ], 39 data = [ 40 target, 41 ], 42 main = "run_and_gather_logs.py", 43 additional_deps = [ 44 "//tensorflow/tools/test:run_and_gather_logs" 45 ]) 46 47 # Create a benchmark test target of a TensorFlow python test (*py_tests) 48 def tf_py_logged_benchmark( 49 name=None, 50 target=None, 51 benchmarks="..", 52 tags=[], 53 test_log_output_prefix=""): 54 # For now generating a py benchmark is the same as generating a C++ 55 # benchmark target. In the future this may change, so we have 56 # two macros just in case 57 tf_cc_logged_benchmark( 58 name=name, 59 target=target, 60 benchmarks=benchmarks, 61 tags=tags, 62 test_log_output_prefix=test_log_output_prefix, 63 benchmark_type="python_benchmark") 64
