added pyre support (#8)

Summary:
This runs the pyre type checking from OSS. It also formats some files + adds some types + small test fixes.

Pull Request resolved: #8

Reviewed By: kiukchung

Differential Revision: D28652691

Pulled By: d4l3k

fbshipit-source-id: 09e4ef1f82c1240a8da19b0018081a36b0f0bafd

Author: d4l3k

.github/workflows/pyre.yaml

@@ -0,0 +1,26 @@
+name: Pyre
+
+on:
+  push:
+    branches:
+      - master
+  pull_request:
+
+jobs:
+  pyre:
+    runs-on: ubuntu-18.04
+    steps:
+      - name: Setup Python
+        uses: actions/setup-python@v2
+        with:
+          python-version: 3.8
+          architecture: x64
+      - name: Checkout TorchX
+        uses: actions/checkout@v2
+      - name: Install Dependencies
+        run: |
+          set -eux
+          pip install -r dev-requirements.txt
+          pip install pyre-check
+      - name: Run Pyre
+        run: scripts/pyre.sh

.pyre_configuration

@@ -0,0 +1,11 @@
+{
+  "source_directories": [
+    "."
+  ],
+  "strict": true,
+  "exclude": [
+    ".*/build/.*",
+    ".*/docs/.*",
+    ".*/setup.py"
+  ]
+}

.watchmanconfig

@@ -0,0 +1,7 @@
+{
+  "root_files": [
+    "torchx",
+    ".pyre_configuration",
+    ".watchmanconfig"
+  ]
+}

dev-requirements.txt

@@ -4,3 +4,7 @@ kfp==1.4.0
 pyre-extensions>=0.0.21
 black>=21.5b1
 isort>=5.8.0
+pytorch-lightning>=0.5.3
+torch>=1.8.1
+torchvision>=0.9.1
+classy-vision>=0.5.0

scripts/pyre.sh

@@ -0,0 +1,11 @@
+#!/bin/sh
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+set -eux
+
+SITE_PACKAGES=$(python -c "import site; print(site.getsitepackages()[0])")
+pyre --search-path "${SITE_PACKAGES}" check

torchx/examples/lightning_classy_vision/lightning_classy_vision.torchx

@@ -25,21 +25,21 @@ arguments:
 import torchx.specs.api as torchx
 import torchx.schedulers.fb.resource as resource
 
-container = torchx.Container(image=args.image).require(resources=resource.get(args.resource))
+container = torchx.Container(image=args.image).require(
+    resources=resource.get(args.resource)
+)
 entrypoint = "main"
 
 trainer_role = (
-    torchx.Role(
-        name="trainer"
-    )
+    torchx.Role(name="trainer")
     .runs(
-      "main",
-      "--output_path",
-      args.output_path,
-      "--load_path",
-      args.load_path,
-      "--log_dir",
-      args.log_dir,
+        "main",
+        "--output_path",
+        args.output_path,
+        "--load_path",
+        args.load_path,
+        "--log_dir",
+        args.log_dir,
     )
     .on(container)
     .replicas(1)

torchx/pipelines/kfp/test/suites.py

@@ -11,7 +11,7 @@
 from itertools import chain
 
 
-def _circleci_parallelism(suite):
+def _circleci_parallelism(suite: unittest.TestSuite) -> unittest.TestSuite:
     """Allow for parallelism in CircleCI for speedier tests.."""
     if int(os.environ.get("CIRCLE_NODE_TOTAL", 0)) <= 1:
         # either not running on circleci, or we're not using parallelism.
@@ -23,14 +23,15 @@ def _circleci_parallelism(suite):
 
     # right now each test is corresponds to a /file/. Certain files are slower than
     # others, so we want to flatten it
+    # pyre-fixme[16]: `TestCase` has no attribute `_tests`.
     tests = [testfile._tests for testfile in suite._tests]
     tests = list(chain.from_iterable(tests))
     random.Random(42).shuffle(tests)
     tests = [t for i, t in enumerate(tests) if i % total == index]
     return unittest.TestSuite(tests)
 
 
-def unittests():
+def unittests() -> unittest.TestSuite:
     """
     Short tests.
 

torchx/runtime/container/test/main_test.py

@@ -5,17 +5,18 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import importlib
 import json
 import os.path
 import tempfile
 import unittest
-from typing import TypedDict, Optional
+from typing import Optional, TypedDict
 
 import yaml
 from torchx.runtime.component import Component
 from torchx.runtime.container.main import main
 from torchx.runtime.plugins import TORCHX_CONFIG_ENV
-from torchx.runtime.storage import temppath, upload_blob, download_blob
+from torchx.runtime.storage import download_blob, temppath, upload_blob
 
 
 class SubConfig(TypedDict):
@@ -126,6 +127,8 @@ def test_config_plugins(self) -> None:
         """
         from torchx.runtime.test import dummy_module
 
+        importlib.reload(dummy_module)
+
         module = "torchx.runtime.test.dummy_module"
         config = {
             "plugins": {

torchx/schedulers/registry.py

@@ -13,7 +13,7 @@
 
 
 def get_schedulers(
-    session_name: str, **scheduler_params
+    session_name: str, **scheduler_params: object
 ) -> Dict[SchedulerBackend, Scheduler]:
     return {
         "local": local_scheduler.create_scheduler(session_name, **scheduler_params),

torchx/test/suites.py

@@ -11,7 +11,7 @@
 from itertools import chain
 
 
-def _circleci_parallelism(suite):
+def _circleci_parallelism(suite: unittest.TestSuite) -> unittest.TestSuite:
     """Allow for parallelism in CircleCI for speedier tests.."""
     if int(os.environ.get("CIRCLE_NODE_TOTAL", 0)) <= 1:
         # either not running on circleci, or we're not using parallelism.
@@ -23,14 +23,15 @@ def _circleci_parallelism(suite):
 
     # right now each test is corresponds to a /file/. Certain files are slower than
     # others, so we want to flatten it
+    # pyre-fixme[16]: `TestCase` has no attribute `_tests`.
     tests = [testfile._tests for testfile in suite._tests]
     tests = list(chain.from_iterable(tests))
     random.Random(42).shuffle(tests)
     tests = [t for i, t in enumerate(tests) if i % total == index]
     return unittest.TestSuite(tests)
 
 
-def unittests():
+def unittests() -> unittest.TestSuite:
     """
     Short tests.