fix(code-bundle): rebuild after source changes

[jeffoodchain]

Fixes flyteorg/flyte#7923

Bug: build_code_bundle memoizes only its arguments → long-lived processes reuse stale source after files are edited

Summary

build_code_bundle and build_code_bundle_from_relative_paths are wrapped in @alru_cache. The cache key contains function arguments such as from_dir, copy_style, and additional_files, but it does not contain the selected files or their contents.

In a long-lived process, the first run() builds and caches a CodeBundle. If source files are then edited and the task is run or forked again with the same bundling arguments, alru_cache returns the original bundle before the working tree is scanned again.

The later run therefore executes the pre-edit source without rebuilding the bundle or emitting a warning.

Impact

• run() → edit source → fork() can execute the original source instead of the edited source.
• Agent debugging loops, interactive sessions, and other long-lived processes can silently retry stale code.
• The file-listing and content-digest code is never reached on an argument-cache hit, so the existing content-addressed cache cannot detect the edit.
• skip_cache=True does not provide a complete guarantee because repeated calls using the same skip_cache=True arguments can themselves be memoized.

Root cause

src/flyte/_code_bundle/bundle.py memoizes the entire bundling operation:

@alru_cache
async def build_code_bundle(
    from_dir: Path,
    *,
    copy_style: CopyFiles = "loaded_modules",
    skip_cache: bool = False,
    additional_files: tuple[str, ...] = (),
) -> CodeBundle:
    ...

The same applies to the explicit-path builder:

@alru_cache
async def build_code_bundle_from_relative_paths(
    relative_paths: tuple[str, ...],
    from_dir: Path,
    *,
    skip_cache: bool = False,
) -> CodeBundle:
    ...

The content-sensitive portion of the operation runs inside the memoized function:

files, digest = list_files_to_bundle(
    from_dir,
    True,
    *ignore,
    copy_style=copy_style,
    additional_files=additional_files or None,
)

On an alru_cache hit, execution returns before list_files_to_bundle or list_relative_files_to_bundle can recompute the digest.

The bundler already has a content-addressed persistent cache after file discovery:

cached = _read_bundle_cache(digest)
if cached:
    hash_digest, remote_path = cached
    return CodeBundle(
        tgz=remote_path,
        destination=extract_dir,
        computed_version=hash_digest,
        files=files,
    )

This change removes the outer argument-only memoization while preserving the digest-backed persistent cache. Each invocation now scans and hashes the current source state; unchanged source still reuses the existing uploaded bundle and skips compression and upload.

Regression tests modify a source file between calls made with identical arguments and verify that:

• the new tarball contains the edited source;
• computed_version changes;
• both code-bundle builders detect edits without callers invoking cache_clear().

Test script result

I ran this locally on my Mac (M1 pro), and here are the logs.

cc @popojk

[flyte] >> Building 1 image...
[flyte]   >> Building image flyte for environment bundle_cache_check
[flyte]   -- Image localhost:30000/flyte:1f0330ec1afa0fcf7873c2b864d56435 already exists, skipping build
[flyte] OK Built image for environment bundle_cache_check: localhost:30000/flyte:1f0330ec1afa0fcf7873c2b864d56435
[flyte] OK Code bundle found in cache, skipping upload
[flyte] >> Waiting for run 'rcdkbdjbd9w8cb9dkdxc'...
[flyte] -- Run 'rcdkbdjbd9w8cb9dkdxc': ActionPhase.QUEUED (0:00:00.140879 secs, attempt 1)
[flyte] -- Run 'rcdkbdjbd9w8cb9dkdxc': ActionPhase.INITIALIZING (0:00:00.726295 secs, attempt 1)
[flyte] -- Run 'rcdkbdjbd9w8cb9dkdxc': ActionPhase.RUNNING (0:00:02.116723 secs, attempt 1)
[flyte] -- Run 'rcdkbdjbd9w8cb9dkdxc': ActionPhase.SUCCEEDED (0:00:03.179049 secs, attempt 1)
[flyte] OK Run 'rcdkbdjbd9w8cb9dkdxc' completed successfully
[flyte] >> Building 1 image...
[flyte]   >> Building image flyte for environment bundle_cache_check
[flyte] OK Built image for environment bundle_cache_check: localhost:30000/flyte:1f0330ec1afa0fcf7873c2b864d56435
[flyte] OK Code bundle found in cache, skipping upload
[flyte] >> Waiting for run 'rhn7rvrxncj7vvzx2hxw'...
[flyte] -- Run 'rhn7rvrxncj7vvzx2hxw': ActionPhase.QUEUED (0:00:00.030412 secs, attempt 1)
[flyte] -- Run 'rhn7rvrxncj7vvzx2hxw': ActionPhase.INITIALIZING (0:00:01.109061 secs, attempt 1)
[flyte] -- Run 'rhn7rvrxncj7vvzx2hxw': ActionPhase.RUNNING (0:00:02.157140 secs, attempt 1)
[flyte] -- Run 'rhn7rvrxncj7vvzx2hxw': ActionPhase.SUCCEEDED (0:00:03.199646 secs, attempt 1)
[flyte] OK Run 'rhn7rvrxncj7vvzx2hxw' completed successfully
[flyte] >> Building 1 image...
[flyte]   >> Building image flyte for environment bundle_cache_check
[flyte] OK Built image for environment bundle_cache_check: localhost:30000/flyte:1f0330ec1afa0fcf7873c2b864d56435
[flyte] OK Code bundle found in cache, skipping upload
[flyte] >> Waiting for run 'r5dg2b292vczp7gzwphj'...
[flyte] -- Run 'r5dg2b292vczp7gzwphj': ActionPhase.QUEUED (0:00:00.022630 secs, attempt 1)
[flyte] -- Run 'r5dg2b292vczp7gzwphj': ActionPhase.INITIALIZING (0:00:01.118085 secs, attempt 1)
[flyte] -- Run 'r5dg2b292vczp7gzwphj': ActionPhase.RUNNING (0:00:02.430836 secs, attempt 1)
[flyte] -- Run 'r5dg2b292vczp7gzwphj': ActionPhase.SUCCEEDED (0:00:04.374561 secs, attempt 1)
[flyte] OK Run 'r5dg2b292vczp7gzwphj' completed successfully

