From d9ea5051104b3580fee2d49c94be2ec45012672f Mon Sep 17 00:00:00 2001
From: Jay Shi <xiaojies@google.com>
Date: Thu, 6 Aug 2020 17:32:29 -0700
Subject: [PATCH] [tf.data] Record the number of times tf.data experiment
 applied to input pipelines.

PiperOrigin-RevId: 325345128
Change-Id: I8de30c47f681a6f41e25e5ade4460f20a9d7bb5d
---
 tensorflow/core/framework/metrics.cc              |  9 +++++++++
 tensorflow/core/framework/metrics.h               |  3 +++
 tensorflow/core/kernels/data/dataset_utils.cc     | 11 -----------
 .../core/kernels/data/optimize_dataset_op.cc      | 15 +++++++++++++++
 4 files changed, 27 insertions(+), 11 deletions(-)

diff --git a/tensorflow/core/framework/metrics.cc b/tensorflow/core/framework/metrics.cc
index 8cbfcd5342a..f5aff3a4e11 100644
--- a/tensorflow/core/framework/metrics.cc
+++ b/tensorflow/core/framework/metrics.cc
@@ -80,6 +80,11 @@ auto* tf_data_bytes_fetched_counter = monitoring::Counter<0>::New(
 auto* tf_data_elements_counter = monitoring::Counter<1>::New(
     "/tensorflow/data/elements", "tf.data elements", "name");
 
+auto* tf_data_experiment_counter = monitoring::Counter<1>::New(
+    "/tensorflow/data/experiment",
+    "The number of times tf.data experiment is applied to input pipelines.",
+    "name");
+
 auto* tf_data_fingerprint_counter = monitoring::Counter<1>::New(
     "/tensorflow/data/fingerprint", "tf.data fingerprint", "name");
 
@@ -179,6 +184,10 @@ void RecordTFDataBytesFetched(int64 num_bytes) {
   tf_data_bytes_fetched_counter->GetCell()->IncrementBy(num_bytes);
 }
 
+void RecordTFDataExperiment(const string& name) {
+  tf_data_experiment_counter->GetCell(name)->IncrementBy(1);
+}
+
 void RecordTFDataFingerprint(const string& name) {
   tf_data_fingerprint_counter->GetCell(name)->IncrementBy(1);
 }
diff --git a/tensorflow/core/framework/metrics.h b/tensorflow/core/framework/metrics.h
index 7bc9a1bda0b..f7c90ce593e 100644
--- a/tensorflow/core/framework/metrics.h
+++ b/tensorflow/core/framework/metrics.h
@@ -56,6 +56,9 @@ monitoring::CounterCell* GetTFDataElementsCounter(const string& name);
 // Records the number of bytes fetched from tf.data.Dataset iterator.
 void RecordTFDataBytesFetched(int64 num_bytes);
 
+// Records the number of times tf.data experiment is applied to input pipelines.
+void RecordTFDataExperiment(const string& name);
+
 // Records the time spent in ItertatorResource::GetNext() in microseconds.
 void RecordTFDataGetNextDuration(uint64 duration_us);
 
diff --git a/tensorflow/core/kernels/data/dataset_utils.cc b/tensorflow/core/kernels/data/dataset_utils.cc
index 4151442d747..66de482467d 100644
--- a/tensorflow/core/kernels/data/dataset_utils.cc
+++ b/tensorflow/core/kernels/data/dataset_utils.cc
@@ -1018,17 +1018,6 @@ std::vector<tstring> SelectOptimizations(
     }
   }
 
-  // Log the experiments that will be applied.
-  if (VLOG_IS_ON(1)) {
-    for (auto& pair : live_experiments) {
-      string experiment = pair.first;
-      if (std::find(optimizations_set.begin(), optimizations_set.end(),
-                    experiment) != optimizations_set.end()) {
-        VLOG(1) << "The experiment \"" << experiment << "\" is applied.";
-      }
-    }
-  }
-
   std::vector<tstring> optimizations;
   optimizations.insert(optimizations.end(), optimizations_set.begin(),
                        optimizations_set.end());
diff --git a/tensorflow/core/kernels/data/optimize_dataset_op.cc b/tensorflow/core/kernels/data/optimize_dataset_op.cc
index a0101435794..a566693ec3d 100644
--- a/tensorflow/core/kernels/data/optimize_dataset_op.cc
+++ b/tensorflow/core/kernels/data/optimize_dataset_op.cc
@@ -101,6 +101,21 @@ void OptimizeDatasetOp::MakeDataset(OpKernelContext* ctx, DatasetBase* input,
           job_name, opt_ins_raw, opt_outs_raw, live_experiments,
           optimizations_enabled, optimizations_disabled, optimizations_default,
           hash_func);
+
+      // Log the experiments that will be applied.
+      if (!live_experiments.empty() && VLOG_IS_ON(1)) {
+        VLOG(1) << "The input pipeline is subject to tf.data experiment. "
+                   "Please see `go/tf-data-experiments` for more details.";
+
+        for (auto& pair : live_experiments) {
+          string experiment = pair.first;
+          if (std::find(optimizations.begin(), optimizations.end(),
+                        experiment) != optimizations.end()) {
+            VLOG(1) << "The experiment \"" << experiment << "\" is applied.";
+            metrics::RecordTFDataExperiment(experiment);
+          }
+        }
+      }
     }
   }