Lint fixes

Author: siju-samuel

torchft/checkpointing/transport_test.py

@@ -60,10 +60,7 @@ def run_multi_recovery_test(
     metadata: str = ""
 
     dist.init_process_group(
-        backend="gloo",
-        rank=0,
-        world_size=1,
-        store=dist.HashStore(),  # pyre-fixme[6]: Expected `Optional[Store]` but got `HashStore`
+        backend="gloo", rank=0, world_size=1, store=dist.HashStore()
     )
 
     def run(rank: int) -> CheckpointTransport[dict[str, object]]:

torchft/manager.py

@@ -58,7 +58,7 @@
 from torchft.checkpointing import CheckpointTransport, HTTPTransport
 from torchft.checkpointing._rwlock import RWLock
 from torchft.futures import future_timeout
-from torchft.utils import get_stream_context
+from torchft.utils import current_stream, get_stream_context
 from torchft.work import _DummyWork
 
 if TYPE_CHECKING:
@@ -790,7 +790,9 @@ def should_commit(self, timeout: Optional[timedelta] = None) -> bool:
             "torchft::manager::should_commit::current_stream::synchronize"
         ):
             if torch.accelerator.is_available():
-                torch.accelerator.current_stream().synchronize()
+
+                # pyre-fixme[16]: no attribute synchronize
+                current_stream().synchronize()
 
         if err := self._pg.errored():
             self.report_error(err)

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, record_event
+from torchft.utils import current_stream, get_stream_context, record_event
 from torchft.work import _DummyWork
 
 if TYPE_CHECKING:
@@ -793,7 +793,8 @@ def abort(self) -> None:
 
     def errored(self) -> Optional[Exception]:
         # force a synchronization to ensure all work is complete
-        torch.accelerator.current_stream().synchronize()
+        # pyre-fixme[16]: no attribute synchronize
+        current_stream().synchronize()
 
         return self._errored
 
@@ -877,6 +878,7 @@ def _create_pg(self, store: Store, rank: int, world_size: int) -> BaseProcessGro
 
         self._errored = None
 
+        # pyre-fixme[16]: no attribute ProcessGroupXCCL
         opts = BaseProcessGroupXCCL.Options()
         # opts.config.blocking = False
 

torchft/process_group_test.py

@@ -58,10 +58,7 @@
 def dummy_init_pg() -> None:
     if not dist.is_initialized():
         dist.init_process_group(
-            backend="gloo",
-            rank=0,
-            world_size=1,
-            store=dist.HashStore(),  # pyre-fixme[6]: Expected `Optional[Store]` but got `HashStore`
+            backend="gloo", rank=0, world_size=1, store=dist.HashStore()
         )
 
 

torchft/utils.py

@@ -54,3 +54,14 @@ def record_event() -> None:
         torch.xpu.current_stream().record_event(torch.xpu.Event())
     else:
         torch.cuda.current_stream().record_event(torch.cuda.Event(interprocess=True))
+
+
+def current_stream() -> None:
+    """
+    This function provides a unified way to get current stream across different
+    accelerator types (CUDA, XPU).
+    """
+    if torch.xpu.is_available():
+        torch.xpu.current_stream()
+    else:
+        torch.cuda.current_stream()