allenai · epwalsh · Oct 10, 2022 · Oct 10, 2022 · Oct 10, 2022 · Oct 10, 2022
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
@@ -219,8 +219,12 @@ jobs:
         run: |
           . .venv/bin/activate
           which python
+          echo "========="
           python --version
+          echo "========="
           pip freeze
+          echo "========="
+          tango info
 
       - name: ${{ matrix.task.name }}
         run: |

diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Fixed
 
 - Fixed bug with `TorchEvalStep` when constructing callbacks.
+- Fixed some import error issues caused when an integration is not installed.
 
 ### Changed
 

diff --git a/tango/__main__.py b/tango/__main__.py
@@ -85,7 +85,7 @@
 import click
 from click_help_colors import HelpColorsCommand, HelpColorsGroup
 
-from tango.common.exceptions import CliRunError
+from tango.common.exceptions import CliRunError, IntegrationMissingError
 from tango.common.logging import (
     cli_logger,
     initialize_logging,
@@ -479,7 +479,7 @@ def info(settings: TangoGlobalSettings):
         is_installed = True
         try:
             import_module_and_submodules(integration)
-        except (ModuleNotFoundError, ImportError):
+        except (IntegrationMissingError, ModuleNotFoundError, ImportError):
             is_installed = False
         if is_installed:
             cli_logger.info(" [green]\N{check mark} %s[/]", name)

diff --git a/tango/common/registrable.py b/tango/common/registrable.py
@@ -188,9 +188,9 @@ def search_modules(cls: Type[_RegistrableT], name: str):
         ):
             return None
 
-        def try_import(module):
+        def try_import(module, recursive: bool = True):
             try:
-                import_module_and_submodules(module)
+                import_module_and_submodules(module, recursive=recursive)
             except IntegrationMissingError:
                 pass
             except ImportError as e:
@@ -200,7 +200,7 @@ def try_import(module):
         integrations = {m.split(".")[-1]: m for m in find_integrations()}
         integrations_imported: Set[str] = set()
         if name in integrations:
-            try_import(integrations[name])
+            try_import(integrations[name], recursive=False)
             integrations_imported.add(name)
             if name in Registrable._registry[cls]:
                 return None
@@ -209,7 +209,7 @@ def try_import(module):
             # Try to guess the integration that it comes from.
             maybe_integration = name.split("::")[0]
             if maybe_integration in integrations:
-                try_import(integrations[maybe_integration])
+                try_import(integrations[maybe_integration], recursive=False)
                 integrations_imported.add(maybe_integration)
                 if name in Registrable._registry[cls]:
                     return None
@@ -248,7 +248,7 @@ def try_import(module):
         # Try importing all other integrations.
         for integration_name, module in integrations.items():
             if integration_name not in integrations_imported:
-                try_import(module)
+                try_import(module, recursive=False)
                 integrations_imported.add(integration_name)
                 if name in Registrable._registry[cls]:
                     return None

diff --git a/tango/common/util.py b/tango/common/util.py
@@ -75,7 +75,9 @@ def resolve_module_name(package_name: str) -> Tuple[str, Path]:
     return package_name, base_path
 
 
-def import_module_and_submodules(package_name: str, exclude: Optional[Set[str]] = None) -> None:
+def import_module_and_submodules(
+    package_name: str, exclude: Optional[Set[str]] = None, recursive: bool = True
+) -> None:
     """
     Import all submodules under the given package.
 
@@ -105,14 +107,15 @@ def import_module_and_submodules(package_name: str, exclude: Optional[Set[str]]
     path = getattr(module, "__path__", [])
     path_string = "" if not path else path[0]
 
-    # walk_packages only finds immediate children, so need to recurse.
-    for module_finder, name, _ in pkgutil.walk_packages(path):
-        # Sometimes when you import third-party libraries that are on your path,
-        # `pkgutil.walk_packages` returns those too, so we need to skip them.
-        if path_string and module_finder.path != path_string:  # type: ignore[union-attr]
-            continue
-        subpackage = f"{package_name}.{name}"
-        import_module_and_submodules(subpackage, exclude=exclude)
+    if recursive:
+        # walk_packages only finds immediate children, so need to recurse.
+        for module_finder, name, _ in pkgutil.walk_packages(path):
+            # Sometimes when you import third-party libraries that are on your path,
+            # `pkgutil.walk_packages` returns those too, so we need to skip them.
+            if path_string and module_finder.path != path_string:  # type: ignore[union-attr]
+                continue
+            subpackage = f"{package_name}.{name}"
+            import_module_and_submodules(subpackage, exclude=exclude)
 
 
 def _parse_bool(value: Union[bool, str]) -> bool:

diff --git a/tango/integrations/beaker/__init__.py b/tango/integrations/beaker/__init__.py
@@ -10,8 +10,8 @@
 from tango.common.exceptions import IntegrationMissingError
 
 try:
-    import beaker
-except ModuleNotFoundError:
+    from beaker import Beaker
+except (ModuleNotFoundError, ImportError):
     raise IntegrationMissingError("beaker", dependencies={"beaker-py"})
 
 from .executor import (

diff --git a/tango/integrations/wandb/__init__.py b/tango/integrations/wandb/__init__.py
@@ -33,15 +33,22 @@
 from .step_cache import WandbStepCache
 from .workspace import WandbWorkspace
 
+try:
+    import torch
+except ModuleNotFoundError:
+    pass
+else:
+    from .torch_train_callback import WandbTrainCallback
+
+    __all__.append("WandbTrainCallback")
+
 try:
     import flax
     import jax
     import tensorflow  # flax has a tensorflow dependency
-    import torch
 except ModuleNotFoundError:
     pass
 else:
     from .flax_train_callback import WandbFlaxTrainCallback
-    from .torch_train_callback import WandbTrainCallback
 
-    __all__.append("WandbTrainCallback")
+    __all__.append("WandbFlaxTrainCallback")
diff --git a/tests/common/registrable_test.py b/tests/common/registrable_test.py
@@ -47,3 +47,6 @@ def test_registering_step_by_reserved_name(self):
             @Step.register("ref")
             class BadStep(Step):
                 pass
+
+    def test_search_modules(self):
+        Step.search_modules("foo-bar-baz-non-existent")