[Tests] Ensure MyPy type checks pass (#9284)

* [Tests] Ensure MyPy type checks pass

There's a few errors that come up when type checking that aren't triggering any failures:
```
Checking MyPy Type defs in the meta schedule package.
python/tvm/meta_schedule/utils.py:23:1: error: Cannot find implementation or library stub for module named "psutil"
python/tvm/meta_schedule/utils.py:23:1: note: See https://mypy.readthedocs.io/en/stable/running_mypy.html#missing-imports
python/tvm/meta_schedule/search_strategy/search_strategy.py: note: In member "__init__" of class "MeasureCandidate":
python/tvm/meta_schedule/search_strategy/search_strategy.py:59:13: error: Module has no attribute "MeasureCandidate"
python/tvm/meta_schedule/search_strategy/search_strategy.py: note: In member "initialize_with_tune_context" of class "SearchStrategy":
python/tvm/meta_schedule/search_strategy/search_strategy.py:83:9: error: Module has no attribute "SearchStrategyInitializeWithTuneContext"
```

To rectify this the `types-psutil` package adds type hints for `mypy` and `# type: ignore` stops `mypy` from trying to figure out types of `_ffi_api` resources.

There's also a few places where variable type definitions are repeated even though they're only required once.

Finally, I've ensured `task_mypy.sh` fails the build since it's stable right now, using `set -e`.

* Add temporary # type : ignore for psutil

python/gen_requirements.py

@@ -198,6 +198,7 @@
                 "sphinx_autodoc_annotation",
                 "sphinx_gallery",
                 "sphinx_rtd_theme",
+                "types-psutil",
             ],
         ),
     ),

python/tvm/meta_schedule/runner/local_runner.py

@@ -239,13 +239,11 @@ def run(self, runner_inputs: List[RunnerInput]) -> List[RunnerFuture]:
                 result: List[float] = future.result()
                 error_message: str = None
             except TimeoutError as exception:
-                result: List[float] = None
-                error_message: str = (
-                    f"LocalRunner: Timeout, killed after {self.timeout_sec} seconds\n"
-                )
+                result = None
+                error_message = f"LocalRunner: Timeout, killed after {self.timeout_sec} seconds\n"
             except Exception as exception:  # pylint: disable=broad-except
-                result: List[float] = None
-                error_message: str = "LocalRunner: An exception occurred\n" + str(exception)
+                result = None
+                error_message = "LocalRunner: An exception occurred\n" + str(exception)
             local_future = LocalRunnerFuture(res=result, error_message=error_message)
             results.append(local_future)
         return results

python/tvm/meta_schedule/search_strategy/replay_trace.py

@@ -41,7 +41,7 @@ class ReplayTrace(SearchStrategy):
     def __init__(self, num_trials_per_iter: int, num_trials_total: int):
         """Constructor"""
         self.__init_handle_by_constructor__(
-            _ffi_api.ReplayTrace,  # pylint: disable=no-member
+            _ffi_api.ReplayTrace,  # type: ignore # pylint: disable=no-member
             num_trials_per_iter,
             num_trials_total,
         )

python/tvm/meta_schedule/search_strategy/search_strategy.py

@@ -56,7 +56,7 @@ def __init__(self, sch: Schedule, args_info: List[ArgInfo]) -> None:
             The argument information.
         """
         self.__init_handle_by_constructor__(
-            _ffi_api.MeasureCandidate,  # pylint: disable=no-member
+            _ffi_api.MeasureCandidate,  # type: ignore # pylint: disable=no-member
             sch,
             args_info,
         )
@@ -80,7 +80,7 @@ def initialize_with_tune_context(
         tune_context : TuneContext
             The tuning context for initialization.
         """
