fix: root action cache key honors Literal.hash, like sub-actions do

Content-based caching worked between tasks but silently degraded to URI-based
caching at the run entrypoint. Re-running a cached root task with byte-identical
inputs uploaded to a fresh URI missed the cache; the same task invoked as a
sub-action hit.

The two paths derive the inputs hash differently. A sub-action's is computed
in-process by the controller via `generate_inputs_repr_for_literal`, which
substitutes `Literal.hash` for the literal's contents. A root action's comes from
`OffloadedInputData.inputs_hash`, which the backend computes over the *marshaled*
inputs (cloud `shared_service/cache/cache_key.go:HashInputs`, called from
`dataproxy/service/dataproxy_flyte.go:UploadInputs`) — folding in the offloaded
blob URI and ignoring `Literal.hash`. Every `flyte.run` uploads to a new URI, so
the key changed on every submission.

`_submit_remote` now recomputes that field over the same representation the
controller uses. The value is exactly what a sub-action would compute, so the two
land on the same cache key: the backend folds the field in as
`sha256(inputsHash + taskName + interfaceHash + cacheVersion)`
(cloud `workflow/service/utils.go:generateCacheKeyFromInputsHash`), which is the
same formula as `generate_cache_key_hash`. The shared byte builder is extracted
into `_named_literals_repr` so the two representations cannot drift.

Cache-ignored inputs are filtered before hashing, matching `filterInputsForHash`
on the backend's upload path. Without that, a task combining
`Cache(ignored_inputs=...)` with a content-hashed input would key on the very
inputs the user asked to exclude.

`generate_content_inputs_hash` returns None unless a hashed input survives that
filtering, leaving the backend's value in place. Without the gate every root run's
key would change and every cache entry already written against the old keys would
become unreachable; with it, only callers who opted into content hashing are
affected.

Overriding the field is safe because the backend treats it as an opaque string:
`generateCacheKeyFromInputsHash` concatenates and re-hashes it, never parsing,
length-checking, or verifying it against the blob. Two backend deployments even
encode it differently — sha256/std-base64 in `cache_key.go`, FNV-64a/url-base64
in `flyte2/dataproxy/service/dataproxy_service.go:hashInputsProto` — which only
works because nothing downstream cares.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QCM3MeG4tcFdFCnEiiHeik

Author: cosmicBboy

examples/advanced/hash_root_action.py

@@ -0,0 +1,87 @@
+"""
+Content-based caching for the *root* task of a run.
+
+`hash_flyte_dataframe.py` shows content-based caching between tasks: a driver produces a
+DataFrame and calls a cached consumer twice, and the second call hits. This example covers the
+other entrypoint — passing a locally-built DataFrame straight into `flyte.run(...)`, so the
+cached task *is* the root action of the run.
+
+The two paths compute the cache key differently. A sub-action's key is computed by the
+controller in-process, which substitutes `Literal.hash` for the literal's contents. A root
+action's key is derived from the offloaded inputs, which reference the upload URI — and every
+`flyte.run` uploads the DataFrame to a fresh URI. So without a content hash the second run
+misses even though the bytes are identical.
+
+Passing `hash_method=` to `DataFrame.from_local_sync` makes both runs agree: the key follows
+the content, not where it happened to land in blob storage.
+
+Run it twice; `check_cache_hit` below asserts the second run returns the first run's value.
+"""
+
+import pandas as pd
+
+import flyte
+from flyte import Cache
+from flyte.io import DataFrame, HashFunction
+
+img = flyte.Image.from_debian_base(name="flyte-root-hash").with_pip_packages("pandas", "pyarrow")
+
+env = flyte.TaskEnvironment(
+    "flyte_root_action_hash",
+    image=img,
+    resources=flyte.Resources(cpu="1", memory="2Gi"),
+)
+
+SAMPLE_DATA = {"id": [1, 2, 3, 4, 5], "value": [100, 200, 300, 400, 500]}
+
+
+def hash_pandas_dataframe(df: pd.DataFrame) -> str:
+    """Content-based hash: the same rows always produce the same digest."""
+    return str(pd.util.hash_pandas_object(df).sum())
+
+
+@env.task(cache=Cache(behavior="override", version_override="v1"))
+async def main(df: DataFrame) -> str:
+    """Cached root task.
+
+    The random number is the cache probe: it is regenerated on every real execution, so two
+    runs returning the same string can only mean the second one was served from the cache.
+    """
+    import random
+
+    pdf = await df.open(pd.DataFrame).all()
+    return f"rows={len(pdf)}, total={pdf['value'].sum()}, random={random.randint(1, 1000000)}"
+
+
+def build_input() -> DataFrame:
+    """The DataFrame to submit, tagged with a content-based hash.
+
+    Without `hash_method` the cache key would follow the (per-run, always new) upload URI and
+    the second run would miss.
+    """
+    return DataFrame.from_local_sync(
+        pd.DataFrame(SAMPLE_DATA),
+        hash_method=HashFunction.from_fn(hash_pandas_dataframe),
+    )
+
+
+if __name__ == "__main__":
+    flyte.init_from_config()
+
+    # Two independent submissions of the same content. Each uploads to its own URI.
+    run1 = flyte.run(main, df=build_input())
+    print(f"Run 1: {run1.url}")
+    run1.wait()
+    result1 = run1.outputs()[0]
+
+    run2 = flyte.run(main, df=build_input())
+    print(f"Run 2: {run2.url}")
+    run2.wait()
+    result2 = run2.outputs()[0]
+
+    print(f"\nRun 1: {result1}")
+    print(f"Run 2: {result2}")
+    if result1 == result2:
+        print("\n✓ Cache hit — the new upload URI did not change the cache key.")
+    else:
+        print("\n✗ Cache miss — the root action's key still tracks the upload URI.")

