diff --git a/flytekit/core/base_task.py b/flytekit/core/base_task.py index 1e544331d0..da92acbe45 100644 --- a/flytekit/core/base_task.py +++ b/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: diff --git a/flytekit/core/local_cache.py b/flytekit/core/local_cache.py index d6c7f93f99..660bdad6fe 100644 --- a/flytekit/core/local_cache.py +++ b/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, diff --git a/tests/flytekit/unit/core/test_local_cache.py b/tests/flytekit/unit/core/test_local_cache.py index 0990541a84..8d2177212c 100644 --- a/tests/flytekit/unit/core/test_local_cache.py +++ b/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()