From 6bf9b8a02e9038949e308ecfdcdbc977d8658209 Mon Sep 17 00:00:00 2001
From: Pranjal Swarup <Pranjal.Swarup@amd.com>
Date: Sun, 15 Dec 2024 18:03:12 -0500
Subject: [PATCH] Merged changes from multiprocess_output: - Execute the merge
 multiprocess script from the rank 0 process - Added scripts folder to path,
 when setup-env is added to the path - Install merge_multiprocess_output.sh to
 /share/rocprofiler-systems/bin dir in the install dir - Added script to merge
 multiprocess output to one file when there are multiprocess proto files
 written

---
 CMakeLists.txt                       |  9 +++++++
 cmake/Templates/setup-env.sh.in      |  1 +
 scripts/merge_multiprocess_output.sh | 39 ++++++++++++++++++++++++++++
 source/lib/core/perfetto.cpp         | 12 +++++++++
 4 files changed, 61 insertions(+)
 create mode 100755 scripts/merge_multiprocess_output.sh

diff --git a/CMakeLists.txt b/CMakeLists.txt
index a3e3b60d..19b71296 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -405,6 +405,10 @@ configure_file(
     ${PROJECT_BINARY_DIR}/${CMAKE_INSTALL_DATAROOTDIR}/modulefiles/${PROJECT_NAME}/${ROCPROFSYS_VERSION}
     @ONLY)
 
+configure_file(
+    ${PROJECT_SOURCE_DIR}/scripts/merge_multiprocess_output.sh
+    ${PROJECT_BINARY_DIR}/${CMAKE_INSTALL_DATAROOTDIR}/${PROJECT_NAME}/merge_multiprocess_output.sh COPYONLY)
+
 install(
     FILES ${PROJECT_BINARY_DIR}/${CMAKE_INSTALL_DATAROOTDIR}/${PROJECT_NAME}/setup-env.sh
           ${PROJECT_BINARY_DIR}/${CMAKE_INSTALL_DATAROOTDIR}/${PROJECT_NAME}/perfetto.cfg
@@ -422,6 +426,11 @@ install(
     DESTINATION ${CMAKE_INSTALL_DATAROOTDIR}/doc/${PROJECT_NAME}
     COMPONENT setup)
 
+install(
+    PROGRAMS ${PROJECT_BINARY_DIR}/${CMAKE_INSTALL_DATAROOTDIR}/${PROJECT_NAME}/merge_multiprocess_output.sh
+    DESTINATION ${CMAKE_INSTALL_DATAROOTDIR}/${PROJECT_NAME}/bin/
+    COMPONENT setup)
+
 # ------------------------------------------------------------------------------#
 #
 # install
diff --git a/cmake/Templates/setup-env.sh.in b/cmake/Templates/setup-env.sh.in
index b6c4a97d..81aac515 100644
--- a/cmake/Templates/setup-env.sh.in
+++ b/cmake/Templates/setup-env.sh.in
@@ -15,6 +15,7 @@ fi
 
 @PROJECT_NAME_UNDERSCORED@_ROOT=${BASEDIR}
 PATH=${BASEDIR}/bin:${PATH}
+PATH=${BASEDIR}/share/rocprofiler-systems/bin:${PATH}
 LD_LIBRARY_PATH=${BASEDIR}/@CMAKE_INSTALL_LIBDIR@:${LD_LIBRARY_PATH}
 PYTHONPATH=${BASEDIR}/@CMAKE_INSTALL_PYTHONDIR@:${PYTHONPATH}
 CMAKE_PREFIX_PATH=${BASEDIR}:${CMAKE_PREFIX_PATH}
diff --git a/scripts/merge_multiprocess_output.sh b/scripts/merge_multiprocess_output.sh
new file mode 100755
index 00000000..9e317f1c
--- /dev/null
+++ b/scripts/merge_multiprocess_output.sh
@@ -0,0 +1,39 @@
+#!/bin/bash
+
+# Check if the folder path is provided
+if [ -z "$1" ]; then
+  echo "Usage: $0 <folder_path>"
+  exit 1
+fi
+
+# Assign the folder path to a variable
+FOLDER_PATH=$1
+
+# Check if the folder exists
+if [ ! -d "$FOLDER_PATH" ]; then
+  echo "Error: Folder '$FOLDER_PATH' does not exist."
+  exit 1
+fi
+
+# Check if there are more than one .proto files
+PROTO_FILES=("$FOLDER_PATH"/*.proto)
+if [ ${#PROTO_FILES[@]} -le 1 ]; then
+  exit 0
+fi
+
+echo "Merging multiprocess files ..."
+# Check if all .proto files have been fully written or wait
+for file in "${PROTO_FILES[@]}"; do
+  while lsof "$file" > /dev/null 2>&1; do
+    echo "..."
+    sleep 1
+  done
+done
+
+# Output file name
+OUTPUT_FILE="merged.proto"
+
+# Merge all .proto files into one file
+cat "$FOLDER_PATH"/*.proto > "$FOLDER_PATH"/"$OUTPUT_FILE"
+
+echo "All multiprocess .proto files in '$FOLDER_PATH' have been merged into '$OUTPUT_FILE'."
\ No newline at end of file
diff --git a/source/lib/core/perfetto.cpp b/source/lib/core/perfetto.cpp
index 427ca6e6..bd296049 100644
--- a/source/lib/core/perfetto.cpp
+++ b/source/lib/core/perfetto.cpp
@@ -28,6 +28,7 @@
 #include "utility.hpp"
 
 #include <chrono>
+#include <filesystem>
 
 namespace rocprofsys
 {
@@ -263,6 +264,17 @@ post_process(tim::manager* _timemory_manager, bool& _perfetto_output_error)
                 _timemory_manager->add_file_output("protobuf", "perfetto", _filename);
         }
         ofs.close();
+
+        if(dmp::rank() == 0){
+            std::filesystem::path file_path(_filename);
+            auto folder_path = file_path.parent_path();
+            // Execute the merge script
+            std::string command = "merge_multiprocess_output.sh '" + std::string(folder_path.c_str()) + "'";
+            int result = system(command.c_str());
+            if (result != 0) {
+                ROCPROFSYS_VERBOSE(0, "Failed to execute merge_multiprocess_output.sh with folder path: %s\n", folder_path.c_str());
+            }
+        }
     }
     else if(dmp::rank() == 0)
     {