lintrunner: added black,isort,rustfmt (#22)

d4l3k · web-flow · commit b626d179ff05 · 2024-12-05T18:29:40.000-08:00
diff --git a/.github/workflows/lint.yaml b/.github/workflows/lint.yaml
@@ -23,12 +23,15 @@ jobs:
 
           sudo apt-get install -y protobuf-compiler
 
+          pip install lintrunner lintrunner-adapters
+          lintrunner init
+
           pip install .[dev] -v
-      - name: Run Python Lint
+      - name: Run lintrunner
         run: |
           set -eux
 
-          black --check .
+          lintrunner --force-color --all-files
       - name: Run Rust Lint
         run: |
           set -eux
diff --git a/.lintrunner.toml b/.lintrunner.toml
@@ -0,0 +1,44 @@
+[[linter]]
+code = 'BLACK-ISORT'
+include_patterns = [
+    '**/*.py',
+]
+exclude_patterns = []
+command = [
+    'python',
+    '-m',
+    'lintrunner_adapters',
+    'run',
+    'black_isort_linter',
+    '--fast',
+    '--',
+    '@{{PATHSFILE}}',
+]
+init_command = [
+    'python',
+    '-m',
+    'lintrunner_adapters',
+    'run',
+    'pip_init',
+    '--dry-run={{DRYRUN}}',
+    'black==24.10.0', # Use 24.x when ruff styles are updated
+    'isort==5.13.2',
+]
+is_formatter = true
+
+[[linter]]
+code = 'RUSTFMT'
+include_patterns = [
+    '**/*.rs',
+]
+command = [
+    'python',
+    '-m',
+    'lintrunner_adapters',
+    'run',
+    'rustfmt_linter',
+    '--binary=rustfmt',
+    '--config-path=.rustfmt.toml',
+    '--',
+    '@{{PATHSFILE}}',
+]
diff --git a/.rustfmt.toml b/.rustfmt.toml
@@ -0,0 +1 @@
+edition = "2021"
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
@@ -60,7 +60,9 @@ We actively welcome your pull requests.
 `torchft` enforces a fairly strict code format with tools such as cargo fmt and black.
 
 ```shell
-scripts/lint.sh
+pip install lintrunner lintrunner-adapters
+lintrunner init
+lintrunner -a
 ```
 
 ### Tests
diff --git a/pyproject.toml b/pyproject.toml
@@ -27,3 +27,7 @@ features = ["pyo3/extension-module"]
 
 [project.scripts]
 torchft_lighthouse = "torchft.torchft:lighthouse_main"
+
+[tool.isort]
+multi_line_output = 3
+combine_as_imports = true
diff --git a/scripts/lint.sh b/scripts/lint.sh
diff --git a/torchft/__init__.py b/torchft/__init__.py
@@ -4,11 +4,11 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from torchft.manager import Manager
-from torchft.process_group import ProcessGroupGloo, ProcessGroupBabyNCCL
+from torchft.data import DistributedSampler
 from torchft.ddp import DistributedDataParallel
+from torchft.manager import Manager
 from torchft.optim import OptimizerWrapper as Optimizer
-from torchft.data import DistributedSampler
+from torchft.process_group import ProcessGroupBabyNCCL, ProcessGroupGloo
 
 __all__ = (
     "DistributedDataParallel",
diff --git a/torchft/checkpointing_test.py b/torchft/checkpointing_test.py
@@ -4,13 +4,12 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from torchft.checkpointing import CheckpointServer
-
 import urllib.error
-
 from unittest import TestCase
 from unittest.mock import MagicMock
 
+from torchft.checkpointing import CheckpointServer
+
 
 class TestCheckpointing(TestCase):
     def test_checkpoint_server(self) -> None:
diff --git a/torchft/data.py b/torchft/data.py
@@ -17,7 +17,6 @@
 from typing import Optional
 
 import torch.distributed as dist
-
 from torch.utils import data
 
 
diff --git a/torchft/data_test.py b/torchft/data_test.py
@@ -6,9 +6,10 @@
 
 from unittest import TestCase
 
-from torchft.data import DistributedSampler
 from torch.utils.data import Dataset
 
+from torchft.data import DistributedSampler
+
 
 class DummyDataset(Dataset):
     def __init__(self, length: int) -> None:
diff --git a/torchft/ddp.py b/torchft/ddp.py
@@ -14,15 +14,15 @@
 
 import os
 import sys
-from typing import Optional, TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional
 from unittest.mock import patch
 
 import torch
 import torch.distributed as dist
 from torch import nn
 from torch.distributed.algorithms.join import Joinable
-
 from torch.nn import parallel
+
 from torchft.process_group import ProcessGroup, ProcessGroupDummy, ProcessGroupGloo
 
 if TYPE_CHECKING:
diff --git a/torchft/ddp_test.py b/torchft/ddp_test.py
@@ -8,16 +8,13 @@
 from unittest.mock import create_autospec
 
 import torch
-from torch import nn
 import torch.distributed as dist
+from torch import nn
 from torch.futures import Future
 
-from torchft.ddp import (
-    PureDistributedDataParallel,
-    DistributedDataParallel,
-)
-from torchft.process_group import ProcessGroupGloo, ProcessGroupBabyGloo
+from torchft.ddp import DistributedDataParallel, PureDistributedDataParallel
 from torchft.manager import Manager
+from torchft.process_group import ProcessGroupBabyGloo, ProcessGroupGloo
 
 
 class TestDDP(TestCase):
diff --git a/torchft/manager.py b/torchft/manager.py
@@ -32,11 +32,12 @@
 import uuid
 from concurrent.futures import ThreadPoolExecutor
 from datetime import timedelta
-from typing import Callable, Dict, List, Optional, TYPE_CHECKING
+from typing import TYPE_CHECKING, Callable, Dict, List, Optional
 
 import torch
 from torch.distributed import PrefixStore, ReduceOp, TCPStore, Work
 from torch.optim import Optimizer
+
 from torchft.checkpointing import CheckpointServer
 
 # pyre-fixme[21]: can't find rust module
diff --git a/torchft/manager_test.py b/torchft/manager_test.py
@@ -5,13 +5,13 @@
 # LICENSE file in the root directory of this source tree.
 
 from unittest import TestCase
-from unittest.mock import create_autospec, MagicMock, patch
+from unittest.mock import MagicMock, create_autospec, patch
 
 import torch
 from torch.distributed import TCPStore
-from torchft.manager import Manager, MANAGER_ADDR_KEY
-from torchft.process_group import _DummyWork, ProcessGroup
 
+from torchft.manager import MANAGER_ADDR_KEY, Manager
+from torchft.process_group import ProcessGroup, _DummyWork
 from torchft.torchft import ManagerClient
 
 
diff --git a/torchft/optim.py b/torchft/optim.py
@@ -12,7 +12,7 @@
 
 """
 
-from typing import Optional, TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional
 
 from torch.optim import Optimizer
 
diff --git a/torchft/optim_test.py b/torchft/optim_test.py
@@ -4,13 +4,14 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from torchft.optim import OptimizerWrapper
-from torchft.manager import Manager
+from unittest import TestCase
+from unittest.mock import MagicMock, create_autospec
+
 from torch.nn import Linear
 from torch.optim import AdamW
 
-from unittest import TestCase
-from unittest.mock import MagicMock, create_autospec
+from torchft.manager import Manager
+from torchft.optim import OptimizerWrapper
 
 
 class TestOptim(TestCase):
diff --git a/torchft/parameter_server.py b/torchft/parameter_server.py
@@ -24,7 +24,6 @@
 
 from torchft.process_group import ProcessGroup
 
-
 logger: logging.Logger = logging.getLogger(__name__)
 
 
diff --git a/torchft/process_group.py b/torchft/process_group.py
@@ -20,7 +20,7 @@
 import threading
 from abc import ABC
 from datetime import timedelta
-from typing import Callable, List, Optional, Tuple, Type, TYPE_CHECKING
+from typing import TYPE_CHECKING, Callable, List, Optional, Tuple, Type
 
 import torch
 import torch.distributed as dist
@@ -32,16 +32,15 @@
 from torch.distributed import (
     BroadcastOptions,
     DeviceMesh,
-    get_rank,
     PrefixStore,
     ProcessGroup as BaseProcessGroup,
     ProcessGroupGloo as BaseProcessGroupGloo,
     ProcessGroupNCCL as BaseProcessGroupNCCL,
     Store,
     TCPStore,
+    get_rank,
 )
-from torch.distributed.distributed_c10d import _world, Work
-
+from torch.distributed.distributed_c10d import Work, _world
 from torch.futures import Future
 
 if TYPE_CHECKING:
diff --git a/torchft/process_group_test.py b/torchft/process_group_test.py
@@ -6,22 +6,19 @@
 
 import os
 from concurrent.futures import ThreadPoolExecutor
-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
-from torch.distributed import _functional_collectives, ReduceOp, TCPStore
+from torch.distributed import ReduceOp, TCPStore, _functional_collectives
 from torch.distributed.device_mesh import init_device_mesh
-from torchft.manager import Manager
 
+from torchft.manager import Manager
 from torchft.process_group import (
-    _DummyWork,
-    _ErrorSwallowingWork,
     ErrorSwallowingProcessGroupWrapper,
-    extend_device_mesh,
     ManagedProcessGroup,
     ProcessGroup,
     ProcessGroupBabyGloo,
@@ -30,6 +27,9 @@
     ProcessGroupGloo,
     ProcessGroupNCCL,
     ProcessGroupWrapper,
+    _DummyWork,
+    _ErrorSwallowingWork,
+    extend_device_mesh,
 )
 
 
diff --git a/train_ddp.py b/train_ddp.py
@@ -4,24 +4,24 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-import sys
 import logging
 import os
+import sys
 
 import torch
-from torch import nn, optim
 import torch.nn.functional as F
 import torchvision
 import torchvision.transforms as transforms
+from torch import nn, optim
 from torchdata.stateful_dataloader import StatefulDataLoader
 
 from torchft import (
-    Manager,
-    ProcessGroupGloo,
-    ProcessGroupBabyNCCL,
     DistributedDataParallel,
-    Optimizer,
     DistributedSampler,
+    Manager,
+    Optimizer,
+    ProcessGroupBabyNCCL,
+    ProcessGroupGloo,
 )
 
 logging.basicConfig(level=logging.INFO)
