[temp] fixed memory issue

Temporary commit to see if the memory issue is fixed.
pangeo-forge · Oct 4, 2021 · 0f9ce09 · 0f9ce09
1 parent 262b240
commit 0f9ce09
Showing 1 changed file with 9 additions and 4 deletions.
diff --git a/pangeo_forge_recipes/recipes/base.py b/pangeo_forge_recipes/recipes/base.py
@@ -125,14 +125,16 @@ def to_dask(self):
         import dask
         from dask.delayed import Delayed
 
-        # TODO: HighlevelGraph layers for each of these mapped inputs.
-        # Cache Input --------------------------------------------------------
         dsk = {}
         token = dask.base.tokenize(self)
 
+        # TODO: HighlevelGraph layers for each of these mapped inputs.
+        # Cache Input --------------------------------------------------------
         # TODO: allow recipes to customize which stages to run
+        cache_input_token = f"cache_input-{token}"
+        dsk[cache_input_token] = self.cache_input
         for i, input_key in enumerate(self.iter_inputs()):
-            dsk[(f"cache_input-{token}", i)] = (self.cache_input, input_key)
+            dsk[(f"cache_input-{token}", i)] = (cache_input_token, input_key)
 
         # Prepare Target ------------------------------------------------------
         dsk[f"checkpoint_0-{token}"] = (lambda *args: None, list(dsk))
@@ -143,10 +145,13 @@ def to_dask(self):
         )
 
         # Store Chunk --------------------------------------------------------
+        store_chunk_token = f"cache_chunk-{token}"
+        dsk[store_chunk_token] = self.store_chunk
+
         keys = []
         for i, chunk_key in enumerate(self.iter_chunks()):
             k = (f"store_chunk-{token}", i)
-            dsk[k] = (_store_chunk, f"prepare_target-{token}", self.store_chunk, chunk_key)
+            dsk[k] = (_store_chunk, f"prepare_target-{token}", store_chunk_token, chunk_key)
             keys.append(k)
 
         # Finalize Target -----------------------------------------------------