Read the local task cache through diskcache-community's SafeDisk

flytekit/core/local_cache.py

@@ -7,14 +7,16 @@
 from flytekit.models.literals import LiteralMap as ModelLiteralMap
 
 if TYPE_CHECKING:
-    from diskcache import Cache
+    from diskcache_community import Cache
 
 joblib = lazy_module("joblib")
 
 # Location on the filesystem where serialized objects will be stored
 # TODO: read from config
 CACHE_LOCATION = "~/.flyte/local-cache"
 
+_PICKLED_MODEL_MODULES = ("flytekit.models.literals", "flytekit.models.types", "flytekit.models.core.types")
+
 
 def _recursive_hash_placement(literal: Literal) -> Literal:
     # Base case, hash gets passed through always if set
@@ -62,11 +64,13 @@ class LocalTaskCache(object):
 
     @staticmethod
     def initialize():
-        # Import diskcache here rather than at module level, so that importing flytekit
+        # Import diskcache_community here rather than at module level, so that importing flytekit
         # does not require it unless the local cache is actually used.
-        from diskcache import Cache
+        from diskcache_community import Cache, SafeDisk
+        from google.protobuf.struct_pb2 import Struct
 
-        LocalTaskCache._cache = Cache(CACHE_LOCATION)
+        disk = SafeDisk.with_types(Struct, modules=_PICKLED_MODEL_MODULES).with_options(on_error="miss")
+        LocalTaskCache._cache = Cache(CACHE_LOCATION, disk=disk)
         LocalTaskCache._initialized = True
 
     @staticmethod

pyproject.toml

@@ -18,7 +18,7 @@ dependencies = [
     "cloudpickle>=2.0.0",
     "croniter>=0.3.20",
     "dataclasses-json>=0.5.2,<0.5.12", # TODO: remove upper-bound after fixing change in contract
-    "diskcache>=5.2.1",
+    "diskcache-community>=5.8.1",
     "docker>=4.0.0",
     "docstring-parser>=0.9.0",
     "flyteidl>=1.16.7,<2.0.0a0",

tests/flytekit/unit/core/test_local_cache.py

@@ -653,8 +653,8 @@ def test_cache_old_version_of_literal_map():
     assert hasattr(loaded_literal_map.literals['o0'], "_offloaded_metadata") is True
 
 
-def test_import_flytekit_does_not_import_diskcache():
-    # Use a fresh interpreter, since other tests in this process may already have imported diskcache.
+def test_import_flytekit_does_not_import_diskcache_community():
+    # Use a fresh interpreter, since other tests in this process may already have imported diskcache_community.
     code = textwrap.dedent(
         """
         import sys
@@ -663,20 +663,20 @@ def test_import_flytekit_does_not_import_diskcache():
         import flytekit.clis.sdk_in_container.pyflyte
         import flytekit.core.base_task
 
-        assert "diskcache" not in sys.modules, "diskcache was imported when loading flytekit"
+        assert "diskcache_community" not in sys.modules, "diskcache_community was imported when loading flytekit"
         """
     )
     subprocess.run([sys.executable, "-c", code], check=True)
 
 
-def test_import_flytekit_without_diskcache():
-    # Setting sys.modules["diskcache"] to None makes "import diskcache" raise ImportError,
-    # which simulates an environment where diskcache is not installed.
+def test_import_flytekit_without_diskcache_community():
+    # Setting sys.modules["diskcache_community"] to None makes "import diskcache_community" raise ImportError,
+    # which simulates an environment where diskcache_community is not installed.
     code = textwrap.dedent(
         """
         import sys
 
-        sys.modules["diskcache"] = None
+        sys.modules["diskcache_community"] = None
 
         import flytekit
         from flytekit.core.local_cache import LocalTaskCache
@@ -686,7 +686,7 @@ def test_import_flytekit_without_diskcache():
         except ImportError:
             pass
         else:
-            raise AssertionError("LocalTaskCache.initialize() should fail without diskcache")
+            raise AssertionError("LocalTaskCache.initialize() should fail without diskcache_community")
         """
     )
     subprocess.run([sys.executable, "-c", code], check=True)