Allow TSL CellReader to work with lazy metrics.

PiperOrigin-RevId: 822757884

tensorflow/core/lib/monitoring/cell_reader_test.cc

@@ -85,6 +85,12 @@ auto* test_percentiles_with_labels = monitoring::PercentileSampler<2>::New(
     GetDefaultPercentiles(), /*max_samples=*/1024,
     monitoring::UnitOfMeasure::kTime);
 
+void IncrementLazyCounter() {
+  static auto* test_counter = monitoring::Counter<0>::New(
+      "/tensorflow/monitoring/test/lazy_counter", "Test lazy counter.");
+  test_counter->GetCell()->IncrementBy(1);
+}
+
 TEST(CellReaderTest, CounterDeltaNoLabels) {
   CellReader<int64_t> cell_reader("/tensorflow/monitoring/test/counter");
   EXPECT_EQ(cell_reader.Delta(), 0);
@@ -1249,7 +1255,7 @@ TEST(CellReaderTest, WrongNumberOfLabels) {
                "has 2 labels");
 }
 
-TEST(CellReaderTest, MetricIsNotFound) {
+TEST(CellReaderTest, MetricIsNotFoundRead) {
   CellReader<int64_t> cell_reader("/metric/does/not/exist");
   CellReader<int64_t> empty_cell_reader("");
   EXPECT_DEATH(cell_reader.Read(), "Metric descriptor is not found");
@@ -1294,6 +1300,24 @@ TEST(CellReaderTest, InvalidType) {
 }
 #endif
 
+TEST(CellReaderTest, MetricIsNotFoundDelta) {
+  CellReader<int64_t> cell_reader("/metric/does/not/exist");
+  CellReader<int64_t> empty_cell_reader("");
+  EXPECT_EQ(cell_reader.Delta(), 0);
+  EXPECT_EQ(empty_cell_reader.Delta(), 0);
+}
+
+TEST(CellReaderTest, LazyInitializationDelta) {
+  CellReader<int64_t> cell_reader("/tensorflow/monitoring/test/lazy_counter");
+  EXPECT_EQ(cell_reader.Delta(), 0);
+}
+
+TEST(CellReaderTest, LazyInitializationDeltaAfterIncrement) {
+  CellReader<int64_t> cell_reader("/tensorflow/monitoring/test/lazy_counter");
+  IncrementLazyCounter();
+  EXPECT_EQ(cell_reader.Delta(), 1);
+}
+
 }  // namespace
 }  // namespace testing
 }  // namespace monitoring

third_party/xla/xla/tsl/lib/monitoring/cell_reader-inl.h

@@ -89,18 +89,23 @@ Percentiles GetValue(const Point& point);
 
 // Returns the latest value for `metric_name`, associated with the `labels`. If
 // the metric has not collected any data, it returns a default value appropriate
-// for `ValueType`. If the metric does not exist, or the wrong number of labels
+// for `ValueType`. If a wrong number of labels is provided, it will crash. If
-// is provided, it will crash.
+// the metric does not exist, it will crash if `return_default_on_not_found` is
+// false, otherwise it will return a default value.
 template <typename ValueType>
 ValueType GetLatestValueOrDefault(const CollectedMetrics& metrics,
                                   const std::string& metric_name,
                                   const std::vector<std::string>& labels,
-                                  const ValueType default_value = ValueType()) {
+                                  bool return_default_on_not_found = false,
+                                  ValueType default_value = ValueType()) {
   absl::StatusOr<Point> latest_point =
       GetLatestPoint(metrics, metric_name, labels);
   if (absl::IsUnavailable(latest_point.status())) {
     return std::move(default_value);
   }
+  if (return_default_on_not_found && absl::IsNotFound(latest_point.status())) {
+    return std::move(default_value);
+  }
   if (!latest_point.ok()) {
     LOG(FATAL) << "Failed to read from tfstreamz: " << latest_point.status();
   }

third_party/xla/xla/tsl/lib/monitoring/cell_reader.h

@@ -17,7 +17,6 @@ limitations under the License.
 
 #include <memory>
 #include <string>
-#include <utility>
 #include <vector>
 
 #include "absl/container/flat_hash_map.h"
@@ -81,9 +80,8 @@ class CellReader {
  public:
   // Constructs a `CellReader` that reads values exported for `metric_name`.
   //
-  // REQUIRES: a tfstreamz with `metric_name` exists. Otherwise, the
+  // NOTE: if a tfstreamz with `metric_name` does not exists, the `CellReader`
-  // `CellReader` will construct without issue, but the `Read` and `Delta` calls
+  // will construct without issue, but the `Read` calls will CHECK-fail.
-  // will CHECK-fail.
   explicit CellReader(const std::string& metric_name);
   virtual ~CellReader() = default;
   CellReader(const CellReader&) = delete;
@@ -100,13 +98,13 @@ class CellReader {
 
   // Returns the difference in the value of this cell since the last time
   // `Delta()` was called for this cell, or when the `CellReader` was created,
-  // whichever was most recent. If the metric has not been modified, it returns
+  // whichever was most recent. If tfstreamz does not exist or the metric has
-  // a default value appropriate for `ValueType`. `Delta` is not supported for
+  // not been modified, it returns a default value appropriate for `ValueType`.
-  // string and bool gauges.
+  // `Delta` is not supported for string and bool gauges.
   //
-  // REQUIRES: The tfstreamz exists, `labels` contains a correct number of
+  // REQUIRES: `labels` contains a correct number of labels per tfstreamz
-  // labels per tfstreamz definition, and the ValueType is not string or bool.
+  // definition, and the ValueType is not string or bool. Otherwise, it will
-  // Otherwise, it will CHECK-fail.
+  // CHECK-fail.
   template <typename... LabelType>
   ValueType Delta(const LabelType&... labels);
 
@@ -149,12 +147,14 @@ ValueType CellReader<ValueType>::Delta(const LabelType&... labels) {
   std::vector<std::string> labels_list{labels...};
   std::unique_ptr<CollectedMetrics> metrics = internal::CollectMetrics();
   ValueType value = internal::GetLatestValueOrDefault<ValueType>(
-      *metrics, metric_name_, labels_list);
+      *metrics, metric_name_, labels_list,
+      /*return_default_on_not_found=*/true);
   auto it = delta_map_.find(labels_list);
   ValueType initial_value;
   if (it == delta_map_.end()) {
     initial_value = internal::GetLatestValueOrDefault<ValueType>(
-        *initial_metrics_, metric_name_, labels_list);
+        *initial_metrics_, metric_name_, labels_list,
+        /*return_default_on_not_found=*/true);
     delta_map_[labels_list] = value;
   } else {
     initial_value = it->second;