fix(code-bundle): rebuild after source changes

Remove argument-only async LRU caching from code bundle
builders so each invocation re-hashes the current source
files. Keep persistent cache reuse for unchanged bundles
and add regression test

Signed-off-by: Jeff Chung <sh1001309@gmail.com>

Author: jeffoodchain

examples/integration_tests.py

@@ -26,18 +26,12 @@
 import pytest
 
 import flyte
-from flyte._code_bundle import build_code_bundle
 
 # =============================================================================
 # FIXTURES
 # =============================================================================
 
 
-@pytest.fixture(autouse=True)
-def clear_lru_caches():
-    build_code_bundle.cache_clear()
-
-
 @pytest.fixture(scope="session")
 def flyte_client():
     """

src/flyte/_code_bundle/bundle.py

@@ -20,8 +20,6 @@
 except ImportError:  # pragma: no cover - Windows local dev; task containers are POSIX
     fcntl = None  # type: ignore[assignment]  # ty: ignore[invalid-assignment]
 
-from async_lru import alru_cache
-
 from flyte._logging import log, logger
 from flyte._status import status
 from flyte._utils import AsyncLRUCache
@@ -184,7 +182,6 @@ async def build_pkl_bundle(
             return CodeBundle(pkl=str(dest), computed_version=str_digest)
 
 
-@alru_cache
 async def build_code_bundle(
     from_dir: Path,
     *ignore: Type[Ignore],
@@ -292,7 +289,6 @@ async def build_code_bundle(
         return CodeBundle(tgz=remote_path, destination=extract_dir, computed_version=hash_digest, files=files)
 
 
-@alru_cache
 async def build_code_bundle_from_relative_paths(
     relative_paths: tuple[str, ...],
     from_dir: Path,

tests/flyte/app/test_app_code_bundling.py

@@ -50,7 +50,7 @@ async def test_code_bundle_consistency_with_include_files(temp_app_directory):
         from_dir=temp_app_directory,
         dryrun=True,
     )
-    build_code_bundle_from_relative_paths.cache_clear()
+
     bundle2 = await build_code_bundle_from_relative_paths(
         relative_paths=include_files,
         from_dir=temp_app_directory,
@@ -70,7 +70,7 @@ async def test_code_bundle_consistency_with_subdirectory_files(temp_app_director
         from_dir=temp_app_directory,
         dryrun=True,
     )
-    build_code_bundle_from_relative_paths.cache_clear()
+
     bundle2 = await build_code_bundle_from_relative_paths(
         relative_paths=include_files,
         from_dir=temp_app_directory,
@@ -89,7 +89,7 @@ async def test_code_bundle_different_files_produce_different_versions(temp_app_d
         from_dir=temp_app_directory,
         dryrun=True,
     )
-    build_code_bundle_from_relative_paths.cache_clear()
+
     bundle2 = await build_code_bundle_from_relative_paths(
         relative_paths=("app.py", "utils.py"),
         from_dir=temp_app_directory,
@@ -106,7 +106,7 @@ async def test_code_bundle_file_content_changes_version(temp_app_directory):
         from_dir=temp_app_directory,
         dryrun=True,
     )
-    build_code_bundle_from_relative_paths.cache_clear()
+
     (temp_app_directory / "app.py").write_text("print('Modified content!')")
     bundle2 = await build_code_bundle_from_relative_paths(
         relative_paths=("app.py",),

tests/flyte/code_bundle/test_build_code_bundle.py

@@ -24,10 +24,7 @@
 
 import pytest
 
-from flyte._code_bundle.bundle import (
-    build_code_bundle,
-    build_code_bundle_from_relative_paths,
-)
+from flyte._code_bundle.bundle import build_code_bundle
 
 TESTDATA_ROOT = Path(__file__).parent / "testdata"
 
@@ -103,20 +100,6 @@ def _load(path: Path, name: str) -> ModuleType:
             sys.modules.pop(name, None)
 
 
-@pytest.fixture(autouse=True)
-def _clear_bundle_cache() -> Iterator[None]:
-    """
-    ``build_code_bundle`` is wrapped in ``alru_cache`` at module level. Clear
-    it before and after each test so cached results from one layout never leak
-    into another — and so determinism tests can re-run the same inputs.
-    """
-    build_code_bundle.cache_clear()
-    build_code_bundle_from_relative_paths.cache_clear()
-    yield
-    build_code_bundle.cache_clear()
-    build_code_bundle_from_relative_paths.cache_clear()
-
-
 # ---------------------------------------------------------------------------
 # Single-file layouts
 # ---------------------------------------------------------------------------
@@ -359,6 +342,40 @@ async def test_build_code_bundle_multi_file_loaded_modules_with_include(importer
 # ---------------------------------------------------------------------------
 
 
+@pytest.mark.asyncio
+async def test_build_code_bundle_refreshes_after_source_file_changes():
+    with tempfile.TemporaryDirectory() as tmp:
+        tmp_dir = Path(tmp)
+        layout = _copy_layout("single_file", tmp_dir)
+        bundle_out = _bundle_out(tmp_dir)
+
+        bundle_before = await build_code_bundle(
+            from_dir=layout,
+            dryrun=True,
+            copy_style="all",
+            copy_bundle_to=bundle_out,
+        )
+
+        updated_source = """\
+def fn(x: int) -> int:
+    return x + 1
+"""
+        (layout / "hello.py").write_text(updated_source)
+
+        bundle_after = await build_code_bundle(
+            from_dir=layout,
+            dryrun=True,
+            copy_style="all",
+            copy_bundle_to=bundle_out,
+        )
+
+        with tarfile.open(Path(bundle_after.tgz), "r:gz") as tar:
+            bundled_file = tar.extractfile("hello.py")
+            assert bundled_file is not None
+            assert bundled_file.read().decode("utf-8") == updated_source
+        assert bundle_after.computed_version != bundle_before.computed_version
+
+
 @pytest.mark.asyncio
 async def test_build_code_bundle_digest_is_content_addressed():
     """
@@ -399,8 +416,6 @@ async def test_build_code_bundle_include_changes_digest():
             copy_style="all",
             copy_bundle_to=_bundle_out(tmp_dir, "out_1"),
         )
-        # Clear so the second call doesn't return the cached result.
-        build_code_bundle.cache_clear()
 
         bundle_with_extra = await build_code_bundle(
             from_dir=layout,