update flight recorder usage

Summary:
we're making changes to flight recorder's record and reset_id api's. update the usage accordingly.

Author: tushar00jain

src/xccl/ProcessGroupXCCL.cpp

@@ -508,7 +508,7 @@ c10::intrusive_ptr<ProcessGroupXCCL::WorkXCCL> ProcessGroupXCCL::initWork(
       xpuEventCacheEnabled_.load());
 
   if (record) {
-    r->trace_id_ = FlightRecorderXCCL::get()->record(
+    auto traceId = FlightRecorderXCCL::get()->record(
         local_id_,
         std::make_tuple(pg_uid_, pg_desc_), // PG name tuple
         seqCollective_,
@@ -522,6 +522,9 @@ c10::intrusive_ptr<ProcessGroupXCCL::WorkXCCL> ProcessGroupXCCL::initWork(
         options_->timeout,
         pgStatus_,
         isP2P);
+
+    r->trace_id_ = traceId.id;
+    r->trace_reset_epoch_ = traceId.reset_epoch;
   }
   return r;
 }
@@ -847,9 +850,11 @@ c10::intrusive_ptr<Work> ProcessGroupXCCL::collective(
       c10::ListType::create(c10::TensorType::get()), devices);
   work->future_->markCompleted(at::IValue(*work->outputs_));
   auto id = work->trace_id_;
+  auto reset_epoch = work->trace_reset_epoch_;
   work->future_->addCallback(
-      [id](at::ivalue::Future&) {
-        FlightRecorderXCCL::get()->retire_id(id, /*compute_duration*/ false);
+      [id, reset_epoch](at::ivalue::Future&) {
+        FlightRecorderXCCL::get()->retire_id(
+            id, reset_epoch, /*compute_duration*/ false);
       },
       /*use_future*/ false);
   work->blockingWait_ = blockingWait_;
@@ -948,7 +953,7 @@ c10::intrusive_ptr<Work> ProcessGroupXCCL::pointToPoint(
     work->outputs_ = std::make_shared<std::vector<at::Tensor>>();
     work->outputs_->push_back(tensor);
 
-    work->trace_id_ = FlightRecorderXCCL::get()->record(
+    auto traceId = FlightRecorderXCCL::get()->record(
         local_id_,
         std::make_tuple(pg_uid_, pg_desc_), // PG name tuple
         seqCollective_,
@@ -962,6 +967,9 @@ c10::intrusive_ptr<Work> ProcessGroupXCCL::pointToPoint(
         options_->timeout,
         pgStatus_,
         true);
+
+    work->trace_id_ = traceId.id;
+    work->trace_reset_epoch_ = traceId.reset_epoch;
   }
 
   if (enableNanCheck_ && opType == OpType::SEND) {
@@ -999,9 +1007,11 @@ c10::intrusive_ptr<Work> ProcessGroupXCCL::pointToPoint(
     }
 
     auto id = work->trace_id_;
+    auto reset_epoch = work->trace_reset_epoch_;
     work->future_->addCallback(
-        [id](at::ivalue::Future&) {
-          FlightRecorderXCCL::get()->retire_id(id, /*compute_duration*/ false);
+        [id, reset_epoch](at::ivalue::Future&) {
+          FlightRecorderXCCL::get()->retire_id(
+              id, reset_epoch, /*compute_duration*/ false);
         },
         /*use_future*/ false);
     setEnqueuedPgStatus(work);

src/xccl/ProcessGroupXCCL.hpp

@@ -121,7 +121,8 @@ class TORCH_API ProcessGroupXCCL : public Backend {
     std::chrono::time_point<std::chrono::steady_clock> workStartTime_;
     uint64_t seq_;
     bool isP2P_;
-    std::optional<uint64_t> trace_id_;
+    std::optional<size_t> trace_id_;
+    std::optional<size_t> trace_reset_epoch_;
     size_t numelIn_ = -1;
     size_t numelOut_ = -1;