diff --git a/src/madengine/scripts/common/post_scripts/gpu_info_post.sh b/src/madengine/scripts/common/post_scripts/gpu_info_post.sh index 5582b986..516e6951 100644 --- a/src/madengine/scripts/common/post_scripts/gpu_info_post.sh +++ b/src/madengine/scripts/common/post_scripts/gpu_info_post.sh @@ -8,8 +8,15 @@ set -e set -x tool=$1 +model_name=$2 + +# Use model name if provided, otherwise fallback to tool name +if [ -n "$model_name" ]; then + OUTPUT=${tool}_${model_name}_output.csv +else + OUTPUT=${tool}_output.csv +fi -OUTPUT=${tool}_output.csv SAVESPACE=/myworkspace/ cd $SAVESPACE diff --git a/src/madengine/scripts/common/post_scripts/trace.sh b/src/madengine/scripts/common/post_scripts/trace.sh index 245dbff2..f86bda1b 100644 --- a/src/madengine/scripts/common/post_scripts/trace.sh +++ b/src/madengine/scripts/common/post_scripts/trace.sh @@ -8,8 +8,15 @@ set -e set -x tool=$1 +model_name=$2 + +# Use model name if provided, otherwise fallback to tool name +if [ -n "$model_name" ]; then + OUTPUT=${tool}_${model_name}_output +else + OUTPUT=${tool}_output +fi -OUTPUT=${tool}_output SAVESPACE=/myworkspace/ mkdir "$OUTPUT" @@ -18,7 +25,20 @@ case "$tool" in rpd) python3 ./rocmProfileData/tools/rpd2tracing.py trace.rpd trace.json - mv trace.rpd trace.json "$OUTPUT" + + # Rename trace files with model name if provided + if [ -n "$model_name" ]; then + if [ -f "trace.rpd" ]; then + mv trace.rpd "trace_${model_name}.rpd" + fi + if [ -f "trace.json" ]; then + mv trace.json "trace_${model_name}.json" + fi + mv "trace_${model_name}.rpd" "trace_${model_name}.json" "$OUTPUT" 2>/dev/null || true + else + mv trace.rpd trace.json "$OUTPUT" 2>/dev/null || true + fi + cp -vLR --preserve=all "$OUTPUT" "$SAVESPACE" ;; diff --git a/src/madengine/tools/run_models.py b/src/madengine/tools/run_models.py index 9e648590..78a58717 100644 --- a/src/madengine/tools/run_models.py +++ b/src/madengine/tools/run_models.py @@ -518,7 +518,7 @@ def get_mount_arg(self, mount_datapaths: typing.List) -> str: return mount_args - def run_pre_post_script(self, model_docker, model_dir, pre_post): + def run_pre_post_script(self, model_docker, model_dir, pre_post, model_name=None): for script in pre_post: script_path = script["path"].strip() model_docker.sh("cp -vLR --preserve=all " + script_path + " " + model_dir, timeout=600) @@ -527,6 +527,13 @@ def run_pre_post_script(self, model_docker, model_dir, pre_post): if "args" in script: script_args = script["args"] script_args.strip() + + # Add model name as an additional argument if provided + if model_name: + # Clean model name for safe file naming (replace / with _) + safe_model_name = model_name.replace("/", "_") + script_args += f" {safe_model_name}" + model_docker.sh("cd " + model_dir + " && bash " + script_name + " " + script_args , timeout=600) def run_model_impl( @@ -779,7 +786,7 @@ def run_model_impl( self.gather_system_env_details(pre_encapsulate_post_scripts, info['name']) # run pre_scripts if pre_encapsulate_post_scripts["pre_scripts"]: - self.run_pre_post_script(model_docker, model_dir, pre_encapsulate_post_scripts["pre_scripts"]) + self.run_pre_post_script(model_docker, model_dir, pre_encapsulate_post_scripts["pre_scripts"], info['name']) scripts_arg = info['scripts'] dir_path = None @@ -868,7 +875,7 @@ def run_model_impl( # run post_scripts if pre_encapsulate_post_scripts["post_scripts"]: - self.run_pre_post_script(model_docker, model_dir, pre_encapsulate_post_scripts["post_scripts"]) + self.run_pre_post_script(model_docker, model_dir, pre_encapsulate_post_scripts["post_scripts"], info['name']) # remove model directory if not self.args.keep_alive and not self.args.keep_model_dir: