Update

[ghstack-poisoned]

Author: fegin

torchft/process_group.py

@@ -20,7 +20,7 @@
 import threading
 from abc import ABC
 from datetime import timedelta
-from typing import Dict, List, Optional, Tuple, Type, TYPE_CHECKING, Union
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Type, Union
 
 import torch
 import torch.distributed as dist
@@ -31,16 +31,16 @@
 from torch.distributed import (
     BroadcastOptions,
     DeviceMesh,
-    get_rank,
-    init_device_mesh,
     PrefixStore,
     ProcessGroup as BaseProcessGroup,
     ProcessGroupGloo as BaseProcessGroupGloo,
     ProcessGroupNCCL as BaseProcessGroupNCCL,
     Store,
     TCPStore,
+    get_rank,
+    init_device_mesh,
 )
-from torch.distributed.distributed_c10d import _world, Work
+from torch.distributed.distributed_c10d import Work, _world
 from torch.futures import Future
 
 if TYPE_CHECKING:

torchft/process_group_test.py

@@ -7,37 +7,32 @@
 import os
 from concurrent.futures import ThreadPoolExecutor
 from typing import Any, Dict, Tuple
-from unittest import skipUnless, TestCase
+from unittest import TestCase, skipUnless
 from unittest.mock import Mock
 
 import torch
 import torch.distributed as dist
 from torch import nn
 from torch._C._distributed_c10d import (
-    _resolve_process_group,
     AllgatherOptions,
     AllreduceOptions,
     BroadcastOptions,
     ReduceOp,
+    _resolve_process_group,
 )
 from torch.distributed import (
-    _functional_collectives,
-    get_world_size,
     ReduceOp,
     TCPStore,
     Work,
+    _functional_collectives,
+    get_world_size,
 )
 from torch.distributed.device_mesh import init_device_mesh
 from torch.testing._internal.common_distributed import MultiProcessTestCase
 
 from torchft.manager import Manager
 from torchft.process_group import (
-    _DummyWork,
-    _ErrorSwallowingWork,
-    _ManagedWork,
     ErrorSwallowingProcessGroupWrapper,
-    extend_device_mesh,
-    ft_init_device_mesh,
     ManagedDeviceMesh,
     ManagedProcessGroup,
     ProcessGroup,
@@ -47,6 +42,11 @@
     ProcessGroupGloo,
     ProcessGroupNCCL,
     ProcessGroupWrapper,
+    _DummyWork,
+    _ErrorSwallowingWork,
+    _ManagedWork,
+    extend_device_mesh,
+    ft_init_device_mesh,
 )