add sanity check

apache · Sep 30, 2021 · 59e5a19 · 59e5a19
1 parent 90602bb
commit 59e5a19
Showing 1 changed file with 20 additions and 0 deletions.
diff --git a/python/tvm/meta_schedule/runner/local_runner.py b/python/tvm/meta_schedule/runner/local_runner.py
@@ -198,6 +198,7 @@ def __init__(
             timeout=timeout_sec,
             initializer=initializer,
         )
+        self._sanity_check()
 
     def run(self, runner_inputs: List[RunnerInput]) -> List[RunnerFuture]:
         results: List[RunnerFuture] = []
@@ -228,6 +229,25 @@ def run(self, runner_inputs: List[RunnerInput]) -> List[RunnerFuture]:
             results.append(local_future)
         return results
 
+    def _sanity_check(self) -> None:
+        def _check(
+            f_alloc_argument,
+            f_run_evaluator,
+            f_cleanup,
+        ) -> None:
+            get_global_func_with_default_on_worker(name=f_alloc_argument, default=None)
+            get_global_func_with_default_on_worker(name=f_run_evaluator, default=None)
+            get_global_func_with_default_on_worker(name=f_cleanup, default=None)
+            tvm.get_global_func("tvm.contrib.random.random_fill", True)
+
+        value = self.pool.submit(
+            _check,
+            self.f_alloc_argument,
+            self.f_run_evaluator,
+            self.f_cleanup,
+        )
+        value.result()
+
     @staticmethod
     def _worker_func(
         _f_alloc_argument: Optional[str],