Bugfix: Debugger breaks if should_save_tensor is called before collections are prepared (#372)

Author: NihalHarish

smdebug/core/hook.py

@@ -550,6 +550,14 @@ def _increment_step(self):
     # Called in the internal AWS codebase to determine
     # if a particular tensor value should be saved
     def should_save_tensor_or_collection(self, tensor_name: str, collection_name: str) -> bool:
+        if self.prepared_collections is False:
+            # always return false if an attempt to save a
+            # tensor is made before the collections are prepared.
+            # this can happen if the fn is called before callbacks are init.
+            self.logger.warning(
+                "Tensors cannot be saved with smdebug before callbacks are initialized."
+            )
+            return False
         if self._is_collection_being_saved_for_step(collection_name):
             return True
         return self.is_tensor_saved_for_step(tensor_name)

tests/tensorflow2/test_should_save_tensor.py

@@ -61,3 +61,11 @@ def test_should_save_tensor_with_custom_collection(out_dir):
         else:
             assert not hook.should_save_tensor_or_collection(layer_name, CollectionKeys.GRADIENTS)
             assert not hook.should_save_tensor_or_collection(layer_name, CollectionKeys.LAYERS)
+
+
+def test_should_save_tensor_behavior_without_prepare_collections(out_dir):
+    """Always return false if an attempt to save a tensor is made before the collections are prepared.
+    This can happen if the fn is called before callbacks are init."""
+    hook = smd.KerasHook(out_dir, save_config=SaveConfig(save_interval=3), save_all=True)
+    assert not hook.should_save_tensor_or_collection("dummy", CollectionKeys.GRADIENTS)
+    assert not hook.should_save_tensor_or_collection("dummy", CollectionKeys.LAYERS)