index: add cache

iterative · May 7, 2023 · 734ade3 · 734ade3
1 parent 6b6084a
commit 734ade3
Showing 1 changed file with 20 additions and 1 deletion.
diff --git a/dvc/repo/index.py b/dvc/repo/index.py
@@ -141,6 +141,7 @@ def _load_data_from_outs(index, prefix, outs):
 def _load_storage_from_out(storage_map, key, out):
     from dvc.config import NoRemoteError
     from dvc_data.index import FileStorage, ObjectStorage
+    from fsspec.utils import tokenize
 
     if out.odb:
         storage_map.add_data(ObjectStorage(key, out.odb))
@@ -156,14 +157,32 @@ def _load_storage_from_out(storage_map, key, out):
                     index=remote.index,
                 )
             )
+
+            token = f"{out.fs.protocol}-{tokenize(out.fs_path)}" 
+            storage_map.add_cache(
+                FileStorage(
+                    key=key,
+                    fs=out.cache.fs,
+                    path=out.cache.fs.path.join(out.cache.path, "fs", token),
+                )
+            )
         else:
             storage_map.add_remote(ObjectStorage(key, remote.odb, index=remote.index))
     except NoRemoteError:
         pass
 
     if out.stage.is_import:
         dep = out.stage.deps[0]
-        storage_map.add_data(FileStorage(key, dep.fs, dep.fs_path))
+        storage_map.add_remote(FileStorage(key, dep.fs, dep.fs_path))
+
+        token = f"{dep.fs.protocol}-{tokenize(dep.fs_path)}"
+        storage_map.add_cache(
+            FileStorage(
+                key=key,
+                fs=out.cache.fs,
+                path=out.cache.fs.path.join(out.cache.path, "fs", token),
+            )
+        )
 
 
 class Index: