From 0a62a3cdc2ecf7bfb22d6da457b59103f168f27f Mon Sep 17 00:00:00 2001
From: hategan <mhategan@gmail.com>
Date: Thu, 27 Feb 2025 16:23:44 -0800
Subject: [PATCH] mpirun can mangle tagged output lines, so use heuristics to
 unmangle them.

---
 src/psij/launchers/scripts/mpi_launch.sh | 18 ++++++------------
 1 file changed, 6 insertions(+), 12 deletions(-)

diff --git a/src/psij/launchers/scripts/mpi_launch.sh b/src/psij/launchers/scripts/mpi_launch.sh
index 31c45233..7c01f89c 100644
--- a/src/psij/launchers/scripts/mpi_launch.sh
+++ b/src/psij/launchers/scripts/mpi_launch.sh
@@ -16,28 +16,22 @@ fi
 pre_launch
 
 filter_out() {
-    sed -nE 's/^\[[^]]+\]<stdout>:(.*)/\1/p'
-}
-
-filter_err() {
-    sed -nE 's/^\[[^]]+\]<stderr>:(.*)/\1/p'
+    # must start with [n,n]<stdout>:
+    # may contain one or more [n,n]<stdout>: lodged into the line
+    sed -nE 's/^\[[^]]+\]<std[oe][ur][tr]>:(.*)/\1/p' | sed -E 's/\[[^]]+\]<std[oe][ur][tr]>://g'
 }
 
 filter_out_5() {
-    sed -nE 's/^\[[^]]+\]<stdout>: (.*)/\1/p'
-}
-
-filter_err_5() {
-    sed -nE 's/^\[[^]]+\]<stderr>: (.*)/\1/p'
+    sed -nE 's/^\[[^]]+\]<std[oe][ur][tr]>: (.*)/\1/p'
 }
 
 set +e
 if [ "$IS_OPENMPI_5" == "1" ]; then
     mpirun --oversubscribe --output TAG -n $_PSI_J_PROCESS_COUNT "$@" \
-        1> >(filter_out_5 > $_PSI_J_STDOUT) 2> >(filter_err_5 > $_PSI_J_STDERR) <$_PSI_J_STDIN
+        1> >(filter_out_5 > $_PSI_J_STDOUT) 2> >(filter_out_5 > $_PSI_J_STDERR) <$_PSI_J_STDIN
 elif [ "$IS_OPENMPI" == "1" ]; then
     mpirun --oversubscribe --tag-output -q -n $_PSI_J_PROCESS_COUNT "$@" \
-        1> >(filter_out > "$_PSI_J_STDOUT") 2> >(filter_err > $_PSI_J_STDERR) <$_PSI_J_STDIN
+        1> >(filter_out > "$_PSI_J_STDOUT") 2> >(filter_out > $_PSI_J_STDERR) <$_PSI_J_STDIN
 else
     mpirun -n $_PSI_J_PROCESS_COUNT "$@" 1>$_PSI_J_STDOUT 2>$_PSI_J_STDERR <$_PSI_J_STDIN
 fi