Introduce a unified Xprof workflow.

PiperOrigin-RevId: 914142746

Author: Pani122

xprof/convert/BUILD

@@ -70,6 +70,7 @@ cc_library(
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/base:no_destructor",
         "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/flags:flag",
         "@com_google_absl//absl/functional:any_invocable",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/strings",
@@ -1670,7 +1671,10 @@ cc_library(
         ":repository",
         ":tool_options",
         ":unified_hlo_stats_processor",
+        ":unified_profile_processor",
+        ":unified_profile_processor_factory",
         "@com_github_grpc_grpc//:grpc++",
+        "@com_google_absl//absl/flags:flag",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",

xprof/convert/unified_hlo_stats_processor.h

@@ -38,8 +38,7 @@ class UnifiedHloStatsProcessor : public BaseOpStatsProcessor {
       const tensorflow::profiler::ToolOptions& options) override;
 };
 
-REGISTER_UNIFIED_PROFILE_PROCESSOR("hlo_stats",
-                                   UnifiedHloStatsProcessor);
+REGISTER_UNIFIED_PROFILE_PROCESSOR("hlo_stats", UnifiedHloStatsProcessor);
 
 }  // namespace xprof
 

xprof/convert/unified_profile_processor_factory.cc

@@ -20,11 +20,15 @@
 #include "absl/base/no_destructor.h"
 #include "absl/functional/any_invocable.h"
 #include "absl/log/log.h"
+#include "absl/strings/match.h"
 #include "absl/strings/string_view.h"
+#include "absl/flags/flag.h"
 #include "absl/synchronization/mutex.h"
 #include "xprof/convert/tool_options.h"
 #include "xprof/convert/unified_profile_processor.h"
 