examples/integration_tests.py

@@ -393,6 +393,35 @@ async def test_advanced_local_tasks(flyte_client):
     await _run_and_wait(parallel_main_no_io, "test_advanced_local_tasks", q="hello")
 
 
+@pytest.mark.integration
+@pytest.mark.asyncio
+async def test_advanced_hash_root_action(flyte_client):
+    """Content-based caching must work when the cached task is the run's root action.
+
+    A sub-action's cache key is computed by the controller, which substitutes `Literal.hash`
+    for the literal's contents; a root action's key is derived from the offloaded inputs, which
+    reference the upload URI. Each `flyte.run` uploads to a fresh URI, so this only passes if
+    the content hash — not the URI — drives the key.
+    """
+    from examples.advanced.hash_root_action import build_input, main
+
+    results = []
+    for i in (1, 2):
+        # Rebuilt each time, so run 2 uploads identical bytes to a brand-new URI.
+        run = await flyte.with_runcontext(log_level=logging.DEBUG).run.aio(main, df=build_input())
+        print(f"\n[test_advanced_hash_root_action] Run {i}: {run.url}")
+        run.wait()
+        detail = await run.action.details()
+        if detail.error_info:
+            raise RuntimeError(f"Run {i} failed with error: {detail.error_info.message}")
+        results.append(run.outputs()[0])
+
+    # The task embeds a fresh random number on every real execution, so identical outputs
+    # mean run 2 was served from run 1's cache entry.
+    assert results[0] == results[1], f"root action cache miss across a new upload URI: {results[0]!r} != {results[1]!r}"
+    print("  Cache hit across a new upload URI\n")
+
+
 @pytest.mark.integration
 @pytest.mark.asyncio
 async def test_advanced_multi_loops(flyte_client):

src/flyte/_internal/runtime/convert.py

@@ -668,6 +668,20 @@ def generate_inputs_repr_for_literal(literal: literals_pb2.Literal) -> bytes:
     return literal.SerializeToString(deterministic=True)
 
 
+def _named_literals_repr(inputs: list[common_pb2.NamedLiteral]) -> bytes:
+    """Deterministic byte representation of named inputs, honoring any `Literal.hash`.
+
+    Names are folded in so argument order and naming matter, with `:`/`;` framing so distinct
+    inputs cannot concatenate into the same bytes.
+    """
+    combined_bytes = b""
+    for named_literal in inputs:
+        name_bytes = named_literal.name.encode("utf-8")
+        literal_bytes = generate_inputs_repr_for_literal(named_literal.value)
+        combined_bytes += name_bytes + b":" + literal_bytes + b";"
+    return combined_bytes
+
+
 def generate_inputs_hash_for_named_literals(
     inputs: list[common_pb2.NamedLiteral],
 ) -> str:
