[JAX] Add xml export for `test_multiprocessing_encoder` and `test_cgemm` (#2210)

* add xml export for test_multiprocessing_encoder and test_cgemm
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

examples/jax/collective_gemm/run_test_cgemm.sh

@@ -4,6 +4,10 @@
 
 NUM_GPUS=${NUM_GPUS:-$(nvidia-smi -L | wc -l)}
 
+: ${TE_PATH:=/opt/transformerengine}
+: ${XML_LOG_DIR:=/logs}
+mkdir -p "$XML_LOG_DIR"
+
 # Check if NVLINK is supported before running tests
 echo "*** Checking NVLINK support***"
 NVLINK_OUTPUT=$(nvidia-smi nvlink --status 2>&1)
@@ -69,7 +73,8 @@ for TEST_FILE in "${TEST_FILES[@]}"; do
       # For process 0: show live output AND save to log file using tee
       echo "=== Live output from process 0 ==="
       pytest -s -c "$TE_PATH/tests/jax/pytest.ini" \
-        -vs "$TE_PATH/examples/jax/collective_gemm/$TEST_FILE" \
+        -vs --junitxml=$XML_LOG_DIR/collective_gemm_${TEST_FILE}.xml \
+        "$TE_PATH/examples/jax/collective_gemm/$TEST_FILE" \
         --num-processes=$NUM_GPUS \
         --process-id=$i 2>&1 | tee "$LOG_FILE" &
       PID=$!
@@ -94,8 +99,11 @@ for TEST_FILE in "${TEST_FILES[@]}"; do
   elif grep -q "FAILED" "${TEST_FILE}_gpu_0.log"; then
     echo "... $TEST_FILE FAILED"
     HAS_FAILURE=1
-  else
+  elif grep -q "PASSED" "${TEST_FILE}_gpu_0.log"; then
     echo "... $TEST_FILE PASSED"
+  else
+    echo "... $TEST_FILE INVALID"
+    HAS_FAILURE=1
   fi
 
   # Remove the log files after processing them

examples/jax/encoder/run_test_multiprocessing_encoder.sh

@@ -15,11 +15,37 @@ TEST_CASES=(
 "test_te_current_scaling_fp8_shardy"
 )
 
+: ${TE_PATH:=/opt/transformerengine}
+: ${XML_LOG_DIR:=/logs}
+mkdir -p "$XML_LOG_DIR"
+
 echo
 echo "*** Executing tests in examples/jax/encoder/test_multiprocessing_encoder.py ***"
 
 HAS_FAILURE=0  # Global failure flag
 
+PIDS=()  # Array to store all process PIDs
+
+# Cleanup function to kill all processes
+cleanup() {
+  for pid in "${PIDS[@]}"; do
+    if kill -0 "$pid" 2>/dev/null; then
+      echo "Killing process $pid"
+      kill -TERM "$pid" 2>/dev/null || true
+    fi
+  done
+  # Wait a bit and force kill if needed
+  sleep 2
+  for pid in "${PIDS[@]}"; do
+    if kill -0 "$pid" 2>/dev/null; then
+      echo "Force killing process $pid"
+      kill -KILL "$pid" 2>/dev/null || true
+    fi
+  done
+}
+
+# Set up signal handlers to cleanup on exit
+trap cleanup EXIT INT TERM
 # Run each test case across all GPUs
 for TEST_CASE in "${TEST_CASES[@]}"; do
   echo
@@ -29,25 +55,40 @@ for TEST_CASE in "${TEST_CASES[@]}"; do
     # Define output file for logs
     LOG_FILE="${TEST_CASE}_gpu_${i}.log"
 
-    # Run pytest and redirect stdout and stderr to the log file
+    # For process 0: show live output AND save to log file using tee
+    if [ $i -eq 0 ]; then
+      echo "=== Live output from process 0 ==="
+      pytest -s -c "$TE_PATH/tests/jax/pytest.ini" \
+        -vs --junitxml=$XML_LOG_DIR/multiprocessing_encoder_${TEST_CASE}.xml \
+        "$TE_PATH/examples/jax/encoder/test_multiprocessing_encoder.py::TestEncoder::$TEST_CASE" \
+        --num-process=$NUM_GPUS \
+        --process-id=$i 2>&1 | tee "$LOG_FILE" &
+      PID=$!
+      PIDS+=($PID)
+    else
       pytest -s -c "$TE_PATH/tests/jax/pytest.ini" \
         -vs "$TE_PATH/examples/jax/encoder/test_multiprocessing_encoder.py::TestEncoder::$TEST_CASE" \
         --num-process=$NUM_GPUS \
         --process-id=$i > "$LOG_FILE" 2>&1 &
+      PID=$!
+      PIDS+=($PID)
+    fi
   done
 
   # Wait for the process to finish
   wait
-  tail -n +7 "${TEST_CASE}_gpu_0.log"
 
   # Check and print the log content accordingly
   if grep -q "SKIPPED" "${TEST_CASE}_gpu_0.log"; then
     echo "... $TEST_CASE SKIPPED"
+  elif grep -q "FAILED" "${TEST_CASE}_gpu_0.log"; then
+    echo "... $TEST_CASE FAILED"
+    HAS_FAILURE=1
   elif grep -q "PASSED" "${TEST_CASE}_gpu_0.log"; then
     echo "... $TEST_CASE PASSED"
   else
+    echo "... $TEST_CASE INVALID"
     HAS_FAILURE=1
-    echo "... $TEST_CASE FAILED"
   fi
 
   # Remove the log file after processing it
@@ -56,4 +97,8 @@ for TEST_CASE in "${TEST_CASES[@]}"; do
 done
 
 wait
+
+# Final cleanup (trap will also call cleanup on exit)
+cleanup
+
 exit $HAS_FAILURE