fix a corner-case bug in memory snapshot uploading (#3504)

Summary:

Fixed two corner case issues in the TorchRec benchmark utilities:

1. **Memory snapshot handling**: Added rank filtering for memory snapshot operations to ensure they only run on rank 0 or when `all_rank_traces` is enabled. This prevents redundant memory snapshots from being taken on all ranks, reducing overhead and storage requirements while still capturing the necessary profiling data.

2. **Shell script robustness**: Added file existence checks before loop iterations in the trace upload script. Previously, if no trace files or memory snapshot files were found, the script would fail silently or produce errors. Now it checks with `ls` first and only proceeds with the loop if files exist, preventing issues when the trace directory is empty or files don't match the expected patterns.

Differential Revision: D86051540

Author: TroyGarden

torchrec/distributed/benchmark/base.py

@@ -749,7 +749,7 @@ def _trace_handler(prof: torch.profiler.profile) -> None:
                     f"{output_dir}/stacks-cuda-{name}.stacks", "self_cuda_time_total"
                 )
 
-        if memory_snapshot:
+        if memory_snapshot and (all_rank_traces or rank == 0):
             torch.cuda.empty_cache()
             torch.cuda.memory._record_memory_history(
                 max_entries=MAX_NUM_OF_MEM_EVENTS_PER_SNAPSHOT
@@ -775,7 +775,7 @@ def _trace_handler(prof: torch.profiler.profile) -> None:
         else:
             torch.cuda.synchronize(rank)
 
-        if memory_snapshot:
+        if memory_snapshot and (all_rank_traces or rank == 0):
             try:
                 torch.cuda.memory._dump_snapshot(
                     f"{output_dir}/memory-{name}-rank{rank}.pickle"