-        _ffi_api.SearchStrategyInitializeWithTuneContext(  # pylint: disable=no-member
+        _ffi_api.SearchStrategyInitializeWithTuneContext(  # type: ignore # pylint: disable=no-member
             self, tune_context
         )
 
@@ -92,11 +92,11 @@ def pre_tuning(self, design_spaces: List[Schedule]) -> None:
         design_spaces : List[Schedule]
             The design spaces for pre-tuning.
         """
-        _ffi_api.SearchStrategyPreTuning(self, design_spaces)  # pylint: disable=no-member
+        _ffi_api.SearchStrategyPreTuning(self, design_spaces)  # type: ignore # pylint: disable=no-member
 
     def post_tuning(self) -> None:
         """Post-tuning for the search strategy."""
-        _ffi_api.SearchStrategyPostTuning(self)  # pylint: disable=no-member
+        _ffi_api.SearchStrategyPostTuning(self)  # type: ignore # pylint: disable=no-member
 
     def generate_measure_candidates(self) -> Optional[List[MeasureCandidate]]:
         """Generate measure candidates from design spaces for measurement.
@@ -106,7 +106,7 @@ def generate_measure_candidates(self) -> Optional[List[MeasureCandidate]]:
         measure_candidates : Optional[List[IRModule]]
             The measure candidates generated, None if finished.
         """
-        return _ffi_api.SearchStrategyGenerateMeasureCandidates(self)  # pylint: disable=no-member
+        return _ffi_api.SearchStrategyGenerateMeasureCandidates(self)  # type: ignore # pylint: disable=no-member
 
     def notify_runner_results(self, results: List[RunnerResult]) -> None:
         """Update the search strategy with profiling results.
@@ -116,7 +116,7 @@ def notify_runner_results(self, results: List[RunnerResult]) -> None:
         results : List[RunnerResult]
             The profiling results from the runner.
         """
-        _ffi_api.SearchStrategyNotifyRunnerResults(self, results)  # pylint: disable=no-member
+        _ffi_api.SearchStrategyNotifyRunnerResults(self, results)  # type: ignore # pylint: disable=no-member
 
 
 @register_object("meta_schedule.PySearchStrategy")
@@ -142,7 +142,7 @@ def f_notify_runner_results(results: List["RunnerResult"]) -> None:
             self.notify_runner_results(results)
 
         self.__init_handle_by_constructor__(
-            _ffi_api.SearchStrategyPySearchStrategy,  # pylint: disable=no-member
+            _ffi_api.SearchStrategyPySearchStrategy,  # type: ignore # pylint: disable=no-member
             f_initialize_with_tune_context,
             f_pre_tuning,
             f_post_tuning,

python/tvm/meta_schedule/task_scheduler/task_scheduler.py

@@ -27,7 +27,7 @@ class TaskScheduler(Object):
 
     def tune(self) -> None:
         """Auto-tuning."""
-        _ffi_api.TaskSchedulerTune(self)  # pylint: disable=no-member
+        _ffi_api.TaskSchedulerTune(self)  # type: ignore # pylint: disable=no-member
 
     def _set_task_stopped(self, task_id: int) -> None:
         """Set specific task to be stopped.
@@ -37,7 +37,7 @@ def _set_task_stopped(self, task_id: int) -> None:
         task_id : int
             The task id to be stopped.
         """
-        _ffi_api.TaskSchedulerSetTaskStopped(self, task_id)  # pylint: disable=no-member
+        _ffi_api.TaskSchedulerSetTaskStopped(self, task_id)  # type: ignore # pylint: disable=no-member
 
     def _is_task_running(self, task_id: int) -> bool:
         """Check whether the task is running.
@@ -52,7 +52,7 @@ def _is_task_running(self, task_id: int) -> bool:
         bool
             Whether the task is running.
         """
-        return _ffi_api.TaskSchedulerIsTaskRunning(self, task_id)  # pylint: disable=no-member
+        return _ffi_api.TaskSchedulerIsTaskRunning(self, task_id)  # type: ignore # pylint: disable=no-member
 
     def _join_running_task(self, task_id: int) -> None:
         """Wait until the task is finished.
@@ -62,7 +62,7 @@ def _join_running_task(self, task_id: int) -> None:
         task_id : int
             The task id to be joined.
         """
-        _ffi_api.TaskSchedulerJoinRunningTask(self, task_id)  # pylint: disable=no-member
+        _ffi_api.TaskSchedulerJoinRunningTask(self, task_id)  # type: ignore # pylint: disable=no-member
 
     def _next_task_id(self) -> int:
         """Fetch the next task id.
@@ -72,7 +72,7 @@ def _next_task_id(self) -> int:
         int
             The next task id.
         """
-        return _ffi_api.TaskSchedulerNextTaskId(self)  # pylint: disable=no-member
+        return _ffi_api.TaskSchedulerNextTaskId(self)  # type: ignore # pylint: disable=no-member
 
 
 @register_object("meta_schedule.PyTaskScheduler")
@@ -98,7 +98,7 @@ def f_next_task_id() -> int:
             return self._next_task_id()
 
         self.__init_handle_by_constructor__(
-            _ffi_api.TaskSchedulerPyTaskScheduler,  # pylint: disable=no-member
+            _ffi_api.TaskSchedulerPyTaskScheduler,  # type: ignore # pylint: disable=no-member
             f_tune,
             f_set_task_stopped,
             f_is_task_running,
@@ -110,13 +110,13 @@ def tune(self) -> None:
         raise NotImplementedError()
 
     def _set_task_stopped(self, task_id: int) -> None:
-        _ffi_api.TaskSchedulerSetTaskStopped(self, task_id)  # pylint: disable=no-member
+        _ffi_api.TaskSchedulerSetTaskStopped(self, task_id)  # type: ignore # pylint: disable=no-member
 
     def _is_task_running(self, task_id: int) -> bool:
-        return _ffi_api.TaskSchedulerIsTaskRunning(self, task_id)  # pylint: disable=no-member
+        return _ffi_api.TaskSchedulerIsTaskRunning(self, task_id)  # type: ignore # pylint: disable=no-member
 
     def _join_running_task(self, task_id: int) -> None:
-        _ffi_api.TaskSchedulerJoinRunningTask(self, task_id)  # pylint: disable=no-member
+        _ffi_api.TaskSchedulerJoinRunningTask(self, task_id)  # type: ignore # pylint: disable=no-member
 
     def _next_task_id(self) -> int:
-        return _ffi_api.TaskSchedulerNextTaskId(self)  # pylint: disable=no-member
+        return _ffi_api.TaskSchedulerNextTaskId(self)  # type: ignore # pylint: disable=no-member

python/tvm/meta_schedule/utils.py

@@ -20,7 +20,7 @@
 import shutil
 from typing import Any, Callable, List, Optional, Union
 
-import psutil
+import psutil  # type: ignore
 import tvm
 from tvm._ffi import get_global_func, register_func
 from tvm.error import TVMError

tests/scripts/task_mypy.sh

@@ -15,6 +15,9 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+
+set -e
+set -u
 set -o pipefail
 
 echo "Checking MyPy Type defs in the TensorIR schedule package."