Store kineto traces to a manifold bucket

Summary: Store kineto traces to a manifold bucket (`hpc_traces`)

Reviewed By: hiwotadese

Differential Revision: D52034905

fbshipit-source-id: 867a3d333781f7777bf37044cb83fc8eea7f3e6c

train/compute/python/tools/et_replay.py

@@ -1331,6 +1331,20 @@ def benchTime(self):
 
         prev_iter = self.numWarmupIters
         if self.profile_replay:
+            try:
+                from aiplatform.monitoring.unitrace.upload_manifold import (
+                    export_trace_func,
+                )
+
+                rank = self.comms_env_params["local_rank"]
+                on_trace_ready = export_trace_func(
+                    "/tmp",
+                    worker_name=f"rank-{rank}",
+                    bucket_name="hpc_traces",
+                    zoomer_request_callsite="hpc",
+                )
+            except ImportError:
+                on_trace_ready = trace_handler
             with torch.profiler.profile(
                 activities=[
                     torch.profiler.ProfilerActivity.CPU,
@@ -1340,7 +1354,7 @@ def benchTime(self):
                     wait=0, warmup=self.numWarmupIters, active=self.numIters
                 ),
                 record_shapes=True,
-                on_trace_ready=trace_handler,
+                on_trace_ready=on_trace_ready,
             ) as prof:
                 for iter in range(self.numWarmupIters + self.numIters):
                     if self.et_profile: