Fixed pyright issues (mosaicml#198)

- Updated pyright version to 1.1.203
- Update python for github tests to 3.9
- Made the pyproject.toml slightly stricter. Added warnings for untyped function signatures and implicit string concat, which should later be converted to errors.
- Removed unused variables and made those errors.
- Set all instance variables on `__init__`
- Fixed implicit None in signatures
Pin numpy to 1.21.5 since 1.22 introduces typing bugs

.github/actions/pyright/action.yaml

@@ -12,7 +12,7 @@ runs:
       working-directory: ${{ inputs.working_directory }}
       run: |
         set -ex
-        sudo npm install -g pyright@1.1.173
+        sudo npm install -g pyright@1.1.203
 
     - name: Run Pyright
       shell: bash

.github/workflows/pr-gate-cpu.yaml

@@ -11,7 +11,7 @@ jobs:
       - uses: actions/checkout@v2
       - uses: actions/setup-python@v2
         with:
-          python-version: 3.8
+          python-version: 3.9
       - name: Cache pip
         uses: actions/cache@v2
         with:

composer/algorithms/alibi/alibi.py

@@ -92,7 +92,8 @@ def apply_alibi(model: torch.nn.Module, heads_per_layer: int, max_sequence_lengt
                                                new_embedding_length=max_sequence_length)
     log.info(f" Position embedding expanded to sequence length {max_sequence_length}, zeroed, and frozen")
 
-    def convert_attention(module: torch.nn.Module, module_index: int = None):
+    def convert_attention(module: torch.nn.Module, module_index: Optional[int] = None):
+        del module_index  # unused
         module = register_alibi(module=module, n_heads=heads_per_layer, max_token_length=max_sequence_length)
         setattr(module, attr_to_replace, MethodType(alibi_attention, module))
         if mask_replacement_function:

composer/algorithms/augmix/augmix.py

@@ -1,6 +1,7 @@
 # Copyright 2021 MosaicML. All Rights Reserved.
 
 from dataclasses import asdict, dataclass
+from typing import Optional
 
 import numpy as np
 import torch
@@ -32,7 +33,7 @@ def initialize_object(self) -> "AugMix":
         return AugMix(**asdict(self))
 
 
-def augment_and_mix(img: ImageType = None,
+def augment_and_mix(img: Optional[ImageType] = None,
                     severity: int = 3,
                     depth: int = -1,
                     width: int = 3,

composer/algorithms/cutmix/cutmix.py

@@ -225,6 +225,9 @@ class CutMix(Algorithm):
 
     def __init__(self, alpha: float):
         self.hparams = CutMixHparams(alpha=alpha)
+        self._indices = torch.Tensor()
+        self._cutmix_lambda = 0.0
+        self._bbox = tuple()
 
     def match(self, event: Event, state: State) -> bool:
         """Runs on Event.INIT and Event.AFTER_DATALOADER
@@ -254,11 +257,11 @@ def cutmix_lambda(self, new_lambda: float) -> None:
         self._cutmix_lambda = new_lambda
 
     @property
-    def bbox(self) -> tuple:
+    def bbox(self) -> Tuple[int, int, int, int]:
         return self._bbox
 
     @bbox.setter
-    def bbox(self, new_bbox: tuple) -> None:
+    def bbox(self, new_bbox: Tuple[int, int, int, int]) -> None:
         self._bbox = new_bbox
 
     def apply(self, event: Event, state: State, logger: Logger) -> None:

composer/algorithms/cutout/cutout.py

@@ -44,7 +44,7 @@ def cutout(X: Tensor, n_holes: int, length: int) -> Tensor:
     w = X.size(3)
 
     mask = torch.ones_like(X)
-    for n in range(n_holes):
+    for _ in range(n_holes):
         y = np.random.randint(h)
         x = np.random.randint(w)
 

composer/algorithms/label_smoothing/label_smoothing.py

@@ -3,6 +3,7 @@
 from dataclasses import asdict, dataclass
 from typing import Optional
 
+import torch
 import yahp as hp
 
 from composer.algorithms.algorithm_hparams import AlgorithmHparams
@@ -37,6 +38,7 @@ class LabelSmoothing(Algorithm):
 
     def __init__(self, alpha: float):
         self.hparams = LabelSmoothingHparams(alpha=alpha)
+        self.original_labels = torch.Tensor()
 
     def match(self, event: Event, state: State) -> bool:
         return event in [Event.BEFORE_LOSS, Event.AFTER_LOSS]

composer/algorithms/mixup/mixup.py

@@ -129,6 +129,8 @@ class MixUp(Algorithm):
 
     def __init__(self, alpha: float):
         self.hparams = MixUpHparams(alpha=alpha)
+        self._interpolation_lambda = 0.0
+        self._indices = torch.Tensor()
 
     def match(self, event: Event, state: State) -> bool:
         """Runs on Event.INIT and Event.AFTER_DATALOADER
@@ -146,7 +148,7 @@ def interpolation_lambda(self) -> float:
         return self._interpolation_lambda
 
     @interpolation_lambda.setter
-    def interpolation_lambda(self, new_int_lamb) -> None:
+    def interpolation_lambda(self, new_int_lamb: float) -> None:
         self._interpolation_lambda = new_int_lamb
 
     @property

composer/algorithms/randaugment/randaugment.py

@@ -1,6 +1,7 @@
 # Copyright 2021 MosaicML. All Rights Reserved.
 
 from dataclasses import asdict, dataclass
+from typing import Optional
 
 import numpy as np
 import torch
@@ -28,7 +29,7 @@ def initialize_object(self) -> "RandAugment":
         return RandAugment(**asdict(self))
 
 
-def randaugment(img: ImageType = None,
+def randaugment(img: Optional[ImageType] = None,
                 severity: int = 9,
                 depth: int = 2,
                 augmentation_set: List = augmentation_sets["all"]) -> ImageType:

composer/algorithms/squeeze_excite/squeeze_excite.py

@@ -68,7 +68,7 @@ def forward(self, input: torch.Tensor) -> torch.Tensor:
 class SqueezeExciteConv2d(torch.nn.Module):
     """Helper class used to add a :class:`SqueezeExcite2d` module after a :class:`~torch.nn.Conv2d`."""
 
-    def __init__(self, *args, latent_channels=.125, conv: torch.nn.Conv2d = None, **kwargs):
+    def __init__(self, *args, latent_channels: float = 0.125, conv: Optional[torch.nn.Conv2d] = None, **kwargs):
         super().__init__()
         self.conv = torch.nn.Conv2d(*args, **kwargs) if conv is None else conv
         self.se = SqueezeExcite2d(num_features=self.conv.out_channels, latent_channels=latent_channels)

composer/callbacks/run_directory_uploader.py

@@ -99,7 +99,7 @@ def __init__(
         upload_staging_folder: Optional[str] = None,
         use_procs: bool = True,
         upload_every_n_batches: int = 100,
-        provider_init_kwargs: Dict[str, Any] = None,
+        provider_init_kwargs: Optional[Dict[str, Any]] = None,
     ) -> None:
         run_directory = get_run_directory()
         if run_directory is None:

composer/core/engine.py

@@ -52,7 +52,7 @@ def __init__(self,
                  state: State,
                  algorithms: Sequence[Algorithm],
                  logger: Optional[Logger] = None,
-                 callbacks: Sequence[Callback] = None):
+                 callbacks: Sequence[Callback] = tuple()):
         if logger is None:
             log.warning("No logger passed to the engine.  Defaulting to an empty logger")
             logger = Logger(state=state, backends=[])

composer/datasets/brats.py

@@ -257,7 +257,6 @@ def get_split(data, idx):
 def get_data_split(path: str):
     from sklearn.model_selection import KFold
 
-    val_cases_list = []
     kfold = KFold(n_splits=5, shuffle=True, random_state=0)
     imgs = load_data(path, "*_x.npy")
     lbls = load_data(path, "*_y.npy")

composer/datasets/lm_datasets.py

@@ -31,7 +31,10 @@ class LMDatasetHparams(DatasetHparams):
     Defines a generic dataset class for autoregressive language models.
     """
 
-    datadir: List[str] = hp.optional("Path to the Huggingface Datasets directory.", default_factory=list)
+    # TODO(moin): Switch datadir to be a string, rather than a list of strings, to be similar to the
+    # other datasets
+    datadir: List[str] = hp.optional(  # type: ignore
+        "Path to the Huggingface Datasets directory.", default_factory=list)
     split: Optional[str] = hp.optional("Whether to use 'train', 'validation' or 'test' split.", default=None)
     tokenizer_name: Optional[str] = hp.optional("The name of the tokenizer to preprocess text with.", default=None)
     num_tokens: int = hp.optional(doc='If desired, the number of tokens to truncate the dataset to.', default=0)

composer/models/base.py

@@ -116,13 +116,13 @@ def metrics(self, train: bool = False) -> Metrics:
         return self.train_acc if train else MetricCollection([self.val_acc, self.val_loss])
 
     def forward(self, batch: BatchPair) -> Tensor:
-        x, y = batch
+        x, _ = batch
         logits = self.module(x)
 
         return logits
 
     def validate(self, batch: BatchPair) -> Tuple[Any, Any]:
         assert self.training is False, "For validation, model must be in eval mode"
-        inputs, targets = batch
+        _, targets = batch
         logits = self.forward(batch)
         return logits, targets

composer/models/efficientnets.py

@@ -4,7 +4,7 @@
 # https://github.com/rwightman/gen-efficientnet-pytorch
 import math
 import re
-from typing import Any, Callable, Dict
+from typing import Any, Callable, Dict, Optional
 
 import torch
 import torch.nn as nn
@@ -14,7 +14,7 @@ def round_channels(
     channels: float,
     width_multiplier: float,
     divisor: int = 8,
-    min_value: int = None,
+    min_value: Optional[int] = None,
 ) -> int:
     """Round number of channels after scaling with width multiplier. This 
        function ensures that channel integers halfway inbetween divisors 
@@ -24,7 +24,8 @@ def round_channels(
         channels (float): Number to round.
         width_multiplier (float): Amount to scale `channels`.
         divisor (int): Number to make the output divisible by.
-        min_value (int): Minimum value the output can be.
+        min_value (int, optional): Minimum value the output can be. If not specified, defaults
+            to the ``divisor``.
     """
 
     if not width_multiplier:

composer/models/gpt2/scaling_laws_generator.py

@@ -331,8 +331,6 @@ def configure_mosaic_yaml(model, scaling_law_predictions):
     template_yaml['schedulers'][1]['cosine_decay']['T_max'] = f"{decay_steps}ba"
     template_yaml['model']['gpt2']['model_config'] = model
 
-    validation_freq = math.floor(final_serial_steps * args.validation_freq)
-
     return template_yaml
 
 

composer/models/unet/unet.py

@@ -61,7 +61,7 @@ def metrics(self, train: bool = False) -> Metrics:
         return self.dice
 
     def forward(self, batch: BatchPair) -> Tensor:
-        x, y = batch
+        x, _ = batch
         context = contextlib.nullcontext if self.training else torch.no_grad
 
         x = x.squeeze(1)  # type: ignore

composer/optim/decoupled_weight_decay.py

@@ -233,7 +233,6 @@ def step(self, closure=None):
             grads = []
             exp_avgs = []
             exp_avg_sqs = []
-            state_sums = []
             max_exp_avg_sqs = []
             state_steps = []
             amsgrad = group['amsgrad']

composer/trainer/checkpoint.py

@@ -27,6 +27,7 @@ class CheckpointLoader:
 
     def __init__(self, checkpoint_filepath: str):
         self.state_dict = torch.load(checkpoint_filepath, map_location='cpu')
+        self.checkpoint_rng_state = None
 
     def load_checkpoint(self, state: State):
         """Initialize state from the loaded checkpoint's data.

composer/trainer/trainer.py

@@ -525,7 +525,7 @@ def _train_loop(self) -> None:
                                       f"found {len(ensure_tuple(state.optimizers))} optimizers")
 
         assert isinstance(state.model, BaseMosaicModel)
-        self.original_model = state.model
+        self.original_model = state.model  # type: ignore  # TODO(ravi) -- update the state to add an original model helper
 
         # place the state, model in the proper devices
         if self.deepspeed_enabled:
@@ -831,7 +831,7 @@ def eval(self, is_batch: bool):
 
             metrics = self._get_metrics_as_collection(is_train=False)
 
-            for i, state.batch in enumerate(itertools.islice(state.eval_dataloader, self._eval_subset_num_batches)):
+            for state.batch in itertools.islice(state.eval_dataloader, self._eval_subset_num_batches):
                 state.batch = self.device.batch_to_device(state.batch)
                 if self._eval_device_transformation_fn is not None:
                     state.batch = self._eval_device_transformation_fn(state.batch)

composer/utils/iter_helpers.pyi

@@ -28,20 +28,5 @@ def ensure_tuple(x: Union[T, Tuple[T, ...], List[T], Dict[Any, T]]) -> Tuple[T,
     ...
 
 
-@overload
-def map_collection(tuple_of_elements: Tuple[T, ...], map_fn: Callable[[T], V], /) -> Tuple[V, ...]:
-    ...
-
-
-@overload
-def map_collection(list_of_elements: List[T], map_fn: Callable[[T], V], /) -> List[V]:
-    ...
-
-
-@overload
-def map_collection(dict_of_elements: Dict[KT, T], map_fn: Callable[[T], V], /) -> Dict[KT, V]:
-    ...
-
-
 def zip_collection(singleton: Any, *others: Any) -> Generator[Tuple[Any, ...], None, None]:
     ...

pyproject.toml

@@ -16,24 +16,36 @@ exclude = [
     "build/**"
 ]
 
-reportUnnecessaryIsInstance = "none"
+reportUnnecessaryIsInstance = "warning"
 reportMissingTypeStubs = "none"
 reportIncompatibleMethodOverride = "warning"
-reportIncompatibleVariableOverride = "warning"
+reportIncompatibleVariableOverride = "error"
 reportUnusedImport = "error"
 reportUnusedClass = "warning"
 reportUnusedFunction = "warning"
-reportUnusedVariable = "warning"
+reportUnusedVariable = "error"
 reportDuplicateImport = "error"
 reportWildcardImportFromLibrary = "error"
 reportUntypedFunctionDecorator = "warning"
 reportPrivateImportUsage = "warning"
 reportUndefinedVariable = "error"
-strictParameterNoneValue = "error"
+strictParameterNoneValue = true
 reportPropertyTypeMismatch = "error"
 reportUntypedNamedTuple = "error"
 reportUnnecessaryCast = "error"
 reportInvalidTypeVarUse = "error"
+reportOverlappingOverload = "error"
+reportUninitializedInstanceVariable = "error"
+reportInvalidStringEscapeSequence = "error"
+reportMissingParameterType = "warning"  # TODO: make this an error
+reportCallInDefaultInitializer = "none"  # TODO: make this an error
+reportUnnecessaryComparison = "warning"
+reportSelfClsParameterName = "error"
+reportImplicitStringConcatenation = "warning"  # TODO: make this an error
+reportInvalidStubStatement = "error"
+reportIncompleteStub = "error"
+reportUnsupportedDunderAll = "error"
+reportUnusedCoroutine = "error"
 
 # Pytest
 [tool.pytest.ini_options]

setup.py

@@ -2,15 +2,16 @@
 
 import os
 import sys
+import textwrap
 
 import setuptools
 from setuptools import setup
 
 
-def package_files(directory):
+def package_files(directory: str):
     # from https://stackoverflow.com/a/36693250
     paths = []
-    for (path, directories, filenames) in os.walk(directory):
+    for (path, _, filenames) in os.walk(directory):
         for filename in filenames:
             paths.append(os.path.join('..', path, filename))
     return paths
@@ -27,6 +28,7 @@ def package_files(directory):
     "torchvision>=0.9.0",
     "torch>=1.9",
     "yahp>=0.0.14",
+    "numpy==1.21.5",
 ]
 extra_deps = {}
 
@@ -106,9 +108,8 @@ def package_files(directory):
 # only visible if user installs with verbose -v flag
 # Printing to stdout as not to interfere with setup.py CLI flags (e.g. --version)
 print("*" * 20, file=sys.stderr)
-print(
-    "\nNOTE: For best performance, we recommend installing Pillow-SIMD "
-    "\nfor accelerated image processing operations. To install:"
-    "\n\n\t pip uninstall pillow && pip install pillow-simd\n",
-    file=sys.stderr)
+print(textwrap.dedent("""NOTE: For best performance, we recommend installing Pillow-SIMD
+    for accelerated image processing operations. To install:
+    \t pip uninstall pillow && pip install pillow-simd"""),
+      file=sys.stderr)
 print("*" * 20, file=sys.stderr)

tests/algorithms/test_cutmix.py

@@ -89,7 +89,7 @@ def test_cutmix(self, fake_data, alpha):
 
     def test_cutmix_algorithm(self, fake_data, alpha, dummy_state, dummy_logger):
         # Generate fake data
-        x_fake, y_fake, _, n_classes = fake_data
+        x_fake, y_fake, _, _ = fake_data
 
         algorithm = CutMixHparams(alpha=alpha).initialize_object()
         state = dummy_state

tests/algorithms/test_mixup.py

@@ -49,7 +49,7 @@ def test_mixup_batch(self, fake_data, alpha):
         interpolation_lambda = gen_interpolation_lambda(alpha)
 
         # Apply mixup
-        x_mix, y_mix, _ = mixup_batch(
+        x_mix, _, _ = mixup_batch(
             x=x_fake,
             y=y_fake,
             interpolation_lambda=interpolation_lambda,