Adds output stream stats to GraphRuntimeInfo

PiperOrigin-RevId: 730870783

Author: MediaPipe Team

Diff for: mediapipe/framework/BUILD

@@ -958,6 +958,8 @@ cc_library(
         "//mediapipe/framework/tool:tag_map",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/log:absl_check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
         "@com_google_absl//absl/synchronization",
     ],
 )
@@ -974,9 +976,10 @@ cc_library(
         ":packet_type",
         ":port",
         ":timestamp",
-        "//mediapipe/framework/port:source_location",
         "//mediapipe/framework/port:status",
+        "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/log:absl_check",
+        "@com_google_absl//absl/log:absl_log",
         "@com_google_absl//absl/synchronization",
     ],
 )
@@ -1473,7 +1476,6 @@ cc_test(
         "//mediapipe/framework/tool:source",
         "@com_google_absl//absl/log:absl_check",
         "@com_google_absl//absl/log:absl_log",
-        "@com_google_absl//absl/memory",
     ],
 )
 

Diff for: mediapipe/framework/calculator_node.cc

@@ -230,26 +230,35 @@ absl::Status CalculatorNode::Initialize(
 }
 
 CalculatorRuntimeInfo CalculatorNode::GetStreamMonitoringInfo() const {
-  CalculatorRuntimeInfo calulator_info;
-  calulator_info.set_calculator_name(DebugName());
+  CalculatorRuntimeInfo calculator_info;
+  calculator_info.set_calculator_name(DebugName());
   {
     absl::MutexLock lock(&runtime_info_mutex_);
-    calulator_info.set_last_process_start_unix_us(
+    calculator_info.set_last_process_start_unix_us(
         absl::ToUnixMicros(last_process_start_ts_));
-    calulator_info.set_last_process_finish_unix_us(
+    calculator_info.set_last_process_finish_unix_us(
         absl::ToUnixMicros(last_process_finish_ts_));
   }
-  const auto monitoring_info = input_stream_handler_->GetMonitoringInfo();
+  const auto input_stream_info = input_stream_handler_->GetMonitoringInfo();
   for (const auto& [stream_name, queue_size, num_packets_added,
-                    minimum_timestamp_or_bound] : monitoring_info) {
-    auto* stream_info = calulator_info.add_input_stream_infos();
+                    minimum_timestamp_or_bound] : input_stream_info) {
+    auto* stream_info = calculator_info.add_input_stream_infos();
     stream_info->set_stream_name(stream_name);
     stream_info->set_queue_size(queue_size);
     stream_info->set_number_of_packets_added(num_packets_added);
     stream_info->set_minimum_timestamp_or_bound(
         minimum_timestamp_or_bound.Value());
   }
-  return calulator_info;
+  const auto output_stream_info = output_stream_handler_->GetMonitoringInfo();
+  for (const auto& [stream_name, num_packets_added,
+                    minimum_timestamp_or_bound] : output_stream_info) {
+    auto* stream_info = calculator_info.add_output_stream_infos();
+    stream_info->set_stream_name(stream_name);
+    stream_info->set_number_of_packets_added(num_packets_added);
+    stream_info->set_minimum_timestamp_or_bound(
+        minimum_timestamp_or_bound.Value());
+  }
+  return calculator_info;
 }
 
 absl::Status CalculatorNode::InitializeOutputSidePackets(

Diff for: mediapipe/framework/calculator_node_test.cc

@@ -20,7 +20,6 @@
 
 #include "absl/log/absl_check.h"
 #include "absl/log/absl_log.h"
-#include "absl/memory/memory.h"
 #include "mediapipe/framework/calculator_framework.h"
 #include "mediapipe/framework/output_side_packet_impl.h"
 #include "mediapipe/framework/port/gmock.h"
@@ -573,5 +572,43 @@ TEST_F(CalculatorNodeTest, CleanupAfterRunTwiceWithTags) {
   TestCleanupAfterRunTwice();
 }
 
+TEST_F(CalculatorNodeTest, ShouldGenerateStreamMonitoringInfo) {
+  InitializeEnvironment(/*use_tags=*/false);
+  MP_ASSERT_OK(PrepareNodeForRun());
+
+  SimulateParentOpenNode();
+  MP_EXPECT_OK(node_->OpenNode());
+
+  OutputStreamShard stream_a_shard;
+  stream_a_shard.SetSpec(stream_a_manager_->Spec());
+  stream_a_shard.Add(new int(1), Timestamp(1));
+  stream_a_manager_->PropagateUpdatesToMirrors(Timestamp(2), &stream_a_shard);
+  EXPECT_EQ(1, schedule_count_);
+  EXPECT_TRUE(node_->TryToBeginScheduling());
+  MP_EXPECT_OK(node_->ProcessNode(cc_));
+  cc_ = nullptr;
+  node_->EndScheduling();
+  EXPECT_EQ(1, schedule_count_);
+
+  const auto stream_monitoring_info = node_->GetStreamMonitoringInfo();
+  EXPECT_EQ(stream_monitoring_info.calculator_name(), "CountCalculator");
+  EXPECT_EQ(stream_monitoring_info.input_stream_infos_size(), 1);
+  EXPECT_THAT(stream_monitoring_info.input_stream_infos(0),
+              mediapipe::EqualsProto(
+                  mediapipe::ParseTextProtoOrDie<InputStreamRuntimeInfo>(R"pb(
+                    stream_name: ":0:stream_a"
+                    number_of_packets_added: 1
+                    minimum_timestamp_or_bound: 2
+                  )pb")));
+  EXPECT_EQ(stream_monitoring_info.output_stream_infos_size(), 1);
+  EXPECT_THAT(stream_monitoring_info.output_stream_infos(0),
+              mediapipe::EqualsProto(
+                  mediapipe::ParseTextProtoOrDie<OutputStreamRuntimeInfo>(R"pb(
+                    stream_name: ":0:stream_b"
+                    number_of_packets_added: 1
+                    minimum_timestamp_or_bound: 2
+                  )pb")));
+}
+
 }  // namespace
 }  // namespace mediapipe

Diff for: mediapipe/framework/graph_runtime_info.proto

@@ -20,7 +20,7 @@ option java_package = "com.google.mediapipe.proto";
 option java_outer_classname = "GraphRuntimeInfoProto";
 
 // The runtime info for an input stream.
-message StreamRuntimeInfo {
+message InputStreamRuntimeInfo {
   // The name of the stream in the format "TAG:index:stream_name"
   string stream_name = 1;
 
@@ -34,6 +34,18 @@ message StreamRuntimeInfo {
   int64 minimum_timestamp_or_bound = 4;
 }
 
+// The runtime info for an output stream.
+message OutputStreamRuntimeInfo {
+  // The name of the stream in the format "TAG:index:stream_name"
+  string stream_name = 1;
+
+  // The total number of packets added to the queue.
+  int32 number_of_packets_added = 2;
+
+  // The minimum timestamp or timestanp bound of the stream.
+  int64 minimum_timestamp_or_bound = 3;
+}
+
 // The runtime info for a calculator.
 message CalculatorRuntimeInfo {
   // The name of the calculator.
@@ -49,7 +61,10 @@ message CalculatorRuntimeInfo {
   int64 timestamp_bound = 4;
 
   // The runtime info for each input stream of the calculator.
-  repeated StreamRuntimeInfo input_stream_infos = 5;
+  repeated InputStreamRuntimeInfo input_stream_infos = 5;
+
+  // The runtime info for each output stream of the calculator.
+  repeated OutputStreamRuntimeInfo output_stream_infos = 6;
 }
 
 // The runtime info for the whole graph.

Diff for: mediapipe/framework/output_stream_handler.cc

@@ -14,7 +14,12 @@
 
 #include "mediapipe/framework/output_stream_handler.h"
 
+#include <string>
+#include <vector>
+
 #include "absl/log/absl_check.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
 #include "absl/synchronization/mutex.h"
 #include "mediapipe/framework/collection_item_id.h"
 #include "mediapipe/framework/output_stream_shard.h"
@@ -121,6 +126,30 @@ std::string OutputStreamHandler::FirstStreamName() const {
   return (*output_stream_managers_.begin())->Name();
 }
 
+std::string OutputStreamHandler::DebugStreamName(CollectionItemId id) const {
+  const auto tag_map = output_stream_managers_.TagMap();
+  const std::string& stream_name = tag_map->Names()[id.value()];
+  const auto& [stream_tag, stream_idx] = tag_map->TagAndIndexFromId(id);
+  return absl::StrCat(stream_tag, ":", stream_idx, ":", stream_name);
+}
+
+std::vector<OutputStreamHandler::OutputStreamMonitoringInfo>
+OutputStreamHandler::GetMonitoringInfo() {
+  std::vector<OutputStreamMonitoringInfo> monitoring_info_vector;
+  for (CollectionItemId id = output_stream_managers_.BeginId();
+       id < output_stream_managers_.EndId(); ++id) {
+    const auto& stream = output_stream_managers_.Get(id);
+    if (!stream) {
+      continue;
+    }
+    monitoring_info_vector.emplace_back(OutputStreamMonitoringInfo(
+        {.stream_name = DebugStreamName(id),
+         .num_packets_added = stream->NumPacketsAdded(),
+         .next_timestamp_bound = stream->NextTimestampBound()}));
+  }
+  return monitoring_info_vector;
+}
+
 void OutputStreamHandler::TryPropagateTimestampBound(Timestamp input_bound) {
   // TODO Some non-range values, such as PostStream(), should also be
   // propagated.

Diff for: mediapipe/framework/output_stream_handler.h

@@ -22,6 +22,7 @@
 #include <unordered_map>
 #include <unordered_set>
 #include <utility>
+#include <vector>
 
 // TODO: Move protos in another CL after the C++ code migration.
 #include "absl/base/thread_annotations.h"
@@ -34,7 +35,6 @@
 #include "mediapipe/framework/output_stream_manager.h"
 #include "mediapipe/framework/packet_set.h"
 #include "mediapipe/framework/port/logging.h"
-#include "mediapipe/framework/port/status.h"
 #include "mediapipe/framework/timestamp.h"
 #include "mediapipe/framework/tool/tag_map.h"
 
@@ -49,6 +49,15 @@ class OutputStreamHandler {
       OutputStreamToSourcesMap;
   typedef internal::Collection<OutputStreamManager*> OutputStreamManagerSet;
 
+  // Struct to return monitoring info via GetMonitoringInfo;
+  struct OutputStreamMonitoringInfo {
+    std::string stream_name;
+    // The total number of packets added to the output stream.
+    int num_packets_added;
+    // The next timestamp bound of the output stream.
+    Timestamp next_timestamp_bound;
+  };
+
   // The constructor of the OutputStreamHandler takes four arguments.
   // The tag_map argument holds the information needed for tag/index retrieval
   // for the output streams; the calculator_context_manager for managing the
@@ -91,7 +100,7 @@ class OutputStreamHandler {
   }
 
   // Calls OutputStreamManager::PrepareForRun(error_callback) per stream, and
-  // resets data memebers.
+  // resets data members.
   void PrepareForRun(const std::function<void(absl::Status)>& error_callback)
       ABSL_LOCKS_EXCLUDED(timestamp_mutex_);
 
@@ -124,6 +133,14 @@ class OutputStreamHandler {
     return output_stream_managers_;
   }
 
+  // Return the stream name for an input stream in the format:
+  // stream_tag:stream_index:stream_name.
+  std::string DebugStreamName(CollectionItemId id) const;
+
+  // Returns a vector of tuples of stream name, number of packets added, and
+  // the next timestamp bound for each stream (for monitoring purposes).
+  std::vector<OutputStreamMonitoringInfo> GetMonitoringInfo();
+
  protected:
   // Checks if the given input bound should be propagated or not. If any output
   // streams with OffsetEnabled() need to have the timestamp bounds updated,

Diff for: mediapipe/framework/output_stream_manager.cc

@@ -14,7 +14,10 @@
 
 #include "mediapipe/framework/output_stream_manager.h"
 
+#include <list>
+
 #include "absl/log/absl_check.h"
+#include "absl/log/absl_log.h"
 #include "absl/synchronization/mutex.h"
 #include "mediapipe/framework/input_stream_handler.h"
 #include "mediapipe/framework/port/status_builder.h"
@@ -40,6 +43,7 @@ void OutputStreamManager::PrepareForRun(
     absl::MutexLock lock(&stream_mutex_);
     next_timestamp_bound_ = Timestamp::PreStream();
     closed_ = false;
+    num_packets_added_ = 0;
   }
 }
 
@@ -165,20 +169,21 @@ Timestamp OutputStreamManager::ComputeOutputTimestampBound(
 void OutputStreamManager::PropagateUpdatesToMirrors(
     Timestamp next_timestamp_bound, OutputStreamShard* output_stream_shard) {
   ABSL_CHECK(output_stream_shard);
+  std::list<Packet>* packets_to_propagate = output_stream_shard->OutputQueue();
   {
     if (next_timestamp_bound != Timestamp::Unset()) {
       absl::MutexLock lock(&stream_mutex_);
       next_timestamp_bound_ = next_timestamp_bound;
+      num_packets_added_ += packets_to_propagate->size();
       VLOG(3) << "Next timestamp bound for output " << output_stream_spec_.name
               << " is " << next_timestamp_bound_;
     }
   }
-  std::list<Packet>* packets_to_propagate = output_stream_shard->OutputQueue();
   VLOG(3) << "Output stream: " << Name()
           << " queue size: " << packets_to_propagate->size();
   VLOG(3) << "Output stream: " << Name()
           << " next timestamp: " << next_timestamp_bound;
-  bool add_packets = !packets_to_propagate->empty();
+  const bool add_packets = !packets_to_propagate->empty();
   bool set_bound =
       (next_timestamp_bound != Timestamp::Unset()) &&
       (!add_packets ||
@@ -218,4 +223,9 @@ void OutputStreamManager::ResetShard(OutputStreamShard* output_stream_shard) {
   output_stream_shard->Reset(next_timestamp_bound, closed);
 }
 
+int OutputStreamManager::NumPacketsAdded() const {
+  absl::MutexLock lock(&stream_mutex_);
+  return num_packets_added_;
+}
+
 }  // namespace mediapipe

Diff for: mediapipe/framework/output_stream_manager.h

@@ -15,10 +15,12 @@
 #ifndef MEDIAPIPE_FRAMEWORK_OUTPUT_STREAM_MANAGER_H_
 #define MEDIAPIPE_FRAMEWORK_OUTPUT_STREAM_MANAGER_H_
 
+#include <cstdint>
 #include <functional>
 #include <string>
 #include <vector>
 
+#include "absl/base/thread_annotations.h"
 #include "absl/synchronization/mutex.h"
 #include "mediapipe/framework/output_stream_shard.h"
 #include "mediapipe/framework/packet.h"
@@ -102,6 +104,10 @@ class OutputStreamManager {
   OutputStreamSpec* Spec() { return &output_stream_spec_; }
   const OutputStreamSpec* Spec() const { return &output_stream_spec_; }
 
+  // Returns the total number of packets added to the output stream. This is
+  // used for monitoring purposes.
+  int NumPacketsAdded() const;
+
  private:
   // The necessary information to locate an InputStreamImpl.
   struct Mirror {
@@ -120,6 +126,9 @@ class OutputStreamManager {
   mutable absl::Mutex stream_mutex_;
   Timestamp next_timestamp_bound_ ABSL_GUARDED_BY(stream_mutex_);
   bool closed_ ABSL_GUARDED_BY(stream_mutex_);
+  // Monotonically increasing total number of packets added to the output
+  // stream. This is used for monitoring purposes.
+  int64_t num_packets_added_ ABSL_GUARDED_BY(stream_mutex_) = 0;
 };
 
 }  // namespace mediapipe

Diff for: mediapipe/framework/output_stream_manager_test.cc

@@ -14,7 +14,9 @@
 
 #include "mediapipe/framework/output_stream_manager.h"
 
+#include <functional>
 #include <memory>
+#include <string>
 
 #include "absl/memory/memory.h"
 #include "mediapipe/framework/input_stream_handler.h"
@@ -32,7 +34,7 @@ namespace {
 
 class OutputStreamManagerTest : public ::testing::Test {
  protected:
-  OutputStreamManagerTest() {}
+  OutputStreamManagerTest() = default;
 
   void SetUp() override {
     packet_type_.Set<std::string>();
@@ -688,5 +690,21 @@ TEST_F(OutputStreamManagerTest, AddPacketAndMovePacket) {
   EXPECT_TRUE(errors_.empty());
 }
 
+TEST_F(OutputStreamManagerTest, ShouldReportNumPacketsAdded) {
+  Timestamp input_timestamp = Timestamp(0);
+  output_stream_shard_.AddPacket(
+      MakePacket<std::string>("packet 1").At(Timestamp(1)));
+  output_stream_shard_.AddPacket(
+      MakePacket<std::string>("packet 2").At(Timestamp(2)));
+
+  Timestamp output_bound = output_stream_manager_->ComputeOutputTimestampBound(
+      output_stream_shard_, input_timestamp);
+  output_stream_manager_->PropagateUpdatesToMirrors(output_bound,
+                                                    &output_stream_shard_);
+  ASSERT_TRUE(errors_.empty());
+  EXPECT_EQ(input_stream_manager_.NumPacketsAdded(), 2);
+  EXPECT_TRUE(output_stream_shard_.IsEmpty());
+}
+
 }  // namespace
 }  // namespace mediapipe