Move torchtnt to arc pyre (#1017)

Summary:
Pull Request resolved: #1017

Turn `python.set_typing(True)` on and remove the old pyre configuration. This opts into the per-target type-checking (e.g. `arc pyre check-changed-targets`).

See https://www.internalfb.com/wiki/Python/Type-annotations-in-python/How_To%3A_Migrate_to_Pyre_Fast_By_Default_Per_Target_Type_Checking/

Reviewed By: diego-urgell

Differential Revision: D78682161

fbshipit-source-id: 5d740b4d82481d06c089d7d9e5c7ae05683ac436

Author: alanhdu

tests/framework/callbacks/test_csv_writer.py

@@ -26,7 +26,6 @@ def get_step_output_rows(
         self,
         state: State,
         unit: PredictUnit[TPredictData],
-        # pyre-fixme[2]: Parameter annotation cannot be `Any`.
         step_output: Any,
     ) -> Union[List[str], List[List[str]]]:
         return [["1"], ["2"]]
@@ -37,7 +36,6 @@ def get_step_output_rows(
         self,
         state: State,
         unit: PredictUnit[TPredictData],
-        # pyre-fixme[2]: Parameter annotation cannot be `Any`.
         step_output: Any,
     ) -> Union[List[str], List[List[str]]]:
         return ["1"]

tests/framework/test_auto_unit.py

@@ -15,7 +15,6 @@
 
 from pyre_extensions import none_throws, ParameterSpecification as ParamSpec
 from torch import nn
-
 from torch.distributed import GradBucket
 from torchtnt.framework._test_utils import (
     DummyAutoUnit,
@@ -456,7 +455,6 @@ def custom_noop_hook(
         ) -> torch.futures.Future[torch.Tensor]:
             nonlocal custom_noop_hook_called
 
-            # pyre-fixme[29]: `Type[torch.futures.Future]` is not a function.
             fut: torch.futures.Future[torch.Tensor] = torch.futures.Future()
             fut.set_result(bucket.buffer())
             custom_noop_hook_called = True

tests/utils/test_memory.py

@@ -220,16 +220,17 @@ def __init__(self) -> None:
             len(tensor_map), 2 * len(inputs.metric_list[0].window_buffer.buffers) + 6
         )
         for metric in inputs.metric_list:
-            self.assertTrue(metric.x in tensor_map)
+            metric = cast(RandomModule, metric)
+            self.assertIn(metric.x, tensor_map)
             self.assertEqual(
                 tensor_map[metric.x], metric.x.size().numel() * metric.x.element_size()
             )
-            self.assertTrue(metric.y[0] in tensor_map)
+            self.assertIn(metric.y[0], tensor_map)
             self.assertEqual(
                 tensor_map[metric.y[0]],
                 metric.y[0].size().numel() * metric.y[0].element_size(),
             )
-            self.assertTrue(metric.y[1] in tensor_map)
+            self.assertIn(metric.y[1], tensor_map)
             self.assertEqual(
                 tensor_map[metric.y[1]],
                 metric.y[1].size().numel() * metric.y[1].element_size(),

torchtnt/framework/auto_unit.py

@@ -131,7 +131,6 @@ class TrainStepResults:
 
     loss: torch.Tensor
     total_grad_norm: Optional[torch.Tensor]
-    # pyre-fixme[4]: Attribute `outputs` of class `TrainStepResults` must have a type other than `Any`.
     outputs: Any
 
 
@@ -371,7 +370,6 @@ def __init__(
             global_mesh=global_mesh,
         )
 
-    # pyre-fixme[3]: Return annotation cannot be `Any`.
     def predict_step(self, state: State, data: TPredictData) -> Any:
         # if detect_anomaly is true, run forward pass under detect_anomaly context
         detect_anomaly = self.detect_anomaly
@@ -394,7 +392,6 @@ def on_predict_step_end(
         state: State,
         data: TPredictData,
         step: int,
-        # pyre-fixme[2]: Parameter annotation cannot be `Any`.
         outputs: Any,
     ) -> None:
         """
@@ -645,7 +642,6 @@ def configure_optimizers_and_lr_scheduler(
         ...
 
     @abstractmethod
-    # pyre-fixme[3]: Return annotation cannot contain `Any`.
     def compute_loss(self, state: State, data: TData) -> Tuple[torch.Tensor, Any]:
         """
         The user should implement this method with their loss computation. This will be called every ``train_step``/``eval_step``.
@@ -662,7 +658,6 @@ def compute_loss(self, state: State, data: TData) -> Tuple[torch.Tensor, Any]:
         """
         ...
 
-    # pyre-fixme[3]: Return annotation cannot contain `Any`.
     def train_step(self, state: State, data: TData) -> Tuple[torch.Tensor, Any]:
         should_update_weights = (
             self.train_progress.num_steps_completed_in_epoch + 1
@@ -779,7 +774,6 @@ def on_train_epoch_end(self, state: State) -> None:
 
         self._is_last_batch = False
 
-    # pyre-fixme[3]: Return annotation cannot contain `Any`.
     def eval_step(self, state: State, data: TData) -> Tuple[torch.Tensor, Any]:
         with self.maybe_autocast_precision:
             # users must override this
@@ -800,7 +794,6 @@ def on_eval_step_end(
         data: TData,
         step: int,
         loss: torch.Tensor,
-        # pyre-fixme[2]: Parameter annotation cannot be `Any`.
         outputs: Any,
     ) -> None:
         """
@@ -816,7 +809,6 @@ def on_eval_step_end(
         """
         pass
 
-    # pyre-fixme[3]: Return annotation cannot contain `Any`.
     def predict_step(self, state: State, data: TData) -> Any:
         with self.maybe_autocast_precision:
             with get_timing_context(state, f"{self.__class__.__name__}.forward"):
@@ -835,7 +827,6 @@ def on_predict_step_end(
         state: State,
         data: TData,
         step: int,
-        # pyre-fixme[2]: Parameter annotation cannot be `Any`.
         outputs: Any,
     ) -> None:
         """

torchtnt/framework/callbacks/base_csv_writer.py

@@ -12,7 +12,6 @@
 from typing import Any, List, TextIO, Union
 
 from pyre_extensions import none_throws
-
 from torchtnt.framework.callback import Callback
 from torchtnt.framework.state import EntryPoint, State
 from torchtnt.framework.unit import TEvalUnit, TPredictUnit, TTrainUnit
@@ -61,7 +60,6 @@ def get_step_output_rows(
         self,
         state: State,
         unit: TPredictUnit,
-        # pyre-fixme: Missing parameter annotation [2]
         step_output: Any,
     ) -> Union[List[str], List[List[str]]]: ...
 

torchtnt/framework/callbacks/module_summary.py

@@ -45,7 +45,6 @@ def __init__(
         process_fn: Callable[
             [List[ModuleSummaryObj]], None
         ] = _log_module_summary_tables,
-        # pyre-fixme
         module_inputs: Optional[
             MutableMapping[str, Tuple[Tuple[Any, ...], Dict[str, Any]]]
         ] = None,

torchtnt/framework/train.py

@@ -247,6 +247,8 @@ def _train_epoch_impl(
             ):
                 _evaluate_impl(
                     state,
+                    # pyre-fixme[6]: For 2nd argument expected `EvalUnit[Any]` but
+                    #  got `TrainUnit[Any]`.
                     train_unit,
                     callback_handler,
                 )
@@ -293,6 +295,8 @@ def _train_epoch_impl(
     ):
         _evaluate_impl(
             state,
+            # pyre-fixme[6]: For 2nd argument expected `EvalUnit[Any]` but got
+            #  `TrainUnit[Any]`.
             train_unit,
             callback_handler,
         )

torchtnt/framework/unit.py

@@ -362,7 +362,6 @@ def on_train_epoch_start(self, state: State) -> None:
         pass
 
     @abstractmethod
-    # pyre-fixme[3]: Return annotation cannot be `Any`.
     def train_step(self, state: State, data: TTrainData) -> Any:
         """Core required method for user to implement. This method will be called at each iteration of the
         train dataloader, and can return any data the user wishes.
@@ -476,7 +475,6 @@ def on_eval_epoch_start(self, state: State) -> None:
         pass
 
     @abstractmethod
-    # pyre-fixme[3]: Return annotation cannot be `Any`.
     def eval_step(self, state: State, data: TEvalData) -> Any:
         """
         Core required method for user to implement. This method will be called at each iteration of the
@@ -597,7 +595,6 @@ def on_predict_epoch_start(self, state: State) -> None:
         pass
 
     @abstractmethod
-    # pyre-fixme[3]: Return annotation cannot be `Any`.
     def predict_step(self, state: State, data: TPredictData) -> Any:
         """
         Core required method for user to implement. This method will be called at each iteration of the

torchtnt/utils/data/iterators.py

@@ -32,7 +32,6 @@
 from python.migrations.py310 import StrEnum310
 
 try:
-    # pyre-ignore[21]: Could not find name `StrEnum` in `enum`
     from enum import StrEnum
 except ImportError:
 

torchtnt/utils/flops.py

@@ -147,7 +147,6 @@ def _conv_backward_flop_jit(
     return flop_count
 
 
-# pyre-fixme [5]
 flop_mapping: Dict[Callable[..., Any], Callable[[Tuple[Any], Tuple[Any]], Number]] = {
     aten.mm: _matmul_flop_jit,
     aten.matmul: _matmul_flop_jit,
@@ -224,8 +223,8 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 
     def __torch_dispatch__(
         self,
-        func: Callable[..., Any],  # pyre-fixme [2] func can be any func
-        types: Tuple[Any],  # pyre-fixme [2]
+        func: Callable[..., Any],
+        types: Tuple[Any],
         args=(),  # pyre-fixme [2]
         kwargs=None,  # pyre-fixme [2]
     ) -> PyTree:
@@ -242,7 +241,6 @@ def __torch_dispatch__(
 
         return rs
 
-    # pyre-fixme [3]
     def _create_backwards_push(self, name: str) -> Callable[..., Any]:
         class PushState(torch.autograd.Function):
             @staticmethod
@@ -265,7 +263,6 @@ def backward(ctx, *grad_outs):
         # using a function parameter.
         return PushState.apply
 
-    # pyre-fixme [3]
     def _create_backwards_pop(self, name: str) -> Callable[..., Any]:
         class PopState(torch.autograd.Function):
             @staticmethod
@@ -289,9 +286,8 @@ def backward(ctx, *grad_outs):
         # using a function parameter.
         return PopState.apply
 
-    # pyre-fixme [3] Return a callable function
     def _enter_module(self, name: str) -> Callable[..., Any]:
-        # pyre-fixme [2, 3]
+        # pyre-fixme [3]
         def f(module: torch.nn.Module, inputs: Tuple[Any]):
             parents = self._parents
             parents.append(name)
@@ -301,9 +297,8 @@ def f(module: torch.nn.Module, inputs: Tuple[Any]):
 
         return f
 
-    # pyre-fixme [3] Return a callable function
     def _exit_module(self, name: str) -> Callable[..., Any]:
-        # pyre-fixme [2, 3]
+        # pyre-fixme [3]
         def f(module: torch.nn.Module, inputs: Tuple[Any], outputs: Tuple[Any]):
             parents = self._parents
             assert parents[-1] == name