add isort

.github/workflows/format.yml

@@ -5,6 +5,17 @@ on:
     branches: [ "main" ]
     types: [opened, synchronize, reopened]
 
+env:
+  CHECK_FILES: >-
+    flagscale/auto_tuner/*.py
+    flagscale/auto_tuner/prune/*.py
+    flagscale/auto_tuner/record/*.py
+    flagscale/auto_tuner/search/*.py
+    flagscale/launcher/*.py
+    flagscale/logger.py
+    flagscale/patches_utils.py
+    flagscale/datasets/sft_dataset.py
+
 jobs:
   format:
     runs-on: ubuntu-latest
@@ -17,16 +28,15 @@ jobs:
         python-version: "3.10"
     - name: Install dependencies
       run: |
-        pip install black
+        pip install black isort
+        
     - name: Run Black
       run: >-
         black --check --diff --include
-        flagscale/auto_tuner/*.py
-        flagscale/auto_tuner/prune/*.py
-        flagscale/auto_tuner/record/*.py
-        flagscale/auto_tuner/search/*.py
-        flagscale/launcher/*.py
-        flagscale/logger.py
-        flagscale/patches_utils.py
-        flagscale/datasets/sft_dataset.py
-        ./
+        $CHECK_FILES
+        ./
+    
+    - name: Run Isort
+      run: |
+        isort . --check --diff 
+        $CHECK_FILES

.pre-commit-config.yaml

@@ -1,19 +1,27 @@
+check_files: &check_files |
+  (?x)^(
+      flagscale/auto_tuner/.*\.py|
+      flagscale/auto_tuner/prune/\..*\.py|
+      flagscale/auto_tuner/record/\..*\.py|
+      flagscale/auto_tuner/search/\..*\.py|
+      flagscale/launcher/\..*\.py|
+      flagscale/logger\.py|
+      flagscale/patches_utils\.py|
+      flagscale/datasets/sft_dataset\.py
+  )$
+
 repos:
   - repo: local
     hooks:
+      - id: isort
+        name: isort
+        entry: isort
+        language: system
+        types: [python]
+        files: *check_files
       - id: black
         name: black
         entry: black
         language: system
         types: [python]
-        files: |
-          (?x)^(
-              flagscale/auto_tuner/.*\.py|
-              flagscale/auto_tuner/prune/\..*\.py|
-              flagscale/auto_tuner/record/\..*\.py|
-              flagscale/auto_tuner/search/\..*\.py|
-              flagscale/launcher/\..*\.py|
-              flagscale/logger\.py|
-              flagscale/patches_utils\.py|
-              flagscale/datasets/sft_dataset\.py
-          )$
+        files: *check_files

flagscale/auto_tuner/generate.py

@@ -1,5 +1,5 @@
-import os
 import copy
+import os
 
 
 class Generator:

flagscale/auto_tuner/prune/history.py

@@ -1,4 +1,5 @@
 import logging
+
 from ..utils import beside
 
 _HISTORY_BASED_PRUNE_FUNC = []

flagscale/auto_tuner/record/recorder.py

@@ -1,7 +1,8 @@
+import logging
 import os
 import re
-import logging
 import subprocess
+
 import pandas as pd
 
 

flagscale/auto_tuner/search/searcher.py

@@ -1,8 +1,8 @@
-import time
 import copy
 import logging
-from ..utils import divisible
+import time
 
+from ..utils import divisible
 
 __BUILT_IN_STRATEGY_DIMS__ = [
     "data_parallel_size",

flagscale/auto_tuner/tuner.py

@@ -1,18 +1,18 @@
-import os
 import copy
-import time
 import datetime
 import logging
+import os
+import time
 
 from omegaconf import DictConfig, OmegaConf
 
-from flagscale.launcher.runner import SSHRunner
 from flagscale.launcher.job_status import JobStatus
+from flagscale.launcher.runner import SSHRunner
 
-from .search import Searcher
-from .prune import Pruner
 from .generate import Generator
+from .prune import Pruner
 from .record import Recorder
+from .search import Searcher
 
 
 class AutoTuner:

flagscale/auto_tuner/utils.py

@@ -1,7 +1,8 @@
-import re
 import os
+import re
 import socket
 import subprocess
+
 from flagscale.launcher.runner import parse_hostfile
 
 

flagscale/launcher/runner.py

@@ -1,17 +1,19 @@
-import os
-import re
-import sys
 import collections
 import copy
-import socket
+import json
+import os
+import re
 import shlex
+import socket
 import subprocess
-import json
-import uuid
+import sys
 import time
-from datetime import datetime
+import uuid
 from abc import ABC, abstractmethod
+from datetime import datetime
+
 from omegaconf import DictConfig, OmegaConf
+
 from ..logger import logger
 from .job_status import JobStatus
 

flagscale/logger.py

@@ -1,5 +1,5 @@
-import sys
 import logging
+import sys
 
 
 class Logger:

megatron/tests/unit_tests/__init__.py

@@ -1,2 +1,3 @@
 import torch._dynamo
-torch._dynamo.config.suppress_errors = True
+torch._dynamo.config.suppress_errors = True
+

megatron/tests/unit_tests/data/test_builder.py

@@ -98,7 +98,7 @@ def __len__(self) -> int:
         def __getitem__(self, idx: int) -> Dict[str, numpy.ndarray]:
             return {"text": self.dataset[self.sample_index[idx]]}
 
-    from tests.unit_tests.data import set_mock_args
+    from tests.unit_tests import set_mock_args
     set_mock_args()
 
     with tempfile.TemporaryDirectory() as temp_dir:

megatron/tests/unit_tests/data/test_gpt_dataset.py

@@ -53,7 +53,7 @@ def test_mock_gpt_dataset():
         tokenizer=tokenizer,
     )
 
-    from tests.unit_tests.data import set_mock_args
+    from tests.unit_tests import set_mock_args
     set_mock_args()
 
     datasets = BlendedMegatronDatasetBuilder(

megatron/tests/unit_tests/data/test_multimodal_dataset.py

@@ -43,7 +43,7 @@ def test_mock_multimodal_dataset():
         tokenizer=_NullTokenizer(vocab_size=_MOCK_VOCAB_SIZE),
     )
 
-    from tests.unit_tests.data import set_mock_args
+    from tests.unit_tests import set_mock_args
     set_mock_args()
 
     datasets = BlendedMegatronDatasetBuilder(