Fix UT failures due to record_event

Author: siju-samuel

torchft/process_group.py

@@ -69,7 +69,7 @@
 from torchft.device_mesh import *  # noqa: F401
 from torchft.futures import context_timeout, stream_timeout
 from torchft.multiprocessing import _MonitoredPipe
-from torchft.utils import get_stream_context
+from torchft.utils import get_stream_context, record_event
 from torchft.work import _DummyWork
 
 if TYPE_CHECKING:
@@ -793,7 +793,7 @@ def abort(self) -> None:
 
     def errored(self) -> Optional[Exception]:
         # force a synchronization to ensure all work is complete
-        torch.cuda.current_stream().synchronize()
+        torch.accelerator.current_stream().synchronize()
 
         return self._errored
 
@@ -1539,13 +1539,7 @@ def _worker(
 
                         # Register event on the stream that we can pass to the main
                         # process.
-                        event = (
-                            torch.accelerator.current_stream().record_event(
-                                torch.Event(interprocess=True)
-                            )
-                            if metadata.stream is not None
-                            else None
-                        )
+                        event = record_event() if metadata.stream is not None else None
 
                     req_pipe.send((op_id, event))
                 elif cmd == "del":
@@ -1562,9 +1556,7 @@ def callback(fut: Future[object], metadata: _OpMetadata) -> None:
                             with metadata.set_stream():
                                 fut.wait()
                                 event = (
-                                    torch.accelerator.current_stream().record_event(
-                                        torch.Event(interprocess=True)
-                                    )
+                                    record_event()
                                     if metadata.stream is not None
                                     else None
                                 )
@@ -1660,13 +1652,7 @@ def _run_func(self, func: str, *args: object, **kwargs: object) -> Work:
         stream_id = (
             torch.accelerator.current_stream().stream_id if is_accelerator else None
         )
-        event = (
-            torch.accelerator.current_stream().record_event(
-                torch.Event(interprocess=True)
-            )
-            if is_accelerator
-            else None
-        )
+        event = record_event() if is_accelerator else None
 
         op_id = self._next_op_id
         self._next_op_id += 1

torchft/utils.py

@@ -41,3 +41,16 @@ def get_stream_context(
             return nullcontext()
     else:
         return nullcontext()
+
+
+def record_event() -> None:
+    """
+    Record an event in the current stream.
+
+    This function provides a unified way to record events across different
+    accelerator types (CUDA, XPU).
+    """
+    if torch.xpu.is_available():
+        torch.xpu.current_stream().record_event(torch.xpu.Event(interprocess=True))
+    else:
+        torch.cuda.current_stream().record_event(torch.cuda.Event(interprocess=True))