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

@@ -483,7 +483,7 @@ c10::intrusive_ptr<ProcessGroupXCCL::WorkXCCL> ProcessGroupXCCL::initWork(
                                 : std::nullopt);
 
   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_,
@@ -497,6 +497,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;
 }
@@ -803,9 +806,10 @@ 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);
+        FlightRecorderXCCL::get()->retire_id(id, reset_epoch, /*compute_duration*/ false);
       },
       /*use_future*/ false);
   work->blockingWait_ = blockingWait_;
@@ -891,7 +895,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_,
@@ -905,6 +909,8 @@ c10::intrusive_ptr<Work> ProcessGroupXCCL::pointToPoint(
         options_->timeout,
         pgStatus_,
         true);
+    work->trace_id_ = traceId.id;
+    work->trace_reset_epoch_ = traceId.reset_epoch;
 
     c10::OptionalDeviceGuard gpuGuard(device);
 
@@ -922,9 +928,10 @@ c10::intrusive_ptr<Work> ProcessGroupXCCL::pointToPoint(
         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);
+          FlightRecorderXCCL::get()->retire_id(id, reset_epoch, /*compute_duration*/ false);
         },
         /*use_future*/ false);
 

src/xccl/ProcessGroupXCCL.hpp

@@ -110,6 +110,7 @@ class TORCH_API ProcessGroupXCCL : public Backend {
     uint64_t seq_;
     bool isP2P_;
     std::optional<uint64_t> trace_id_;
+    std::optional<uint64_t> trace_reset_epoch_;
     size_t numelIn_ = -1;
     size_t numelOut_ = -1;