Helper function to generate trace file name (#3456)

Summary:
Pull Request resolved: #3456

Add a helper function to generate trace file name, which will be later used in manifold uploader.

Also add a `.gz` suffix to the trace file name, this will enable trace compression according to https://github.com/pytorch/pytorch/blob/v2.8.0/torch/profiler/profiler.py#L257

Reviewed By: TroyGarden

Differential Revision: D84548870

fbshipit-source-id: a81b490e51761b21b050277517f62c0d1653fb49

Author: Shuangping Liu

torchrec/distributed/benchmark/base.py

@@ -571,6 +571,11 @@ def main(self) -> None:
             )
 
 
+def create_trace_file_name(profile_name: str, rank: int) -> str:
+    """Create a unique trace file name for the given rank and profile name."""
+    return f"trace-{profile_name}-rank{rank}.json.gz"
+
+
 def init_argparse_and_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser()
 
@@ -727,9 +732,10 @@ def _trace_handler(prof: torch.profiler.profile) -> None:
             if not all_rank_traces and rank > 0:
                 # only save trace for rank 0 when all_rank_traces is disabled
                 return
-            trace_file = f"{output_dir}/trace-{name}-rank{rank}.json"
+            trace_file = f"{output_dir}/{create_trace_file_name(name, rank)}"
             logger.info(f" PROFILE[{name}].chrome_trace:{trace_file}")
             prof.export_chrome_trace(trace_file)
+
             if export_stacks:
                 prof.export_stacks(
                     f"{output_dir}/stacks-cpu-{name}.stacks", "self_cpu_time_total"