+ABSL_FLAG(bool, enable_unified_xprof, false, "Enable unified Xprof workflow");
+
 namespace xprof {
 
 UnifiedProfileProcessorFactory& UnifiedProfileProcessorFactory::GetInstance() {
@@ -46,8 +50,17 @@ std::unique_ptr<UnifiedProfileProcessor>
 UnifiedProfileProcessorFactory::Create(
     absl::string_view tool_name,
     const tensorflow::profiler::ToolOptions& options) const {
+  absl::string_view clean_tool_name = tool_name;
+  if (absl::EndsWith(clean_tool_name, ".json")) {
+    clean_tool_name.remove_suffix(5);
+  } else if (absl::EndsWith(clean_tool_name, ".pb")) {
+    clean_tool_name.remove_suffix(3);
+  } else if (absl::EndsWith(clean_tool_name, ".pbtxt")) {
+    clean_tool_name.remove_suffix(6);
+  }
+
   absl::MutexLock lock(mu_);
-  auto it = creators_.find(tool_name);
+  auto it = creators_.find(clean_tool_name);
   if (it == creators_.end()) {
     LOG(ERROR) << "No UnifiedProfileProcessor registered for tool: "
                << tool_name;

xprof/convert/unified_profile_processor_factory.h

@@ -60,7 +60,7 @@ class RegisterUnifiedProfileProcessor {
       UnifiedProfileProcessorFactory::Creator creator);
 };
 
-#define REGISTER_UNIFIED_PROFILE_PROCESSOR(tool_name, ClassName)               \
+#define REGISTER_UNIFIED_PROFILE_PROCESSOR(tool_name, ClassName)    \
   ABSL_ATTRIBUTE_UNUSED static const ::xprof::RegisterUnifiedProfileProcessor \
       register_##ClassName(                                                  \
           tool_name, [](const tensorflow::profiler::ToolOptions& options) {  \

xprof/convert/unified_profile_processor_factory_test.cc

@@ -66,7 +66,7 @@ class MacroRegisteredProcessor : public UnifiedProfileProcessor {
 };
 
 REGISTER_UNIFIED_PROFILE_PROCESSOR("macro_registered_tool",
-                                 MacroRegisteredProcessor);
+                                   MacroRegisteredProcessor);
 
 TEST(UnifiedProfileProcessorTest, BaseClassMethods) {
   DummyProcessor processor;

xprof/convert/xplane_to_tools_data_with_profile_processor.cc

@@ -5,9 +5,11 @@
 #include <variant>
 #include <vector>
 
+#include "absl/flags/flag.h"
 #include "absl/log/log.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
 #include "absl/time/clock.h"
 #include "absl/time/time.h"
@@ -17,17 +19,24 @@
 #include "xla/tsl/platform/errors.h"
 #include "xla/tsl/platform/threadpool.h"
 #include "tsl/platform/path.h"
+#include "absl/flags/declare.h"
 #include "xprof/convert/profile_processor.h"
 #include "xprof/convert/profile_processor_factory.h"
 #include "xprof/convert/repository.h"
 #include "xprof/convert/tool_options.h"
+#include "xprof/convert/unified_profile_processor.h"
+#include "xprof/convert/unified_profile_processor_factory.h"
 #include "plugin/xprof/protobuf/worker_service.pb.h"
 #include "plugin/xprof/worker/grpc_utils.h"
 #include "plugin/xprof/worker/stub_factory.h"
 
+ABSL_DECLARE_FLAG(bool, enable_unified_xprof);
+
 namespace tensorflow {
 namespace profiler {
+
 namespace {
+using xprof::UnifiedProfileProcessor;
 
 constexpr absl::string_view kXplaneFileName = ".xplane.pb";
 
@@ -133,6 +142,43 @@ absl::Status ProcessSession(xprof::ProfileProcessor* processor,
   return absl::OkStatus();
 }
 
+absl::Status RunUnifiedMapReduce(const SessionSnapshot& session_snapshot,
+                                 absl::string_view tool_name,
+                                 UnifiedProfileProcessor* processor,
+                                 const ToolOptions& options) {
+  const int num_hosts = session_snapshot.XSpaceSize();
+  std::vector<absl::StatusOr<std::string>> map_outputs(num_hosts);
+
+  {
+    tsl::thread::ThreadPool thread_pool(tsl::Env::Default(), __FUNCTION__,
+                                        num_hosts);
+    for (int i = 0; i < num_hosts; ++i) {
+      thread_pool.Schedule([&session_snapshot, &tool_name, &options,
+                            &map_outputs, i] {
+        std::string hostname = session_snapshot.GetHostname(i);
+        std::string xspace_path = GetXSpaceFilePath(session_snapshot, hostname);
+        map_outputs[i] = CallWorkerService(xspace_path, tool_name, options);
+      });
+    }
+  }
+
+  std::vector<std::string> map_output_files;
+  map_output_files.reserve(num_hosts);
+  for (int i = 0; i < num_hosts; ++i) {
+    TF_RETURN_IF_ERROR(map_outputs[i].status());
+    map_output_files.push_back(*std::move(map_outputs[i]));
+  }
+  LOG(INFO) << "Started reducing outputs for tool: " << tool_name
+            << " num_hosts: " << num_hosts;
+  absl::Time start_time = absl::Now();
+  absl::Status reduce_status =
+      processor->Reduce(session_snapshot, map_output_files);
+  absl::Duration reduce_time = absl::Now() - start_time;
+  LOG(INFO) << "Finished reducing outputs for tool: " << tool_name
+            << " num_hosts: " << num_hosts << " time taken: " << reduce_time;
+  return reduce_status;
+}
+
 }  // namespace
 
 absl::StatusOr<std::string> ConvertMultiXSpacesToToolDataWithProfileProcessor(
@@ -146,6 +192,28 @@ absl::StatusOr<std::string> ConvertMultiXSpacesToToolDataWithProfileProcessor(
 
   absl::Time start_time = absl::Now();
 
+  if (absl::GetFlag(FLAGS_enable_unified_xprof)) {
+    auto unified_processor =
+        xprof::UnifiedProfileProcessorFactory::GetInstance().Create(
+            tool_name, options);
+    if (!unified_processor) {
+      LOG(WARNING) << "Unified processor not found for tool: " << tool_name
+                   << ", falling back to legacy.";
+    } else {
+      LOG(INFO) << "Using unified workflow for tool: " << tool_name;
+      if (unified_processor->ShouldUseWorkerService(session_snapshot,
+                                                    options)) {
+        TF_RETURN_IF_ERROR(RunUnifiedMapReduce(session_snapshot, tool_name,
+                                               unified_processor.get(),
+                                               options));
+      } else {
+        TF_RETURN_IF_ERROR(
+            unified_processor->ProcessSession(session_snapshot, options));
+      }
+      return unified_processor->GetData();
+    }
+  }
+
   auto processor =
       xprof::ProfileProcessorFactory::GetInstance().Create(tool_name, options);
   if (!processor) {