diff --git a/wandb/integration/kfp/kfp_patch.py b/wandb/integration/kfp/kfp_patch.py index 367b03afd0f..33e9227c824 100644 --- a/wandb/integration/kfp/kfp_patch.py +++ b/wandb/integration/kfp/kfp_patch.py @@ -43,19 +43,17 @@ def full_path_exists(full_func): - def get_parent_child_pairs(full_func): - components = full_func.split(".") - parents, children = [], [] - for i, _ in enumerate(components[:-1], 1): - parent = ".".join(components[:i]) - child = components[i] - parents.append(parent) - children.append(child) - return zip(parents, children) - - for parent, child in get_parent_child_pairs(full_func): - module = wandb.util.get_module(parent) - if not module or not hasattr(module, child) or getattr(module, child) is None: + split = full_func.split(".") + get_module = wandb.util.get_module + for i in range(1, len(split)): + parent = ".".join(split[:i]) + child = split[i] + module = get_module(parent) + if module is None: + return False + # Only do attribute checks if module is not None + attr = getattr(module, child, None) + if attr is None: return False return True diff --git a/wandb/util.py b/wandb/util.py index 264e526ed0d..f36ce7087a1 100644 --- a/wandb/util.py +++ b/wandb/util.py @@ -65,6 +65,7 @@ from wandb.sdk.lib import filesystem, runid from wandb.sdk.lib.json_util import dump, dumps from wandb.sdk.lib.paths import FilePathStr, StrPath +from functools import lru_cache if TYPE_CHECKING: import wandb.sdk.internal.settings_static @@ -239,12 +240,13 @@ def get_module( :param (bool) lazy: If True, return a lazy loader for the module. :return: (module|None) If import succeeds, the module will be returned. """ + # Only cache when importable not previously reported if name not in _not_importable: try: - if not lazy: - return import_module(name) + if lazy: + return _get_module_lazy_cached(name) else: - return import_module_lazy(name) + return _get_module_nonlazy_cached(name) except Exception: _not_importable.add(name) msg = f"Error importing optional module {name}" @@ -2030,6 +2032,14 @@ def get_core_path() -> str: return str(bin_path) +@lru_cache(maxsize=128) +def _get_module_lazy_cached(name: str) -> Any: + return import_module_lazy(name) + +@lru_cache(maxsize=128) +def _get_module_nonlazy_cached(name: str) -> Any: + return import_module(name) + class NonOctalStringDumper(yaml.Dumper): """Prevents strings containing non-octal values like "008" and "009" from being converted to numbers in in the yaml string saved as the sweep config."""