add option to run registry bootstrap in initialization

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

Reviewed By: itomatik

Differential Revision: D39631903

fbshipit-source-id: 1668a8b06260d02b40208b3dda3cbade0a12bc16

.github/workflows/workflows.yaml

@@ -31,12 +31,16 @@ jobs:
           pip install git+https://github.com/facebookresearch/detectron2.git
           pip install git+https://github.com/facebookresearch/mobile-vision
           pip install scikit-learn
-          pip install pytest nbval
+          pip install pytest pytest-xdist nbval
           pip install -e .
 
       - name: Run pytest
         run: |
-          python -m unittest discover -v -s tests
+          python -m pytest -n 4 --durations=15 -sv tests/ --ignore=tests/skip_init/
+
+      - name: Run pytest (skip initializer)
+        run: |
+          D2GO_IMPORT_SKIP_INITIALIZATION=1 python -m pytest -n 4 --durations=15 -sv tests/skip_init/
 
       - name: Test Notebooks
         run: |

d2go/__init__.py

 #!/usr/bin/env python3
 # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
 
+import os
 
 from d2go.initializer import initialize_all
 
 # NOTE: by default a list of initializations will run whenever D2Go is first imported,
 # so that users don't need to do any manual iniitialization other than importing `d2go`.
-initialize_all()
+
+# Environment variable can be used to skip initialization for special cases like unit test
+skip_initialization = os.environ.get("D2GO_IMPORT_SKIP_INITIALIZATION", "0") == "1"
+
+if not skip_initialization:
+    initialize_all()

d2go/initializer.py

 #!/usr/bin/env python3
 # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+from d2go.registry.bootstrap import bootstrap_registries
 from mobile_cv.common.misc.oss_utils import fb_overwritable
 
 _INITIALIZED = False
 
 
-def initialize_all():
+def initialize_all(boostrap_registries: bool = False) -> None:
     global _INITIALIZED
     if _INITIALIZED:
         return
     _INITIALIZED = True
 
-    _initialize_all()
+    _initialize_all(boostrap_registries=boostrap_registries)
 
 
-def _initialize_all():
+def _initialize_all(boostrap_registries: bool) -> None:
     _setup_env()
     _register_builtin_datasets()
+
     _populate_registries()
+    if boostrap_registries:
+        bootstrap_registries(enable_cache=True, catch_exception=True)
 
 
 # fmt: off

d2go/registry/bootstrap.py

@@ -54,6 +54,7 @@ def _get_cache_dir() -> str:
     except (OSError, AssertionError):
         tmp_dir = os.path.join(tempfile.gettempdir(), "d2go_cache")
         logger.warning(f"{cache_dir} is not accessible! Using {tmp_dir} instead!")
+        os.makedirs(tmp_dir, exist_ok=True)
         cache_dir = tmp_dir
     return cache_dir
 
@@ -363,8 +364,9 @@ def lazy_on_bootstrap(f: Callable) -> Callable:
 
 def _load_cached_results(filename: str) -> Dict[str, CachedResult]:
     with open(filename) as f:
-        loaded = yaml.safe_load(f)
-    assert isinstance(loaded, dict), f"Wrong format: {filename}"
+        content = f.read()
+        loaded = yaml.safe_load(content)
+    assert isinstance(loaded, dict), f"Wrong format: {content}"
     results = {
         filename: CachedResult(**result_dic) for filename, result_dic in loaded.items()
     }

tests/skip_init/test_registries_bootstrap.py

+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+import importlib
+import sys
+import unittest
+
+from d2go.initializer import initialize_all
+from d2go.registry import bootstrap
+
+# manual initialize without bootstrap
+initialize_all(boostrap_registries=False)
+
+
+def _unimport(package_name):
+    # remove sub modules from sys
+    modules = [
+        key
+        for key in sys.modules
+        if (
+            (key == package_name or key.startswith(package_name + "."))
+            # previent the parent package of this file being removed
+            and not __name__.startswith(key)
+        )
+    ]
+    for key in sorted(modules, reverse=True):
+        sys.modules.pop(key)
+
+    # invalidate the cache of removed sub modules
+    importlib.invalidate_caches()
+
+
+class TestRegistryBootstrap(unittest.TestCase):
+    def setUp(self):
+        # NOTE: reload this file since the imported modules (eg. `d2go.registry.bootstrap`)
+        # might be "unimported" during `tearDown`.
+        importlib.reload(sys.modules[__name__])
+
+    def tearDown(self):
+        # NOTE: "unimport" bootstrapped libraries, so that each test runs like starting
+        # a new python program.
+
+        # TODO: match list with the bootstrapped packages
+        _unimport("d2go.registry")
+        _unimport("mobile_cv")
+        _unimport("detectron2")
+
+    def test_bootstrap_core_lib(self):
+        self.assertFalse(bootstrap._IS_BOOTSTRAPPED)
+        bootstrap.bootstrap_registries(enable_cache=False, catch_exception=True)
+        self.assertTrue(bootstrap._IS_BOOTSTRAPPED)
+
+    def test_bootstrap_with_cache(self):
+        self.assertFalse(bootstrap._IS_BOOTSTRAPPED)
+        bootstrap.bootstrap_registries(enable_cache=True, catch_exception=True)
+        self.assertTrue(bootstrap._IS_BOOTSTRAPPED)