Fix recursive local module discovery in bundle serialization (#19124)

Author: desertaxle

AGENTS.md

@@ -74,7 +74,6 @@ prefect config view            # Inspect configuration
 
 - GitHub issues are used for tracking issues (use the `gh` cli)
 - Pre-commit hooks required (never use `--no-verify`)
-- There are some slower pre-push hooks that may modify files on `git push`; when that happens, run `git commit --amend` to bring those into the prior commit (never use `--amend` in any other situation unless asked)
 - Dependencies: updates to client-side deps in `@pyproject.toml` require parallel changes ing `@client/pyproject.toml`
 - AGENTS.md always symlinked to CLAUDE.md
 - the redis lease storage lives in @src/integrations/prefect-redis/

src/prefect/_experimental/bundles/__init__.py

@@ -209,28 +209,49 @@ def _discover_local_dependencies(
 
     module_name = flow_module.__name__
 
+    # Process the flow's module and all its dependencies recursively
+    _process_module_dependencies(flow_module, module_name, local_modules, visited)
+
+    return local_modules
+
+
+def _process_module_dependencies(
+    module: ModuleType,
+    module_name: str,
+    local_modules: set[str],
+    visited: set[str],
+) -> None:
+    """
+    Recursively process a module and discover its local dependencies.
+
+    Args:
+        module: The module to process.
+        module_name: The name of the module.
+        local_modules: Set to accumulate discovered local modules.
+        visited: Set of already visited modules to avoid infinite recursion.
+    """
     # Skip if we've already processed this module
     if module_name in visited:
-        return local_modules
+        return
     visited.add(module_name)
 
-    # Check if the flow's module itself is local
-    module_file = getattr(flow_module, "__file__", None)
+    # Check if this is a local module
+    module_file = getattr(module, "__file__", None)
     if not module_file or not _is_local_module(module_name, module_file):
-        return local_modules
+        return
 
     local_modules.add(module_name)
 
     # Get the source code of the module
     try:
-        source_code = inspect.getsource(flow_module)
+        source_code = inspect.getsource(module)
     except (OSError, TypeError):
-        # Can't get source for the flow's module
-        return local_modules
+        # Can't get source for this module
+        return
 
     imports = _extract_imports_from_source(source_code)
 
-    # Check each import to see if it's local
+    # Check each import to see if it's local and recursively process it
     for import_name in imports:
         # Skip if already visited
         if import_name in visited:
@@ -249,30 +270,13 @@ def _discover_local_dependencies(
             else:
                 imported_module = importlib.import_module(import_name)
         except (ImportError, AttributeError):
-            # Can't import or module has no __file__, skip it
-            continue
-
-        imported_file = getattr(imported_module, "__file__", None)
-        if not imported_file or not _is_local_module(import_name, imported_file):
+            # Can't import, skip it
             continue
 
-        local_modules.add(import_name)
-        visited.add(import_name)
-
-        # Recursively check this module's dependencies
-        try:
-            imported_source = inspect.getsource(imported_module)
-        except (OSError, TypeError):
-            # Can't get source for this module, skip its dependencies
-            continue
-
-        nested_imports = _extract_imports_from_source(imported_source)
-        for nested_import in nested_imports:
-            if nested_import not in visited and _is_local_module(nested_import):
-                local_modules.add(nested_import)
-                visited.add(nested_import)
-
-    return local_modules
+        # Recursively process this imported module
+        _process_module_dependencies(
+            imported_module, import_name, local_modules, visited
+        )
 
 
 @contextmanager

tests/experimental/test_bundles.py

@@ -2,6 +2,7 @@
 import signal
 import subprocess
 import sys
+from pathlib import Path
 from typing import Any, Literal
 from unittest.mock import MagicMock, patch
 
@@ -490,7 +491,9 @@ def mock_unregister(module):
                         assert len(unregistered) == 1
                         assert unregistered[0] == mock_module
 
-    def test_pickle_local_modules_handles_import_errors(self, caplog):
+    def test_pickle_local_modules_handles_import_errors(
+        self, caplog: pytest.LogCaptureFixture
+    ):
         """Test that import errors are handled gracefully."""
 
         @flow
@@ -506,3 +509,98 @@ def test_flow():
 
             # Check that a debug message was logged about the failure
             assert "Failed to register module nonexistent_module" in caplog.text
+
+    def test_discover_deeply_nested_local_dependencies(self, tmp_path: Path):
+        """Test that local dependencies are discovered recursively through multiple levels.
+
+        This tests the scenario where:
+        - flow_module imports module_b
+        - module_b imports module_c
+        - module_c imports module_d
+
+        All modules should be discovered, including module_d which is 3 levels deep.
+        """
+        # Create temporary package structure with deep nesting
+        package_root = tmp_path / "test_packages"
+        package_root.mkdir()
+
+        # Create flow_module package
+        flow_pkg = package_root / "flow_module"
+        flow_pkg.mkdir()
+        (flow_pkg / "__init__.py").write_text("")
+
+        # Create module_b package
+        module_b_pkg = package_root / "module_b"
+        module_b_pkg.mkdir()
+        (module_b_pkg / "__init__.py").write_text("")
+
+        # Create module_c package
+        module_c_pkg = package_root / "module_c"
+        module_c_pkg.mkdir()
+        (module_c_pkg / "__init__.py").write_text("")
+
+        # Create module_d package (deepest level)
+        module_d_pkg = package_root / "module_d"
+        module_d_pkg.mkdir()
+        (module_d_pkg / "__init__.py").write_text("")
+
+        # Create module_d with a simple function
+        (module_d_pkg / "utils.py").write_text("""
+def function_d():
+    return "d"
+""")
+
+        # Create module_c that imports from module_d
+        (module_c_pkg / "utils.py").write_text("""
+from module_d.utils import function_d
+
+def function_c():
+    return function_d()
+""")
+
+        # Create module_b that imports from module_c
+        (module_b_pkg / "utils.py").write_text("""
+from module_c.utils import function_c
+
+def function_b():
+    return function_c()
+""")
+
+        # Create flow_module that imports from module_b
+        (flow_pkg / "my_flow.py").write_text("""
+from module_b.utils import function_b
+from prefect import flow
+
+@flow
+def test_flow():
+    return function_b()
+""")
+
+        # Add package_root to sys.path so modules can be imported
+        sys.path.insert(0, str(package_root))
+
+        try:
+            # Import the flow module and get the flow
+            import flow_module.my_flow
+
+            flow_obj = flow_module.my_flow.test_flow
+
+            # Discover dependencies
+            deps = _discover_local_dependencies(flow_obj)
+
+            # All four modules should be discovered
+            assert "flow_module.my_flow" in deps, (
+                "Flow module itself should be discovered"
+            )
+            assert "module_b.utils" in deps, "First-level import should be discovered"
+            assert "module_c.utils" in deps, "Second-level import should be discovered"
+            assert "module_d.utils" in deps, "Third-level import should be discovered"
+
+        finally:
+            # Clean up sys.path and sys.modules
+            sys.path.remove(str(package_root))
+            for module in list(sys.modules.keys()):
+                if module.startswith(
+                    ("flow_module", "module_b", "module_c", "module_d")
+                ):
+                    del sys.modules[module]