From 576c2f91388fa34d3a077642de73956ce783edd9 Mon Sep 17 00:00:00 2001
From: garima dhaked <garima.dhaked@amd.com>
Date: Fri, 7 Aug 2026 04:42:44 -0600
Subject: [PATCH] collect aie_dtrace/dtrace_dump files and keep zip/data
 indices aligned

Add aie_dtrace_*.ct and dtrace_dump_*.json to the files collected into
vai_profile.zip. When vai_profile.zip already exists, write the archive
as vai_profile_<n>.zip using the first free index that is also free for
vai_profile_data_<n>, so the zip and its extracted data directory always
share the same suffix.

Signed-off-by: garima dhaked <garima.dhaked@amd.com>
Signed-off-by: Garima Dhaked <garima.dhaked@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
---
 src/mldebug/scripts/vaiprofile-collect-data | 23 ++++++++++++++-------
 1 file changed, 16 insertions(+), 7 deletions(-)

diff --git a/src/mldebug/scripts/vaiprofile-collect-data b/src/mldebug/scripts/vaiprofile-collect-data
index 3d050c1..dcd5590 100755
--- a/src/mldebug/scripts/vaiprofile-collect-data
+++ b/src/mldebug/scripts/vaiprofile-collect-data
@@ -4,9 +4,10 @@
 # Packages VAIProfile inputs from a hardware run directory into vai_profile.zip (top-level files only),
 # then moves collected files into vai_profile_data (or vai_profile_data_<n> if the name exists).
 # Collects: aie_event_runtime_config_ctx_<id>_run_<id>.json, aie_trace_ctx_*_run_*_inf_*_strm_*.txt,
-# dtrace_dump_ctx*.py, record_timer_*.json (incl. record_timer_inference_*, record_timer_ts*, etc.),
+# dtrace_dump_ctx*.py, aie_dtrace_*.ct, dtrace_dump_*.json,
+# record_timer_*.json (incl. record_timer_inference_*, record_timer_ts*, etc.),
 # onnxruntime_profile_*.json.
-# vai_profile.zip remains at the run directory root.
+# vai_profile.zip remains at the run directory root (or vai_profile_<n>.zip if the name already exists).
 # Usage: vaiprofile-collect-data <hardware_run_directory>
 # Requires: zip (Info-ZIP) on PATH.
 
@@ -52,6 +53,12 @@ done
 for f in dtrace_dump_ctx*.py; do
   files+=("$f")
 done
+for f in aie_dtrace_*.ct; do
+  files+=("$f")
+done
+for f in dtrace_dump_*.json; do
+  files+=("$f")
+done
 for f in record_timer_*.json; do
   files+=("$f")
 done
@@ -68,17 +75,19 @@ if ((${#files[@]} == 0)); then
   exit 1
 fi
 
+# Pick a single index shared by vai_profile_<n>.zip and vai_profile_data_<n>,
+# so the zip and its extracted data directory always use the same suffix.
 out="${root}/vai_profile.zip"
-rm -f "$out"
-zip -j -q "$out" -- "${files[@]}"
-
 profile_dir="${root}/vai_profile_data"
-if [[ -e "$profile_dir" ]]; then
+if [[ -e "$out" || -e "$profile_dir" ]]; then
   n=1
-  while [[ -e "${root}/vai_profile_data_${n}" ]]; do
+  while [[ -e "${root}/vai_profile_${n}.zip" || -e "${root}/vai_profile_data_${n}" ]]; do
     ((n++)) || true
   done
+  out="${root}/vai_profile_${n}.zip"
   profile_dir="${root}/vai_profile_data_${n}"
 fi
+zip -j -q "$out" -- "${files[@]}"
+
 mkdir -p "$profile_dir"
 mv -- "${files[@]}" "$profile_dir/"