diff --git a/flytekit/core/tracker.py b/flytekit/core/tracker.py index f2e6437bed..c30af041d5 100644 --- a/flytekit/core/tracker.py +++ b/flytekit/core/tracker.py @@ -55,7 +55,10 @@ def _get_module_from_main(globals) -> Optional[str]: # make sure current directory is in the PYTHONPATH. sys.path.insert(0, str(curdir)) - return import_module_from_file(module_name, file) + try: + return import_module_from_file(module_name, file) + except ModuleNotFoundError: + return None @staticmethod def _find_instance_module(): @@ -63,12 +66,17 @@ def _find_instance_module(): while frame: if frame.f_code.co_name == "" and "__name__" in frame.f_globals: if frame.f_globals["__name__"] != "__main__": - return frame.f_globals["__name__"], frame.f_globals["__file__"] - # if the remote_deploy command is invoked in the same module as where - # the app is defined, get the module from the file name + return frame.f_globals["__name__"], None + + # Try to find the module and filename in the case that we're in the __main__ module + # This is useful in cases that use FlyteRemote to load tasks/workflows that are defined + # in the same file as where FlyteRemote is being invoked to register and execute Flyte + # entities. One such case is with the `eager` decorator in the flytekit.experimental module. mod = InstanceTrackingMeta._get_module_from_main(frame.f_globals) if mod is None: return None, None + + # This is used in find_lhs to find the trackable instances when the module is loaded from the __file__. return mod.__name__, mod.__file__ frame = frame.f_back return None, None @@ -137,27 +145,32 @@ def find_lhs(self) -> str: # continue looping through m. logger.warning("Caught ValueError {} while attempting to auto-assign name".format(err)) - # try to find object in module when the tracked instance is defined in the __main__ module - module = import_module_from_file(self._instantiated_in, self._module_file) - - def _candidate_name_matches(candidate) -> bool: - if not hasattr(candidate, "name") or not hasattr(self, "name"): - return False - return candidate.name == self.name - - for k in dir(module): - try: - candidate = getattr(module, k) - # consider the variable equivalent to self if it's of the same type, name - if ( - type(candidate) == type(self) - and _candidate_name_matches(candidate) - and candidate.instantiated_in == self.instantiated_in - ): - self._lhs = k - return k - except ValueError as err: - logger.warning(f"Caught ValueError {err} while attempting to auto-assign name") + # Try to find object in module when the tracked instance is defined in the __main__ module. + # This section tries to find the matching object in the module when the module is loaded from the __file__. + if self._module_file is not None: + + # Since the module loaded from the file is different from the original module that defined self, we need + # to match by variable name and type. + module = import_module_from_file(self._instantiated_in, self._module_file) + + def _candidate_name_matches(candidate) -> bool: + if not hasattr(candidate, "name") or not hasattr(self, "name"): + return False + return candidate.name == self.name + + for k in dir(module): + try: + candidate = getattr(module, k) + # consider the variable equivalent to self if it's of the same type and name + if ( + type(candidate) == type(self) + and _candidate_name_matches(candidate) + and candidate.instantiated_in == self.instantiated_in + ): + self._lhs = k + return k + except ValueError as err: + logger.warning(f"Caught ValueError {err} while attempting to auto-assign name") logger.error(f"Could not find LHS for {self} in {self._instantiated_in}") raise _system_exceptions.FlyteSystemException(f"Error looking for LHS in {self._instantiated_in}")