Add logging for storing cache in pipeline and blocks

TeamEpochGithub · May 15, 2024 · c3f6206 · c3f6206
1 parent 53ac91b
commit c3f6206
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 16 deletions.
diff --git a/epochalyst/pipeline/model/training/training.py b/epochalyst/pipeline/model/training/training.py
@@ -63,6 +63,7 @@ def train(self, x: Any, y: Any, cache_args: CacheArgs | None = None, **train_arg
         x, y = super().train(x, y, **train_args)
 
         if cache_args:
+            self.log_to_terminal(f"Storing cache for x and y to {cache_args["storage_path"]}")
             self._store_cache(name=self.get_hash() + "x", data=x, cache_args=cache_args)
             self._store_cache(name=self.get_hash() + "y", data=y, cache_args=cache_args)
 
@@ -115,7 +116,9 @@ def predict(self, x: Any, cache_args: CacheArgs | None = None, **pred_args: Any)
 
         x = super().predict(x, **pred_args)
 
-        self._store_cache(self.get_hash() + "p", x, cache_args) if cache_args else None
+        if cache_args:
+            self.log_to_terminal(f"Storing cache for x to {cache_args["storage_path"]}")
+            self._store_cache(self.get_hash() + "p", x, cache_args)
 
         # Set steps to original in case class is called again
         self.steps = self.all_steps

diff --git a/epochalyst/pipeline/model/training/training_block.py b/epochalyst/pipeline/model/training/training_block.py
@@ -76,16 +76,18 @@ def train(self, x: Any, y: Any, cache_args: CacheArgs | None = None, **train_arg
 
         x, y = self.custom_train(x, y, **train_args)
 
-        self._store_cache(
-            name=self.get_hash() + "x",
-            data=x,
-            cache_args=cache_args,
-        ) if cache_args else None
-        self._store_cache(
-            name=self.get_hash() + "y",
-            data=y,
-            cache_args=cache_args,
-        ) if cache_args else None
+        if cache_args:
+            self.log_to_terminal(f"Storing cache for x and y to {cache_args["storage_path"]}")
+            self._store_cache(
+                name=self.get_hash() + "x",
+                data=x,
+                cache_args=cache_args,
+            )
+            self._store_cache(
+                name=self.get_hash() + "y",
+                data=y,
+                cache_args=cache_args,
+            )
 
         return x, y
 
@@ -116,11 +118,13 @@ def predict(self, x: Any, cache_args: CacheArgs | None = None, **pred_args: Any)
 
         x = self.custom_predict(x, **pred_args)
 
-        self._store_cache(
-            name=self.get_hash() + "p",
-            data=x,
-            cache_args=cache_args,
-        ) if cache_args else None
+        if cache_args:
+            self.log_to_terminal(f"Store cache for predictions to {cache_args["storage_path"]}")
+            self._store_cache(
+                name=self.get_hash() + "p",
+                data=x,
+                cache_args=cache_args,
+            )
 
         return x
 

diff --git a/epochalyst/pipeline/model/transformation/transformation.py b/epochalyst/pipeline/model/transformation/transformation.py
@@ -105,6 +105,7 @@ def transform(self, data: Any, cache_args: CacheArgs | None = None, **transform_
         data = super().transform(data, **transform_args)
 
         if cache_args:
+            self.log_to_terminal(f"Storing cache for pipeline to {cache_args["storage_path"]}")
             self._store_cache(self.get_hash(), data, cache_args)
 
         # Set steps to original in case class is called again

diff --git a/epochalyst/pipeline/model/transformation/transformation_block.py b/epochalyst/pipeline/model/transformation/transformation_block.py
@@ -73,6 +73,7 @@ def transform(self, data: Any, cache_args: CacheArgs | None = None, **transform_
 
         data = self.custom_transform(data, **transform_args)
         if cache_args:
+            self.log_to_terminal(f"Storing cache to {cache_args["storage_path"]}")
             self._store_cache(name=self.get_hash(), data=data, cache_args=cache_args)
         return data