@@ -685,16 +699,52 @@ def generate_inputs_hash_for_named_literals(
     if not inputs:
         return ""
 
-    # Build the byte representation by concatenating each literal's representation
-    combined_bytes = b""
-    for named_literal in inputs:
-        # Add the name to ensure order matters
-        name_bytes = named_literal.name.encode("utf-8")
-        literal_bytes = generate_inputs_repr_for_literal(named_literal.value)
-        # Combine name and literal bytes with a separator to avoid collisions
-        combined_bytes += name_bytes + b":" + literal_bytes + b";"
+    return hash_data(_named_literals_repr(inputs))
+
 
-    return hash_data(combined_bytes)
+def literal_carries_hash(literal: literals_pb2.Literal) -> bool:
+    """Whether `literal` (or anything nested in it) carries a user-supplied content hash.
+
+    Mirrors the cases `generate_inputs_repr_for_literal` substitutes a hash for, so the two
+    stay in step: a literal that would change the repr is exactly one this reports True for.
+    """
+    if literal.hash:
+        return True
+    if literal.HasField("collection"):
+        return any(literal_carries_hash(nested) for nested in literal.collection.literals)
+    if literal.HasField("map"):
+        return any(literal_carries_hash(nested) for nested in literal.map.literals.values())
+    return False
+
+
+def generate_content_inputs_hash(
+    inputs: common_pb2.Inputs,
+    ignored_input_vars: List[str],
+) -> Optional[str]:
+    """Content-addressed replacement for `OffloadedInputData.inputs_hash`, or None to defer.
+
+    The backend fills that field by hashing the *marshaled* inputs, which folds in the offloaded
+    blob URI and ignores `Literal.hash` (cloud `shared_service/cache/cache_key.go:HashInputs`) —
+    so a root action re-run with identical content at a fresh upload URI misses the cache, while
+    the same task invoked as a sub-action hits.
+
+    Returns exactly what a sub-action would compute, so both land on the same cache key: the
+    backend derives it as `sha256(inputsHash + taskName + interfaceHash + cacheVersion)` (cloud
+    `workflow/service/utils.go:generateCacheKeyFromInputsHash`), the same formula as
+    `generate_cache_key_hash` below. Overriding a server-computed field is safe because it is
+    only ever concatenated and re-hashed — never parsed, nor checked against the blob.
+
+    `ignored_input_vars` is filtered out first, matching the backend's `filterInputsForHash`;
+    without it, inputs excluded via `Cache(ignored_inputs=...)` would leak into the key. None
+    when no hashed input survives that filtering — the common case, where deferring to the
+    backend keeps existing cache keys, and the entries written against them, valid.
+    """
+    if not inputs or not inputs.literals:
+        return None
+    literals = [named for named in inputs.literals if named.name not in ignored_input_vars]
+    if not any(literal_carries_hash(named.value) for named in literals):
+        return None
+    return generate_inputs_hash_for_named_literals(literals)
 
 
 def generate_inputs_hash_from_proto(inputs: common_pb2.Inputs) -> str:

src/flyte/_run.py

@@ -688,6 +688,7 @@ async def _submit_remote(
         from flyteidl2.workflow import run_service_pb2
 
         import flyte.errors
+        from flyte._internal.runtime.convert import generate_content_inputs_hash
         from flyte.remote import Run
 
         try:
@@ -711,6 +712,25 @@ async def _submit_remote(
                 upload_resp = await get_client().dataproxy_service.upload_inputs(upload_req)
                 offloaded_input_data = upload_resp.offloaded_input_data
 
+                # The hash the server derives from the marshaled inputs folds in the offloaded
+                # blob URI and ignores `Literal.hash`, so content-based caching silently degrades
+                # to URI-based caching at the run entrypoint: identical content uploaded to a
+                # fresh URI misses. Sub-actions don't have this problem — the controller hashes
+                # the same inputs through `generate_inputs_repr_for_literal`, which substitutes
+                # the content hash. Recompute over that representation so the root action agrees.
+                # Returns None (leaving the server's value alone) unless a hashed input survives
+                # cache-ignore filtering, so cache keys for everyone else are untouched.
+                #
+                # `task_spec` is populated on every path into here, including the by-reference
+                # one where `task_id` is also set (`task_spec = task.pb2.spec` on a fetched
+                # task), so the ignore list is always the registered task's own.
+                md = task_spec.task_template.metadata if task_spec is not None else None
+                content_hash = generate_content_inputs_hash(
+                    proto_inputs, list(md.cache_ignore_input_vars) if md else []
+                )
+                if content_hash is not None:
+                    offloaded_input_data.inputs_hash = content_hash
+
             create_req = run_service_pb2.CreateRunRequest(
                 run_id=run_id,
                 project_id=project_id,

tests/flyte/internal/runtime/test_convert.py

@@ -14,6 +14,8 @@
     LiteralMap,
     Primitive,
     Scalar,
+    StructuredDataset,
+    StructuredDatasetMetadata,
 )
 from flyteidl2.core.types_pb2 import (
     BlobType,
@@ -1593,3 +1595,139 @@ async def test_convert_inputs_no_kickoff_key_is_noop():
         out = await convert.convert_inputs_to_native(Inputs(proto_inputs=proto), interface)
 
     assert out == {"x": 7}
+
+
+# ---------------------------------------------------------------------------
+# Content-addressed inputs hash for root actions
+#
+# The backend derives `OffloadedInputData.inputs_hash` from the marshaled inputs, which folds
+# in the offloaded blob URI and ignores `Literal.hash`. Sub-actions don't go through that path
+# — the controller hashes via `generate_inputs_repr_for_literal`, which substitutes the content
+# hash — so content-based caching worked for sub-actions but degraded to URI-based caching at
+# the run entrypoint. `generate_content_inputs_hash` closes that gap.
+# ---------------------------------------------------------------------------
+
+_CONTENT_HASH = "sha256-of-geoparquet-bytes"
+
+
+def _sd_literal(uri: str, hash_val: str | None = None) -> Literal:
+    return Literal(
+        scalar=Scalar(
+            structured_dataset=StructuredDataset(
+                uri=uri,
+                metadata=StructuredDatasetMetadata(structured_dataset_type=StructuredDatasetType(format="parquet")),
+            )
+        ),
+        hash=hash_val,
+    )
+
+
+def _int_literal(v: int) -> Literal:
+    return Literal(scalar=Scalar(primitive=Primitive(integer=v)))
+
+
+def _named_inputs(**kwargs: Literal) -> _task_common_pb2.Inputs:
+    return _task_common_pb2.Inputs(
+        literals=[_task_common_pb2.NamedLiteral(name=name, value=lit) for name, lit in kwargs.items()]
+    )
+
+
+def test_content_inputs_hash_ignores_upload_uri():
+    """Identical content re-uploaded to a fresh URI must produce the same key."""
+    run1 = _named_inputs(aoi=_sd_literal("s3://bkt/run-1/abc/0", _CONTENT_HASH))
+    run2 = _named_inputs(aoi=_sd_literal("s3://bkt/run-2/xyz/0", _CONTENT_HASH))
+
+    assert convert.generate_content_inputs_hash(run1, []) == convert.generate_content_inputs_hash(run2, [])
+
+
+def test_content_inputs_hash_tracks_content():
+    same_uri_other_content = _named_inputs(aoi=_sd_literal("s3://bkt/run-1/abc/0", "a-different-digest"))
+    baseline = _named_inputs(aoi=_sd_literal("s3://bkt/run-1/abc/0", _CONTENT_HASH))
+
+    assert convert.generate_content_inputs_hash(baseline, []) != convert.generate_content_inputs_hash(
+        same_uri_other_content, []
+    )
+
+
+def test_content_inputs_hash_is_name_sensitive():
+    """Same literal bound to a different parameter is a different call."""
+    as_aoi = _named_inputs(aoi=_sd_literal("s3://bkt/1", _CONTENT_HASH))
+    as_other = _named_inputs(other=_sd_literal("s3://bkt/1", _CONTENT_HASH))
+
+    assert convert.generate_content_inputs_hash(as_aoi, []) != convert.generate_content_inputs_hash(as_other, [])
+
+
+@pytest.mark.parametrize(
+    "name,inputs",
+    [
+        ("empty", _task_common_pb2.Inputs()),
+        ("plain scalar", _named_inputs(x=Literal(scalar=Scalar(primitive=Primitive(integer=5))))),
+        ("dataframe without a hash", _named_inputs(aoi=_sd_literal("s3://bkt/run-1/abc/0"))),
+    ],
+)
+def test_content_inputs_hash_defers_when_no_input_is_hashed(name, inputs):
+    """None means "leave the backend's value alone", which keeps already-written cache entries
+    reachable for the overwhelmingly common case of no content hashes at all."""
+    assert convert.generate_content_inputs_hash(inputs, []) is None
+
+
+def test_content_inputs_hash_equals_the_sub_action_hash():
+    """The value must be exactly what the controller computes for a sub-action.
+
+    The backend folds this field into the cache key as
+    `sha256(inputsHash + taskName + interfaceHash + cacheVersion)`
+    (cloud `workflow/service/utils.go:generateCacheKeyFromInputsHash`), which is the same
+    formula as `generate_cache_key_hash`. Equal inputs hashes therefore mean equal cache keys,
+    so a root action and a sub-action of the same task share cache entries.
+    """
+    inputs = _named_inputs(aoi=_sd_literal("s3://bkt/1", _CONTENT_HASH))
+
+    assert convert.generate_content_inputs_hash(inputs, []) == convert.generate_inputs_hash_from_proto(inputs)
+
+
+def test_content_inputs_hash_excludes_cache_ignored_inputs():
+    """Matches `filterInputsForHash` on the backend's upload path.
+
+    Without this, a task combining `Cache(ignored_inputs=...)` with a content-hashed input
+    would key on the very inputs the user asked to exclude.
+    """
+    run1 = _named_inputs(aoi=_sd_literal("s3://bkt/1", _CONTENT_HASH), seed=_int_literal(1))
+    run2 = _named_inputs(aoi=_sd_literal("s3://bkt/1", _CONTENT_HASH), seed=_int_literal(2))
+
+    assert convert.generate_content_inputs_hash(run1, ["seed"]) == convert.generate_content_inputs_hash(run2, ["seed"])
+    # ...and without the ignore list, the differing input does move the key.
+    assert convert.generate_content_inputs_hash(run1, []) != convert.generate_content_inputs_hash(run2, [])
+
+
+def test_content_inputs_hash_defers_when_only_ignored_inputs_are_hashed():
+    """Nothing left to fix once the hashed input is filtered out — leave the backend's value."""
+    inputs = _named_inputs(aoi=_sd_literal("s3://bkt/1", _CONTENT_HASH), seed=_int_literal(1))
+
+    assert convert.generate_content_inputs_hash(inputs, ["aoi"]) is None
+
+
+@pytest.mark.parametrize(
+    "name,wrap",
+    [
+        ("collection", lambda lit: Literal(collection=LiteralCollection(literals=[lit]))),
+        ("map", lambda lit: Literal(map=LiteralMap(literals={"k": lit}))),
+    ],
+)
+def test_content_inputs_hash_sees_nested_hashes(name, wrap):
+    """`Literal.hash` nested in a collection/map counts, matching what the repr substitutes."""
+    run1 = _named_inputs(aoi=wrap(_sd_literal("s3://bkt/run-1/abc/0", _CONTENT_HASH)))
+    run2 = _named_inputs(aoi=wrap(_sd_literal("s3://bkt/run-2/xyz/0", _CONTENT_HASH)))
+
+    assert convert.generate_content_inputs_hash(run1, []) is not None
+    assert convert.generate_content_inputs_hash(run1, []) == convert.generate_content_inputs_hash(run2, [])
+
+
+def test_root_and_sub_action_agree_on_uri_independence():
+    """The property the fix is really about: both paths now ignore a changed upload URI."""
+    run1 = _named_inputs(aoi=_sd_literal("s3://bkt/run-1/abc/0", _CONTENT_HASH))
+    run2 = _named_inputs(aoi=_sd_literal("s3://bkt/run-2/xyz/0", _CONTENT_HASH))
+
+    # sub-action path (controller-side), unchanged by this fix
+    assert convert.generate_inputs_hash_from_proto(run1) == convert.generate_inputs_hash_from_proto(run2)
+    # root-action path (client-side), previously URI-sensitive via the server's hash
+    assert convert.generate_content_inputs_hash(run1, []) == convert.generate_content_inputs_hash(run2, [])