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

Summary:
Pull Request resolved: #3504

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.

Reviewed By: spmex

Differential Revision: D86051540

fbshipit-source-id: 6371d3dc20285903cb59c1dc74017f70f1f9fceb

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"