fix(cache): recompute tasks with missing local outputs

flytekit/core/base_task.py

@@ -327,7 +327,6 @@ def local_execute(
                 outputs_literal_map = LocalTaskCache.get(
                     self.name, self.metadata.cache_version, input_literal_map, self.metadata.cache_ignore_input_vars
                 )
-                # The cache returns None iff the key does not exist in the cache
                 if outputs_literal_map is None:
                     logger.info("Cache miss, task will be executed now")
                 else:

flytekit/core/local_cache.py

@@ -1,10 +1,13 @@
-from typing import Optional, Tuple
+from typing import Iterable, Optional, Tuple
 
 from diskcache import Cache
 from flyteidl.core.literals_pb2 import LiteralMap
+from fsspec.utils import get_protocol
 
 from flytekit import lazy_module
-from flytekit.models.literals import Literal, LiteralCollection
+from flytekit.core.local_fsspec import FlyteLocalFileSystem
+from flytekit.loggers import logger
+from flytekit.models.literals import Blob, Literal, LiteralCollection, Schema, StructuredDataset
 from flytekit.models.literals import LiteralMap as ModelLiteralMap
 
 joblib = lazy_module("joblib")
@@ -50,6 +53,33 @@ def _calculate_cache_key(
     return f"{task_name}-{cache_version}-{joblib.hash(hashed_inputs)}"
 
 
+def _get_missing_local_artifact_uri(literal: Literal) -> Optional[str]:
+    children: Iterable[Literal]
+    if literal.collection:
+        children = literal.collection.literals
+    elif literal.map:
+        children = literal.map.literals.values()
+    elif literal.scalar and literal.scalar.union:
+        children = (literal.scalar.union.value,)
+    else:
+        if literal.scalar:
+            value: object = literal.scalar.value
+            if isinstance(value, (Blob, Schema, StructuredDataset)):
+                uri: str = value.uri
+                if get_protocol(uri) in FlyteLocalFileSystem.protocol:
+                    try:
+                        FlyteLocalFileSystem().info(path=uri)
+                    except (FileNotFoundError, NotADirectoryError):
+                        return uri
+        return None
+
+    for child in children:
+        missing_uri: Optional[str] = _get_missing_local_artifact_uri(literal=child)
+        if missing_uri is not None:
+            return missing_uri
+    return None
+
+
 class LocalTaskCache(object):
     """
     This class implements a persistent store able to cache the result of local task executions.
@@ -73,6 +103,7 @@ def clear():
     def get(
         task_name: str, cache_version: str, input_literal_map: ModelLiteralMap, cache_ignore_input_vars: Tuple[str, ...]
     ) -> Optional[ModelLiteralMap]:
+        """Return cached outputs, treating missing local artifact URIs as a cache miss."""
         if not LocalTaskCache._initialized:
             LocalTaskCache.initialize()
         serialized_obj = LocalTaskCache._cache.get(
@@ -85,17 +116,25 @@ def get(
         # If the serialized object is a model file, first convert it back to a proto object (which will force it to
         # use the installed flyteidl proto messages) and then convert it to a model object. This will guarantee
         # that the object is in the correct format.
+        literal_map: ModelLiteralMap
         if isinstance(serialized_obj, ModelLiteralMap):
-            return ModelLiteralMap.from_flyte_idl(ModelLiteralMap.to_flyte_idl(serialized_obj))
+            literal_map = ModelLiteralMap.from_flyte_idl(ModelLiteralMap.to_flyte_idl(serialized_obj))
         elif isinstance(serialized_obj, bytes):
             # If it is a bytes object, then it is a serialized proto object.
             # We need to convert it to a model object first.o
             pb_literal_map = LiteralMap()
             pb_literal_map.ParseFromString(serialized_obj)
-            return ModelLiteralMap.from_flyte_idl(pb_literal_map)
+            literal_map = ModelLiteralMap.from_flyte_idl(pb_literal_map)
         else:
             raise ValueError(f"Unexpected object type {type(serialized_obj)}")
 
+        for literal in literal_map.literals.values():
+            missing_uri: Optional[str] = _get_missing_local_artifact_uri(literal=literal)
+            if missing_uri is not None:
+                logger.warning(f"Ignoring local cache for {task_name}: missing artifact {missing_uri}")
+                return None
+        return literal_map
+
     @staticmethod
     def set(
         task_name: str,

tests/flytekit/unit/core/test_local_cache.py

@@ -19,12 +19,27 @@
 from flytekit.core.dynamic_workflow_task import dynamic
 from flytekit.core.hash import HashMethod
 from flytekit.core.local_cache import LocalTaskCache, _calculate_cache_key, _recursive_hash_placement
+from flytekit.core.local_fsspec import FlyteLocalFileSystem
 from flytekit.core.task import TaskMetadata, task
 from flytekit.core.testing import task_mock
 from flytekit.core.type_engine import TypeEngine
 from flytekit.core.workflow import workflow
-from flytekit.models.literals import Literal, LiteralCollection, LiteralMap, Primitive, Scalar
-from flytekit.models.types import LiteralType, SimpleType
+from flytekit.models.core.types import BlobType
+from flytekit.models.literals import (
+    Blob,
+    BlobMetadata,
+    Literal,
+    LiteralCollection,
+    LiteralMap,
+    Primitive,
+    Scalar,
+    Schema,
+    StructuredDataset,
+    StructuredDatasetMetadata,
+    Union,
+)
+from flytekit.models.types import LiteralType, SchemaType, SimpleType, StructuredDatasetType
+from flytekit.types.file import FlyteFile
 from flytekit.types.schema import FlyteSchema
 
 # Global counter used to validate number of calls to cache
@@ -649,3 +664,154 @@ def test_cache_old_version_of_literal_map():
     # Now load the same object from the cache and confirm that the `_offloaded_metadata` attribute is now present
     loaded_literal_map = LocalTaskCache.get("t.produce_dc", "1", LiteralMap(literals={}), ())
     assert hasattr(loaded_literal_map.literals['o0'], "_offloaded_metadata") is True
+
+
+@pytest.mark.serial
+def test_cached_workflow_recreates_missing_file(tmp_path: pathlib.Path) -> None:
+    calls = 0
+    artifact = tmp_path / "artifact.txt"
+
+    @task(cache=True, cache_version="missing-file-v1")
+    def produce() -> FlyteFile:
+        nonlocal calls
+        calls += 1
+        artifact.write_text("cached contents")
+        return FlyteFile(path=str(artifact))
+
+    @workflow
+    def cached_workflow() -> FlyteFile:
+        return produce()
+
+    assert pathlib.Path(cached_workflow().download()).read_text() == "cached contents"
+    assert pathlib.Path(cached_workflow().download()).read_text() == "cached contents"
+    assert calls == 1
+
+    artifact.rename(tmp_path / "renamed.txt")
+
+    assert pathlib.Path(cached_workflow().download()).read_text() == "cached contents"
+    assert calls == 2
+
+
+@pytest.mark.serial
+@pytest.mark.parametrize("kind", ["blob", "schema", "structured_dataset"])
+@pytest.mark.parametrize("container", ["scalar", "collection", "map", "union"])
+@pytest.mark.parametrize("legacy", [False, True])
+@pytest.mark.parametrize("file_uri", [False, True])
+def test_cache_misses_for_removed_local_artifacts(
+    tmp_path: pathlib.Path, kind: str, container: str, legacy: bool, file_uri: bool
+) -> None:
+    artifact = tmp_path / "artifact"
+    uri = artifact.as_uri() if file_uri else str(artifact)
+    literal: Literal
+    literal_type: LiteralType
+    if kind == "blob":
+        artifact.write_text("cached contents")
+        blob_type = BlobType(format="", dimensionality=BlobType.BlobDimensionality.SINGLE)
+        literal = Literal(
+            scalar=Scalar(
+                blob=Blob(
+                    metadata=BlobMetadata(type=blob_type),
+                    uri=uri,
+                )
+            )
+        )
+        literal_type = LiteralType(blob=blob_type)
+    elif kind == "schema":
+        artifact.mkdir()
+        schema_type = SchemaType(columns=[])
+        literal = Literal(scalar=Scalar(schema=Schema(uri=uri, type=schema_type)))
+        literal_type = LiteralType(schema=schema_type)
+    else:
+        artifact.mkdir()
+        dataset_type = StructuredDatasetType()
+        literal = Literal(
+            scalar=Scalar(
+                structured_dataset=StructuredDataset(
+                    uri=uri,
+                    metadata=StructuredDatasetMetadata(structured_dataset_type=dataset_type),
+                )
+            )
+        )
+        literal_type = LiteralType(structured_dataset_type=dataset_type)
+    if container == "collection":
+        literal = Literal(collection=LiteralCollection(literals=[literal]))
+    elif container == "map":
+        literal = Literal(map=LiteralMap(literals={"artifact": literal}))
+    elif container == "union":
+        literal = Literal(scalar=Scalar(union=Union(value=literal, stored_type=literal_type)))
+
+    inputs = LiteralMap(literals={})
+    outputs = LiteralMap(literals={"o0": literal})
+    if legacy:
+        cache_key: str = _calculate_cache_key(
+            task_name="missing-artifact", cache_version="v1", input_literal_map=inputs
+        )
+        LocalTaskCache._cache.set(key=cache_key, value=outputs)
+    else:
+        LocalTaskCache.set(
+            task_name="missing-artifact", cache_version="v1", input_literal_map=inputs,
+            cache_ignore_input_vars=(), value=outputs,
+        )
+
+    assert LocalTaskCache.get(
+        task_name="missing-artifact", cache_version="v1", input_literal_map=inputs,
+        cache_ignore_input_vars=(),
+    ) == outputs
+    if artifact.is_dir():
+        artifact.rmdir()
+    else:
+        artifact.unlink()
+    assert LocalTaskCache.get(
+        task_name="missing-artifact", cache_version="v1", input_literal_map=inputs,
+        cache_ignore_input_vars=(),
+    ) is None
+
+
+@pytest.mark.serial
+def test_remote_artifact_does_not_need_a_local_path() -> None:
+    inputs = LiteralMap(literals={})
+    outputs = LiteralMap(
+        literals={
+            "o0": Literal(
+                scalar=Scalar(
+                    blob=Blob(
+                        metadata=BlobMetadata(type=BlobType(format="", dimensionality=BlobType.BlobDimensionality.SINGLE)),
+                        uri="s3://bucket/artifact",
+                    )
+                )
+            )
+        }
+    )
+    LocalTaskCache.set(
+        task_name="remote-artifact", cache_version="v1", input_literal_map=inputs,
+        cache_ignore_input_vars=(), value=outputs,
+    )
+    assert LocalTaskCache.get(
+        task_name="remote-artifact", cache_version="v1", input_literal_map=inputs,
+        cache_ignore_input_vars=(),
+    ) == outputs
+
+
+@pytest.mark.serial
+def test_local_artifact_permission_errors_are_not_cache_misses(
+    tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    artifact = tmp_path / "artifact.txt"
+    artifact.write_text("cached contents")
+
+    @task(cache=True, cache_version="permission-v1")
+    def produce() -> FlyteFile:
+        return FlyteFile(path=str(artifact))
+
+    @workflow
+    def cached_workflow() -> FlyteFile:
+        return produce()
+
+    cached_workflow()
+
+    def fail_info(self: FlyteLocalFileSystem, path: str, **kwargs: typing.Any) -> typing.NoReturn:
+        raise PermissionError("artifact access denied")
+
+    monkeypatch.setattr(target=FlyteLocalFileSystem, name="info", value=fail_info)
+    with pytest.raises(PermissionError, match="artifact access denied"):
+        cached_workflow()