reset flight recorder trace (#283)

Summary:

- call FR api to reset the trace after every quorum - we reset so that after every quorum, we start a fresh FR trace since the pg's could have changed and we already dumped FR trace from previous errors
- change the env var that's used to determine the file after every quorum

Differential Revision: D84260745

Author: tushar00jain

torchft/manager.py

@@ -88,6 +88,8 @@
 # crash if call to quorum fails, all replicas will crash.
 QUORUM_RETRIES_ENV: str = "TORCHFT_QUORUM_RETRIES"
 
+TORCH_FR_DUMP_TEMP_FILE_ENV: str = "TORCH_FR_DUMP_TEMP_FILE"
+
 T = TypeVar("T")
 
 
@@ -223,6 +225,9 @@ def __init__(
         self._load_state_dict_fns: Dict[str, Callable[[object], None]] = {}
         self._user_state_dicts: Dict[str, Callable[[], object]] = {}
 
+        self._original_fr_dump_temp_file: Optional[str] = os.environ.get(
+            TORCH_FR_DUMP_TEMP_FILE_ENV
+        )
         self._replica_id = replica_id
 
         # Protects state dict
@@ -666,8 +671,16 @@ def _async_quorum(
             # We use the replica rank and world as we want all replicas in the PG.
             try:
                 with torch.profiler.record_function("torchft::manager::_pg::configure"):
+                    # Reset GPU state for Flight Recorder
                     if torch.accelerator.is_available():
                         torch.accelerator.synchronize()
+                    torch._C._distributed_c10d._reset_fr_recording_nccl()
+
+                    if self._original_fr_dump_temp_file is not None:
+                        os.environ[TORCH_FR_DUMP_TEMP_FILE_ENV] = (
+                            f"{self._original_fr_dump_temp_file}_quorum_{quorum_id}/replica_{self._replica_id}_"
+                        )
+
                     self._pg.configure(
                         store_prefixed_addr,
                         self._replica_id if self._replica_id is not None else "0",