support lazy_func during bootstrap

Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/287

In previous D36798327 (https://github.com/facebookresearch/d2go/commit/1f45cf04c69c63f51387b073ae931832b7a9510f), we still have 7 files which have exception during bootstrap. Instead of tweaking the the mock logic to support those scenarios, since they're all related to import-time computation, I think it's reasonable to "delay" those computations, therefore this diff introduces the `lazy_func` decorator, which delays it for bootstrap.

After this change, now all files are be bootstrapped without exception, next step we'll change `catch_exception` to False and put a unit test on this to prevent committing code that is not bootstrap-able.

Reviewed By: tglik

Differential Revision: D36814147

fbshipit-source-id: e344a5ccee76b1d9962117d300926a88c066e051

d2go/modeling/modeldef/modeldef.py

@@ -5,6 +5,7 @@
 import copy
 
 from d2go.modeling.modeldef.fbnet_modeldef_registry import FBNetV2ModelArch
+from d2go.registry.bootstrap import lazy_on_bootstrap
 from mobile_cv.arch.fbnet_v2.modeldef_utils import _ex, e1, e1p, e2, e3, e4, e6
 
 
@@ -14,6 +15,7 @@ def _mutated_tuple(tp, pos, value):
     return tuple(tp_list)
 
 
+@lazy_on_bootstrap
 def _repeat_last(stage, n=None):
     """
     Repeat the last "layer" of given stage, i.e. a (op_type, c, s, n_repeat, ...)

d2go/registry/bootstrap.py

@@ -14,7 +14,7 @@ import traceback
 from collections import defaultdict
 from dataclasses import asdict, dataclass
 from enum import Enum
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Callable, Dict, List, Optional, Tuple
 
 import pkg_resources
 import yaml
@@ -329,6 +329,27 @@ def break_bootstrap():
     return
 
 
+def lazy_on_bootstrap(f: Callable) -> Callable:
+    """
+    A decorator to mark a function as "lazy" during bootstrap, such that the decorated
+    function will skip the execution and immediately return a MagicMock object during
+    the bootstrap (the decorator is a non-op outside of bootstrap). This can be used to
+    hide un-executable code (usually related to import-time computation) during the
+    bootstrap.
+
+    For registration related import-time computation, please consider using the
+    `LazyRegisterable` since it will also save time for the normal import.
+    """
+
+    def wrapped(*args, **kwargs):
+        if _INSIDE_BOOTSTRAP:
+            return MoreMagicMock()
+        else:
+            return f(*args, **kwargs)
+
+    return wrapped
+
+
 def _load_cached_results(filename: str) -> Dict[str, CachedResult]:
     with open(filename) as f:
         loaded = yaml.safe_load(f)