Adding test metrics -- a set of lightweight counters exposed for e.g. C++ and Python unit tests, benchmarks, etc. Also updating unit tests involving single-threaded executor to confirm the proper executor was triggered.
PiperOrigin-RevId: 411118723
Change-Id: I4d879d225647dfb6a31d54bba1ad1a27655b5ec7
diff --git a/tensorflow/core/common_runtime/function.cc b/tensorflow/core/common_runtime/function.cc
index 9ae4492..e34d7ce 100644
--- a/tensorflow/core/common_runtime/function.cc
+++ b/tensorflow/core/common_runtime/function.cc
@@ -35,6 +35,7 @@
#include "tensorflow/core/framework/collective.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/function_handle_cache.h"
+#include "tensorflow/core/framework/metrics.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/op.h"
@@ -1036,6 +1037,12 @@
if (executor_type.empty() && IsSingleThreadedExecutorCompatible(g.get())) {
executor_type = "SINGLE_THREADED_EXECUTOR";
}
+
+ metrics::IncrementTestCounter("flr_executor",
+ (executor_type == "SINGLE_THREADED_EXECUTOR")
+ ? "single_threaded"
+ : "default");
+
TF_RETURN_IF_ERROR(NewExecutor(executor_type, params, *g, &exec));
{
// Guard item since it is already inserted in items_.
diff --git a/tensorflow/core/common_runtime/function_test.cc b/tensorflow/core/common_runtime/function_test.cc
index da0a5a7..d925a31 100644
--- a/tensorflow/core/common_runtime/function_test.cc
+++ b/tensorflow/core/common_runtime/function_test.cc
@@ -38,6 +38,7 @@
#include "tensorflow/core/common_runtime/step_stats_collector.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/function_testlib.h"
+#include "tensorflow/core/framework/metrics.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/op_requires.h"
@@ -448,6 +449,8 @@
TEST_F(FunctionLibraryRuntimeTest, XTimesTwo_ConsumeArgument_DefaultExecutor) {
Init({test::function::XTimesTwo()});
+ auto default_executor = metrics::TestDelta("flr_executor", "default");
+ auto single_threaded = metrics::TestDelta("flr_executor", "single_threaded");
FunctionLibraryRuntime::Handle handle;
TF_CHECK_OK(flr0_->Instantiate(
"XTimesTwo", test::function::Attrs({{"T", DT_FLOAT}}), &handle));
@@ -469,11 +472,15 @@
EXPECT_FALSE(x.IsInitialized());
TF_CHECK_OK(flr0_->ReleaseHandle(handle));
+ EXPECT_GT(default_executor.Get(), 0);
+ EXPECT_EQ(single_threaded.Get(), 0);
}
TEST_F(FunctionLibraryRuntimeTest,
XTimesTwo_ConsumeArgument_SingleThreadedExecutor) {
Init({test::function::XTimesTwo()});
+ auto default_executor = metrics::TestDelta("flr_executor", "default");
+ auto single_threaded = metrics::TestDelta("flr_executor", "single_threaded");
FunctionLibraryRuntime::InstantiateOptions instantiate_opts;
instantiate_opts.executor_type = "SINGLE_THREADED_EXECUTOR";
FunctionLibraryRuntime::Handle handle;
@@ -498,6 +505,8 @@
EXPECT_FALSE(x.IsInitialized());
TF_CHECK_OK(flr0_->ReleaseHandle(handle));
+ EXPECT_EQ(default_executor.Get(), 0);
+ EXPECT_GT(single_threaded.Get(), 0);
}
TEST_F(FunctionLibraryRuntimeTest, XTimesN) {
diff --git a/tensorflow/core/framework/metrics.cc b/tensorflow/core/framework/metrics.cc
index f3c56f7..3386648 100644
--- a/tensorflow/core/framework/metrics.cc
+++ b/tensorflow/core/framework/metrics.cc
@@ -226,6 +226,10 @@
"at the start of a call to TPUExecute. Timer starts at RunGraph "
"invocation and ends when TPUExecute args are ready on the current task.");
+auto* test_counters =
+ monitoring::Counter<2>::New("/tensorflow/core/test_counters",
+ "Counters used for testing.", "name", "label");
+
} // namespace
monitoring::Counter<2>* GetGraphOptimizationCounter() {
@@ -459,5 +463,23 @@
graph_unused_outputs->GetCell(op_name)->IncrementBy(1);
}
+void IncrementTestCounter(const string& name, const string& label) {
+ test_counters->GetCell(name, label)->IncrementBy(1);
+}
+
+const monitoring::CounterCell* TestCounter(const string& name,
+ const string& label) {
+ return test_counters->GetCell(name, label);
+}
+
+TestDelta::TestDelta(const string& name, const string& label)
+ : cell_(TestCounter(name, label)) {
+ Reset();
+}
+
+void TestDelta::Reset() { last_value_ = cell_->value(); }
+
+int64 TestDelta::Get() { return cell_->value() - last_value_; }
+
} // namespace metrics
} // namespace tensorflow
diff --git a/tensorflow/core/framework/metrics.h b/tensorflow/core/framework/metrics.h
index 0c3e3d0..3cdaf43 100644
--- a/tensorflow/core/framework/metrics.h
+++ b/tensorflow/core/framework/metrics.h
@@ -275,6 +275,25 @@
// Updates the metrics stored about time BFC allocator spents during delay.
void UpdateBfcAllocatorDelayTime(const uint64 delay_usecs);
+// Increments (by 1) a simple integer counter that is exposed for testing.
+void IncrementTestCounter(const string& name, const string& label);
+
+// Read-only access to a counter for testing.
+const monitoring::CounterCell* TestCounter(const string& name,
+ const string& label);
+
+// Read-only wrapper for a TestCounter to track increments between calls.
+class TestDelta {
+ public:
+ TestDelta(const string& name, const string& label);
+ void Reset();
+ int64 Get();
+
+ private:
+ const monitoring::CounterCell* cell_;
+ int64 last_value_;
+};
+
} // namespace metrics
} // namespace tensorflow
diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD
index 0783a77..f8e6ee0 100644
--- a/tensorflow/python/BUILD
+++ b/tensorflow/python/BUILD
@@ -208,6 +208,7 @@
"//tensorflow/python/eager:remote",
"//tensorflow/python/framework",
"//tensorflow/python/framework:_pywrap_python_op_gen",
+ "//tensorflow/python/framework:_test_metrics_util",
"//tensorflow/python/framework:combinations",
"//tensorflow/python/framework:config",
"//tensorflow/python/framework:errors",
diff --git a/tensorflow/python/eager/run_eager_op_as_function_test.py b/tensorflow/python/eager/run_eager_op_as_function_test.py
index ccc96cc..52add10 100644
--- a/tensorflow/python/eager/run_eager_op_as_function_test.py
+++ b/tensorflow/python/eager/run_eager_op_as_function_test.py
@@ -177,5 +177,30 @@
cs = critical_section_ops.CriticalSection(shared_name="cs")
cs.execute(lambda: 1.0)
+
+class RunEagerOpAsFunctionInternalsTest(test.TestCase):
+
+ @test_util.enable_eager_op_as_function
+ def testSimpleGraphUsesDefaultExecutor(self):
+ if context.num_gpus():
+ self.skipTest("CPU-only test (requires unpartitioned graph).")
+
+ default_executor = test_util.TestDelta("flr_executor", "default")
+ single_threaded = test_util.TestDelta("flr_executor", "single_threaded")
+ array_ops.fill([2], constant_op.constant(7, dtype=dtypes.int64))
+ assert default_executor.Get() > 0
+ assert single_threaded.Get() == 0
+
+ @test_util.enable_eager_op_as_function
+ def testPartitionedGraphUsesDefaultExecutor(self):
+ if not context.num_gpus():
+ self.skipTest("GPU-only test (requires partitioned graph).")
+
+ default_executor = test_util.TestDelta("flr_executor", "default")
+ single_threaded = test_util.TestDelta("flr_executor", "single_threaded")
+ array_ops.fill([2], constant_op.constant(7, dtype=dtypes.int64))
+ assert default_executor.Get() > 0
+ assert single_threaded.Get() == 0
+
if __name__ == "__main__":
test.main()
diff --git a/tensorflow/python/framework/BUILD b/tensorflow/python/framework/BUILD
index d4bfbc1..1f30fdd 100644
--- a/tensorflow/python/framework/BUILD
+++ b/tensorflow/python/framework/BUILD
@@ -1288,6 +1288,7 @@
"//tensorflow_estimator/python/estimator:__subpackages__",
],
deps = [
+ ":_test_metrics_util",
":errors",
":for_generated_wrappers",
":gpu_util",
@@ -2118,6 +2119,17 @@
],
)
+tf_python_pybind_extension(
+ name = "_test_metrics_util",
+ srcs = ["test_metrics_util.cc"],
+ module_name = "_test_metrics_util",
+ deps = [
+ "//tensorflow/core:framework",
+ "//third_party/python_runtime:headers",
+ "@pybind11",
+ ],
+)
+
# copybara:uncomment_begin(google-only)
# py_proto_library(
# name = "cpp_shape_inference_proto_py_pb2",
diff --git a/tensorflow/python/framework/test_metrics_util.cc b/tensorflow/python/framework/test_metrics_util.cc
new file mode 100644
index 0000000..ffb6fc1
--- /dev/null
+++ b/tensorflow/python/framework/test_metrics_util.cc
@@ -0,0 +1,30 @@
+/* Copyright 2021 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"
+#include "tensorflow/core/framework/metrics.h"
+
+namespace {
+
+// Read an internal TensorFlow counter that has been exposed for testing.
+int64_t test_counter_value(std::string name, std::string label) {
+ return tensorflow::metrics::TestCounter(name, label)->value();
+}
+
+} // anonymous namespace
+
+PYBIND11_MODULE(_test_metrics_util, m) {
+ m.def("test_counter_value", &test_counter_value);
+};
diff --git a/tensorflow/python/framework/test_util.py b/tensorflow/python/framework/test_util.py
index 5619528..fe68a9d 100644
--- a/tensorflow/python/framework/test_util.py
+++ b/tensorflow/python/framework/test_util.py
@@ -49,6 +49,7 @@
from tensorflow.python.eager import context
from tensorflow.python.eager import def_function
from tensorflow.python.eager import tape
+from tensorflow.python.framework import _test_metrics_util
from tensorflow.python.framework import config
from tensorflow.python.framework import device as pydev
from tensorflow.python.framework import dtypes
@@ -3828,3 +3829,20 @@
yield
finally:
def_function.run_functions_eagerly(initial_state)
+
+
+class TestDelta(object):
+ """A utility class to track increments to test counters."""
+
+ def __init__(self, name, label):
+ self.name = name
+ self.label = label
+ self.Reset()
+
+ def Reset(self):
+ self.last_value = _test_metrics_util.test_counter_value(
+ self.name, self.label)
+
+ def Get(self):
+ value = _test_metrics_util.test_counter_value(self.name, self.label)
+ return value - self.last_value