===== run #1  initial =====
run #1  initial -> remote returned: 'ActionOutputs(o0="v1-ORIGINAL")'  (http://localhost:30080/v2/domain/development/project/flytesnacks/runs/rcdkbdjbd9w8cb9dkdxc)

>>> rewrote VALUE to v2-EDITED (same process, no restart)

===== run #2  after edit =====
run #2  after edit -> remote returned: 'ActionOutputs(o0="v2-EDITED")'  (http://localhost:30080/v2/domain/development/project/flytesnacks/runs/rhn7rvrxncj7vvzx2hxw)

===== run #3  unchanged, expect a cache hit =====
run #3  unchanged, expect a cache hit -> remote returned: 'ActionOutputs(o0="v2-EDITED")'  (http://localhost:30080/v2/domain/development/project/flytesnacks/runs/r5dg2b292vczp7gzwphj)

================ result ================
run #1 saw the original value : True
edit picked up on rerun       : True  PASS
unchanged rerun is consistent : True  PASS

(restored bundle_cache_check.py)

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-09-02T09:53:24.722296Z	2c4af12	New commits

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[popojk]

@jeffoodchain I used this script to test in devbox and it works, could you please:

1. test it in local devbox and paste the result in PR description
2. sign off your commits
   cc @cosmicBboy @pingsutw for second eye

import pathlib
import re

import flyte

# The probe constant. The script rewrites this exact line between runs; the task
# reads it as a module global, so the container reports whatever was bundled.
VALUE = "v1-ORIGINAL"

_VALUE_LINE = re.compile(r'^VALUE = ".*"$', re.MULTILINE)
_SELF = pathlib.Path(__file__)

image = (
    flyte.Image.from_debian_base()
    # Same pin the other devbox examples use; drop it if your devbox does not need it.
    .with_pip_packages("cryptography==44.0.3", "pyOpenSSL==25.1.0")
)

env = flyte.TaskEnvironment(
    name="bundle_cache_check",
    resources=flyte.Resources(cpu=1, memory="1Gi"),
    image=image,
)


@env.task
def report() -> str:
    return VALUE


def run_once(label: str) -> str:
    print(f"\n===== {label} =====")
    # mode="remote" is not optional here: a local run never builds a code bundle at
    # all, so leaving the mode to config defaults can make every check below pass
    # while exercising nothing.
    run = flyte.with_runcontext(mode="remote").run(report)
    run.wait()
    out = str(run.outputs())
    print(f"{label} -> remote returned: {out!r}  ({run.url})")
    return out


def set_value(new_value: str) -> None:
    src = _SELF.read_text()
    _SELF.write_text(_VALUE_LINE.sub(f'VALUE = "{new_value}"', src, count=1))


if __name__ == "__main__":
    flyte.init_from_config()
    original = _SELF.read_text()
    try:
        first = run_once("run #1  initial")

        set_value("v2-EDITED")
        print("\n>>> rewrote VALUE to v2-EDITED (same process, no restart)")
        second = run_once("run #2  after edit")

        third = run_once("run #3  unchanged, expect a cache hit")

        rebuilt = "v2-EDITED" in second
        stable = second == third
        print("\n================ result ================")
        print(f"run #1 saw the original value : {'v1-ORIGINAL' in first}")
        print(f"edit picked up on rerun       : {rebuilt}  {'PASS' if rebuilt else 'FAIL <- issue #7923'}")
        print(f"unchanged rerun is consistent : {stable}  {'PASS' if stable else 'FAIL'}")
    finally:
        _SELF.write_text(original)
        print(f"\n(restored {_SELF.name})")

[jeffoodchain]

@popojk thank you. I've added the result to PR description and also signoff.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 874068d9bc

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Retain a content-aware cache for dry-run bundles

When dryrun=True and copy_bundle_to is unset, removing this cache makes every call rescan and recompress the source, allocate a new flyte-tmp-* directory, and leave the copied archive there; the SQLite cache cannot help because its lookup is explicitly skipped for dry runs. This is a production path in connectors/_connector.py, where each remote-storage execution calls build_code_bundle(..., dryrun=True), so a long-lived connector worker now accumulates one full code bundle per execution and can eventually exhaust local disk while repeatedly blocking on compression. Preserve reuse after the content digest is recomputed, or clean up the generated bundle after the connector uploads it.

Useful? React with 👍 / 👎.

[jeffoodchain]

I think this might be an issue, @popojk do you think we should solve this in this PR as well?

[popojk]

@jeffoodchain Please also solve this in this PR