Adds and integrates GraphRuntimeInfoLogger into CalculatorGraph.

PiperOrigin-RevId: 703446598

mediapipe/framework/BUILD

@@ -369,6 +369,7 @@ cc_library(
         "//mediapipe/framework/port:source_location",
         "//mediapipe/framework/port:status",
         "//mediapipe/framework/tool:fill_packet_set",
+        "//mediapipe/framework/tool:graph_runtime_info_logger",
         "//mediapipe/framework/tool:packet_generator_wrapper_calculator",
         "//mediapipe/framework/tool:status_util",
         "//mediapipe/framework/tool:tag_map",

mediapipe/framework/calculator.proto

@@ -217,6 +217,19 @@ message ProfilerConfig {
   string calculator_filter = 18;
 }
 
+// Configuration for the runtime info logger. It collects runtime information
+// and statistics about calculators and their input streams at the configured
+// capture rate and writes them to LOG(INFO). It can be used to inspect a
+// stalled graph by understanding which calculators are waiting for input
+// packets to triger their Process() method.
+message GraphRuntimeInfoConfig {
+  // If true, the runtime info logger is enabled and runs in the background.
+  bool enable_graph_runtime_info = 1;
+  // The period in seconds at which the runtime info logger is updated. The
+  // default value is 10 secods.
+  uint32 capture_period_msec = 2;
+}
+
 // Describes the topology and function of a MediaPipe Graph.  The graph of
 // Nodes must be a Directed Acyclic Graph (DAG) except as annotated by
 // "back_edge" in InputStreamInfo.  Use a mediapipe::CalculatorGraph object to
@@ -392,6 +405,9 @@ message CalculatorGraphConfig {
   // calculators from running.  If false, max_queue_size for an input stream
   // is adjusted when throttling prevents all calculators from running.
   bool report_deadlock = 21;
+  // Enable the collection of runtime information and statistics about
+  // calculators and their input streams.
+  GraphRuntimeInfoConfig runtime_info = 22;
   // Config for this graph's InputStreamHandler.
   // If unspecified, the framework will automatically install the default
   // handler, which works as follows.

mediapipe/framework/calculator_graph.cc

@@ -460,6 +460,23 @@ absl::Status CalculatorGraph::Initialize(
 #endif
 
   initialized_ = true;
+
+  // Emscripten only supports single threaded applications.
+  const auto& runtime_info_logger_config =
+      validated_graph_->Config().runtime_info();
+#if !defined(__EMSCRIPTEN__)
+  if (runtime_info_logger_config.enable_graph_runtime_info()) {
+    MP_RETURN_IF_ERROR(graph_runtime_info_logger_.StartInBackground(
+        runtime_info_logger_config,
+        [this]() { return GetGraphRuntimeInfo(); }));
+  }
+#else
+  // TODO - remove once graph runtime infos are supported in
+  // Emscripten.
+  if (runtime_info_logger_config.enable_graph_runtime_info()) {
+    ABSL_LOG(WARNING) << "Graph runtime infos are not supported in Emscripten.";
+  }
+#endif  // defined(__EMSCRIPTEN__)
   return absl::OkStatus();
 }
 

mediapipe/framework/calculator_graph.h

@@ -58,6 +58,10 @@
 #include "mediapipe/framework/timestamp.h"
 #include "mediapipe/framework/validated_graph_config.h"
 
+#if !defined(__EMSCRIPTEN__)
+#include "mediapipe/framework/tool/graph_runtime_info_logger.h"
+#endif  // !defined(__EMSCRIPTEN__)
+
 namespace mediapipe {
 
 #if !MEDIAPIPE_DISABLE_GPU
@@ -770,6 +774,11 @@ class CalculatorGraph {
   std::shared_ptr<ProfilingContext> profiler_;
 
   internal::Scheduler scheduler_;
+
+#if !defined(__EMSCRIPTEN__)
+  // Collects runtime information about the graph in the background.
+  tool::GraphRuntimeInfoLogger graph_runtime_info_logger_;
+#endif  // !defined(__EMSCRIPTEN__)
 };
 
 }  // namespace mediapipe

mediapipe/framework/tool/BUILD

@@ -978,6 +978,41 @@ cc_test(
     ],
 )
 
+cc_library(
+    name = "graph_runtime_info_logger",
+    srcs = ["graph_runtime_info_logger.cc"],
+    hdrs = ["graph_runtime_info_logger.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":graph_runtime_info_utils",
+        "//mediapipe/framework:calculator_cc_proto",
+        "//mediapipe/framework:graph_runtime_info_cc_proto",
+        "//mediapipe/framework/port:ret_check",
+        "//mediapipe/framework/port:threadpool",
+        "@com_google_absl//absl/functional:any_invocable",
+        "@com_google_absl//absl/log:absl_check",
+        "@com_google_absl//absl/log:absl_log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/synchronization",
+        "@com_google_absl//absl/time",
+    ],
+)
+
+cc_test(
+    name = "graph_runtime_info_logger_test",
+    size = "small",
+    srcs = ["graph_runtime_info_logger_test.cc"],
+    deps = [
+        ":graph_runtime_info_logger",
+        "//mediapipe/framework:calculator_cc_proto",
+        "//mediapipe/framework/port:gtest_main",
+        "//mediapipe/framework/port:status_matchers",
+        "@com_google_absl//absl/synchronization",
+        "@com_google_absl//absl/time",
+    ],
+)
+
 cc_library(
     name = "graph_runtime_info_utils",
     srcs = ["graph_runtime_info_utils.cc"],

mediapipe/framework/tool/graph_runtime_info_logger.cc

@@ -0,0 +1,62 @@
+#include "mediapipe/framework/tool/graph_runtime_info_logger.h"
+
+#include <string>
+#include <utility>
+
+#include "absl/functional/any_invocable.h"
+#include "absl/log/absl_check.h"
+#include "absl/log/absl_log.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/time/time.h"
+#include "mediapipe/framework/calculator.pb.h"
+#include "mediapipe/framework/port/ret_check.h"
+#include "mediapipe/framework/tool/graph_runtime_info_utils.h"
+
+namespace mediapipe::tool {
+
+constexpr absl::Duration kDefaultCaptureInterval = absl::Seconds(10);
+
+GraphRuntimeInfoLogger::GraphRuntimeInfoLogger()
+    : thread_pool_("GraphRuntimeInfoLogger", 1) {}
+
+GraphRuntimeInfoLogger::~GraphRuntimeInfoLogger() { Stop(); };
+
+absl::Status GraphRuntimeInfoLogger::StartInBackground(
+    const mediapipe::GraphRuntimeInfoConfig& config,
+    absl::AnyInvocable<absl::StatusOr<GraphRuntimeInfo>()>
+        get_runtime_info_fn) {
+  get_runtime_info_fn_ = std::move(get_runtime_info_fn);
+  RET_CHECK(!is_running_.HasBeenNotified());
+  ABSL_CHECK_EQ(thread_pool_.num_threads(), 1);
+  thread_pool_.StartWorkers();
+  absl::Duration interval =
+      config.capture_period_msec() > 0
+          ? absl::Milliseconds(config.capture_period_msec())
+          : kDefaultCaptureInterval;
+  thread_pool_.Schedule([this, interval]() mutable {
+    is_running_.Notify();
+    while (!shutdown_signal_.HasBeenNotified()) {
+      const auto runtime_info = get_runtime_info_fn_();
+      if (!runtime_info.ok()) {
+        ABSL_LOG(DFATAL) << "Failed to get graph runtime info: "
+                         << runtime_info.status();
+        return;
+      }
+      const auto runtime_info_str = GetGraphRuntimeInfoString(*runtime_info);
+      if (!runtime_info_str.ok()) {
+        ABSL_LOG(DFATAL) << "Failed to render graph runtime info: "
+                         << runtime_info_str.status();
+        return;
+      }
+      ABSL_LOG(INFO) << *runtime_info_str;
+      shutdown_signal_.WaitForNotificationWithTimeout(interval);
+    }
+  });
+  is_running_.WaitForNotification();
+  return absl::OkStatus();
+}
+
+void GraphRuntimeInfoLogger::Stop() { shutdown_signal_.Notify(); }
+
+}  // namespace mediapipe::tool

mediapipe/framework/tool/graph_runtime_info_logger.h

@@ -0,0 +1,51 @@
+// Copyright 2024 The MediaPipe Authors.
+//
+// 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.
+
+#ifndef MEDIAPIPE_FRAMEWORK_TOOL_GRAPH_RUNTIME_INFO_LOGGER_H_
+#define MEDIAPIPE_FRAMEWORK_TOOL_GRAPH_RUNTIME_INFO_LOGGER_H_
+
+#include "absl/functional/any_invocable.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/synchronization/notification.h"
+#include "mediapipe/framework/calculator.pb.h"
+#include "mediapipe/framework/graph_runtime_info.pb.h"
+#include "mediapipe/framework/port/threadpool.h"
+
+namespace mediapipe::tool {
+
+// Periodically collects the graph runtime info and output it to LOG(INFO).
+class GraphRuntimeInfoLogger {
+ public:
+  GraphRuntimeInfoLogger();
+  ~GraphRuntimeInfoLogger();
+
+  // Starts the collector in the background. Can be called only once.
+  absl::Status StartInBackground(
+      const mediapipe::GraphRuntimeInfoConfig& config,
+      absl::AnyInvocable<absl::StatusOr<GraphRuntimeInfo>()>
+          get_runtime_info_fn);
+
+ private:
+  void Stop();
+
+  absl::Notification shutdown_signal_;
+  absl::Notification is_running_;
+  ThreadPool thread_pool_;
+  absl::AnyInvocable<absl::StatusOr<GraphRuntimeInfo>()> get_runtime_info_fn_;
+};
+
+}  // namespace mediapipe::tool
+
+#endif  // MEDIAPIPE_FRAMEWORK_TOOL_GRAPH_RUNTIME_INFO_LOGGER_H_

mediapipe/framework/tool/graph_runtime_info_logger_test.cc

@@ -0,0 +1,42 @@
+// Copyright 2024 The MediaPipe Authors.
+//
+// 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 "mediapipe/framework/tool/graph_runtime_info_logger.h"
+
+#include "absl/synchronization/notification.h"
+#include "absl/time/time.h"
+#include "mediapipe/framework/calculator.pb.h"
+#include "mediapipe/framework/port/gmock.h"
+#include "mediapipe/framework/port/gtest.h"
+#include "mediapipe/framework/port/status_matchers.h"
+
+namespace mediapipe::tool {
+namespace {
+
+TEST(GraphRuntimeInfoLoggerTest, ShouldCaptureRuntimeInfo) {
+  mediapipe::GraphRuntimeInfoConfig config;
+  config.set_enable_graph_runtime_info(true);
+
+  absl::Notification callback_called;
+  GraphRuntimeInfoLogger logger;
+  MP_ASSERT_OK(logger.StartInBackground(config, [&]() {
+    callback_called.Notify();
+    return GraphRuntimeInfo();
+  }));
+  EXPECT_TRUE(
+      callback_called.WaitForNotificationWithTimeout(absl::Seconds(10)));
+}
+
+}  // namespace
+}  // namespace mediapipe::tool