lintrunner: enable pyre

.github/workflows/docs.yaml

@@ -8,7 +8,7 @@ on:
 
 jobs:
   build:
-    runs-on: ubuntu-20.04
+    runs-on: ubuntu-latest
     steps:
       - name: Setup Python
         uses: actions/setup-python@v3

.github/workflows/lint.yaml

@@ -8,7 +8,7 @@ on:
 
 jobs:
   lint:
-    runs-on: ubuntu-20.04
+    runs-on: ubuntu-latest
     steps:
       - name: Setup Python
         uses: actions/setup-python@v3
@@ -31,7 +31,12 @@ jobs:
         run: |
           set -eux
 
-          lintrunner --force-color --all-files
+          lintrunner --skip PYRE --force-color --all-files
+      - name: Run pyre
+        run: |
+          set -eux
+
+          pyre check
       - name: Run Rust Lint
         run: |
           set -eux

.lintrunner.toml

@@ -42,3 +42,26 @@ command = [
     '--',
     '@{{PATHSFILE}}',
 ]
+
+[[linter]]
+code = 'PYRE'
+include_patterns = [
+    '**/*.py',
+    '**/*.pyi',
+]
+command = [
+    'python3',
+    'tools/linter/adapters/pyre_linter.py',
+    '--',
+    '@{{PATHSFILE}}'
+]
+init_command = [
+    'python',
+    '-m',
+    'lintrunner_adapters',
+    'run',
+    'pip_init',
+    '--dry-run={{DRYRUN}}',
+    'pyre-check==0.9.23',
+]
+is_formatter = false

.watchmanconfig

@@ -0,0 +1,8 @@
+{
+  "root_files": [
+    "torchft",
+    "*.py",
+    ".pyre_configuration",
+    ".watchmanconfig"
+  ]
+}

tools/linter/adapters/pyre_linter.py

@@ -0,0 +1,124 @@
+import argparse
+import concurrent.futures
+import json
+import logging
+import os
+import subprocess
+import sys
+from enum import Enum
+from pathlib import Path
+from typing import Any, List, NamedTuple, Optional, Set, TypedDict
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+
+class LintSeverity(str, Enum):
+    ERROR = "error"
+    WARNING = "warning"
+    ADVICE = "advice"
+    DISABLED = "disabled"
+
+
+class LintMessage(NamedTuple):
+    path: Optional[str]
+    line: Optional[int]
+    char: Optional[int]
+    code: str
+    severity: LintSeverity
+    name: str
+    original: Optional[str]
+    replacement: Optional[str]
+    description: Optional[str]
+
+
+class PyreResult(TypedDict):
+    line: int
+    column: int
+    stop_line: int
+    stop_column: int
+    path: str
+    code: int
+    name: str
+    description: str
+    concise_description: str
+
+
+def run_pyre() -> List[PyreResult]:
+    proc = subprocess.run(
+        ["pyre", "--output=json", "incremental"],
+        capture_output=True,
+    )
+    return json.loads(proc.stdout)
+
+
+def check_pyre(
+    filenames: Set[str],
+) -> List[LintMessage]:
+    try:
+        results = run_pyre()
+
+        return [
+            LintMessage(
+                path=result["path"],
+                line=result["line"],
+                char=result["column"],
+                code="pyre",
+                severity=LintSeverity.WARNING,
+                name=result["name"],
+                description=result["description"],
+                original=None,
+                replacement=None,
+            )
+            for result in results
+        ]
+    except Exception as err:
+        return [
+            LintMessage(
+                path=None,
+                line=None,
+                char=None,
+                code="pyre",
+                severity=LintSeverity.ADVICE,
+                name="command-failed",
+                original=None,
+                replacement=None,
+                description=(f"Failed due to {err.__class__.__name__}:\n{err}"),
+            )
+        ]
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Checks files with pyre",
+        fromfile_prefix_chars="@",
+    )
+    parser.add_argument(
+        "--verbose",
+        action="store_true",
+        help="verbose logging",
+    )
+    parser.add_argument(
+        "filenames",
+        nargs="+",
+        help="paths to lint",
+    )
+    args = parser.parse_args()
+
+    logging.basicConfig(
+        format="<%(processName)s:%(levelname)s> %(message)s",
+        level=(
+            logging.NOTSET
+            if args.verbose
+            else logging.DEBUG if len(args.filenames) < 1000 else logging.INFO
+        ),
+        stream=sys.stderr,
+    )
+
+    lint_messages = check_pyre(set(args.filenames))
+
+    for lint_message in lint_messages:
+        print(json.dumps(lint_message._asdict()), flush=True)
+
+
+if __name__ == "__main__":
+    main()

torchft/checkpointing.py

@@ -16,20 +16,19 @@
 import socket
 import threading
 import urllib.request
-from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
-from typing import Callable
+from http.server import BaseHTTPRequestHandler
+from typing import Callable, Generic, TypeVar
 
 import torch
 
-logger: logging.Logger = logging.getLogger(__name__)
+from torchft.http import _IPv6HTTPServer
 
+logger: logging.Logger = logging.getLogger(__name__)
 
-class _IPv6HTTPServer(ThreadingHTTPServer):
-    address_family = socket.AF_INET6
-    request_queue_size = 1024
+T = TypeVar("T")
 
 
-class CheckpointServer:
+class CheckpointServer(Generic[T]):
     """
     This is an HTTP server that can be used to transfer checkpoints
     between workers.
@@ -41,7 +40,7 @@ class CheckpointServer:
         state_dict: a callable that returns the state dict to be transferred
     """
 
-    def __init__(self, state_dict: Callable[[], object]) -> None:
+    def __init__(self, state_dict: Callable[[], T]) -> None:
         self._checkpoint_lock = threading.Lock()
         self._disallowed = False
         self._step = -1
@@ -88,7 +87,7 @@ def err(self, msg: str) -> None:
         self._thread.start()
 
     @classmethod
-    def load_from_address(cls, address: str) -> object:
+    def load_from_address(cls, address: str) -> T:
         """
         Loads a checkpoint from the given address.
 

torchft/data.py

@@ -47,7 +47,6 @@ def __init__(
         dataset: data.Dataset,
         replica_group: int,
         num_replica_groups: int,
-        *args,
         rank: Optional[int] = None,
         num_replicas: Optional[int] = None,
         **kwargs,
@@ -69,5 +68,8 @@ def __init__(
         self.global_world_size = num_replicas * num_replica_groups
 
         super().__init__(
-            dataset, *args, rank=self.global_rank, num_replicas=self.global_world_size
+            dataset,
+            rank=self.global_rank,
+            num_replicas=self.global_world_size,
+            **kwargs,
         )

torchft/ddp_test.py

@@ -44,7 +44,7 @@ def allreduce_grad(tensor: torch.Tensor) -> Future[torch.Tensor]:
 
             call_count += 1
 
-            fut = Future()
+            fut = Future()  # pyre-fixme[29]: not a function
             fut.set_result(tensor)
             return fut
 

torchft/http.py

@@ -0,0 +1,7 @@
+import socket
+from http.server import ThreadingHTTPServer
+
+
+class _IPv6HTTPServer(ThreadingHTTPServer):
+    address_family = socket.AF_INET6
+    request_queue_size = 1024