Raise Error For Invalid Collection Config (aws#162)

Author: NihalHarish

smdebug/core/collection.py

@@ -34,10 +34,10 @@ class CollectionKeys:
 
     OPTIMIZER_VARIABLES = "optimizer_variables"
     TENSORFLOW_SUMMARIES = "tensorflow_summaries"
+    METRICS = "metrics"
 
     # XGBOOST
     HYPERPARAMETERS = "hyperparameters"
-    METRICS = "metrics"
     PREDICTIONS = "predictions"
     LABELS = "labels"
     FEATURE_IMPORTANCE = "feature_importance"
@@ -65,6 +65,50 @@ class CollectionKeys:
 
 NON_HISTOGRAM_COLLECTIONS = SCALAR_COLLECTIONS.union(SUMMARIES_COLLECTIONS)
 
+DEFAULT_TF_COLLECTIONS = {
+    CollectionKeys.ALL,
+    CollectionKeys.DEFAULT,
+    CollectionKeys.WEIGHTS,
+    CollectionKeys.BIASES,
+    CollectionKeys.GRADIENTS,
+    CollectionKeys.LOSSES,
+    CollectionKeys.METRICS,
+    CollectionKeys.INPUTS,
+    CollectionKeys.OUTPUTS,
+    CollectionKeys.SM_METRICS,
+    CollectionKeys.OPTIMIZER_VARIABLES,
+}
+
+DEFAULT_PYTORCH_COLLECTIONS = {
+    CollectionKeys.ALL,
+    CollectionKeys.DEFAULT,
+    CollectionKeys.WEIGHTS,
+    CollectionKeys.BIASES,
+    CollectionKeys.GRADIENTS,
+    CollectionKeys.LOSSES,
+}
+
+DEFAULT_MXNET_COLLECTIONS = {
+    CollectionKeys.ALL,
+    CollectionKeys.DEFAULT,
+    CollectionKeys.WEIGHTS,
+    CollectionKeys.BIASES,
+    CollectionKeys.GRADIENTS,
+    CollectionKeys.LOSSES,
+}
+
+DEFAULT_XGBOOST_COLLECTIONS = {
+    CollectionKeys.ALL,
+    CollectionKeys.DEFAULT,
+    CollectionKeys.HYPERPARAMETERS,
+    CollectionKeys.PREDICTIONS,
+    CollectionKeys.LABELS,
+    CollectionKeys.FEATURE_IMPORTANCE,
+    CollectionKeys.AVERAGE_SHAP,
+    CollectionKeys.FULL_SHAP,
+    CollectionKeys.TREES,
+}
+
 
 class Collection:
     """

smdebug/core/hook.py

@@ -36,6 +36,7 @@
 from smdebug.core.state_store import StateStore
 from smdebug.core.utils import flatten, get_tb_worker, match_inc, size_and_shape
 from smdebug.core.writer import FileWriter
+from smdebug.exceptions import InvalidCollectionConfiguration
 
 try:
     from smexperiments.metrics import SageMakerFileMetricsWriter
@@ -311,10 +312,17 @@ def _get_collections_with_tensor(self, tensor_name) -> Set["Collection"]:
             self.tensor_to_collections[tensor_name] = matched_colls
         return self.tensor_to_collections[tensor_name]
 
+    @abstractmethod
+    def _get_default_collections(self):
+        pass
+
     def _prepare_collections(self):
         """Populate collections_to_save and ensure every collection has
         a save_config and reduction_config."""
         for c_name, c in self.collection_manager.get_collections().items():
+            if c_name not in self._get_default_collections():
+                if bool(c.include_regex) is False and bool(c.tensor_names) is False:
+                    raise InvalidCollectionConfiguration(c_name)
             if c in self._collections_to_save:
                 continue
             elif self._should_collection_be_saved(CollectionKeys.ALL):

smdebug/exceptions.py

@@ -2,6 +2,15 @@
 from smdebug.core.modes import ModeKeys as modes
 
 
+class InvalidCollectionConfiguration(Exception):
+    def __init__(self, c_name):
+        self.c_name = c_name
+
+    def __str__(self):
+        return f"Collection {self.c_name} has not been configured. \
+        Please fill in tensor_name or include_regex"
+
+
 class StepNotYetAvailable(Exception):
     def __init__(self, step, mode):
         self.step = step

smdebug/mxnet/collection.py

@@ -1,4 +1,5 @@
 # First Party
+from smdebug.core.collection import DEFAULT_MXNET_COLLECTIONS
 from smdebug.core.collection import Collection as BaseCollection
 from smdebug.core.collection import CollectionKeys
 from smdebug.core.collection_manager import CollectionManager as BaseCollectionManager
@@ -21,6 +22,8 @@ def __init__(self, create_default=True):
             self._register_default_collections()
 
     def _register_default_collections(self):
+        for c in DEFAULT_MXNET_COLLECTIONS:
+            self.create_collection(c)
         self.get(CollectionKeys.WEIGHTS).include("^(?!gradient).*weight")
         self.get(CollectionKeys.BIASES).include("^(?!gradient).*bias")
         self.get(CollectionKeys.GRADIENTS).include("^gradient")

smdebug/mxnet/hook.py

@@ -2,7 +2,7 @@
 import mxnet as mx
 
 # First Party
-from smdebug.core.collection import CollectionKeys
+from smdebug.core.collection import DEFAULT_MXNET_COLLECTIONS, CollectionKeys
 from smdebug.core.hook import CallbackHook
 from smdebug.core.json_config import DEFAULT_WORKER_NAME
 from smdebug.mxnet.collection import CollectionManager
@@ -113,6 +113,9 @@ def _export_model(self):
                     f"due to the mxnet exception: {e}"
                 )
 
+    def _get_default_collections(self):
+        return DEFAULT_MXNET_COLLECTIONS
+
     # This hook is invoked by trainer prior to running the forward pass.
     def forward_pre_hook(self, block, inputs):
         if self.writer is not None:

smdebug/pytorch/collection.py

@@ -1,4 +1,5 @@
 # First Party
+from smdebug.core.collection import DEFAULT_PYTORCH_COLLECTIONS
 from smdebug.core.collection import Collection as BaseCollection
 from smdebug.core.collection import CollectionKeys
 from smdebug.core.collection_manager import CollectionManager as BaseCollectionManager
@@ -36,6 +37,8 @@ def __init__(self, create_default=True):
             self._register_default_collections()
 
     def _register_default_collections(self):
+        for c in DEFAULT_PYTORCH_COLLECTIONS:
+            self.create_collection(c)
         self.get(CollectionKeys.WEIGHTS).include("^(?!gradient).*weight")
         self.get(CollectionKeys.BIASES).include("^(?!gradient).*bias")
         self.get(CollectionKeys.GRADIENTS).include("^gradient")

smdebug/pytorch/hook.py

@@ -5,7 +5,7 @@
 import torch.distributed as dist
 
 # First Party
-from smdebug.core.collection import CollectionKeys
+from smdebug.core.collection import DEFAULT_PYTORCH_COLLECTIONS, CollectionKeys
 from smdebug.core.hook import CallbackHook
 from smdebug.core.json_config import DEFAULT_WORKER_NAME
 from smdebug.pytorch.collection import CollectionManager
@@ -103,15 +103,18 @@ def _log_params(self, module):
     def _export_model(self):
         pass
 
+    def _get_default_collections(self):
+        return DEFAULT_PYTORCH_COLLECTIONS
+
     def _prepare_collections(self):
-        super()._prepare_collections()
         for coll in self.collection_manager.collections.values():
             for m, (include_inputs, include_outputs) in coll.modules.items():
                 module_name = self.module_maps[m]
                 if include_inputs:
                     coll.include(module_name + "_input_")
                 if include_outputs:
                     coll.include(module_name + "_output_")
+        super()._prepare_collections()
 
     # This hook is invoked by trainer prior to running the forward pass.
     def forward_pre_hook(self, module, inputs):

smdebug/tensorflow/base_hook.py

@@ -7,6 +7,7 @@
 from tensorflow.python.distribute.distribute_lib import _DefaultDistributionStrategy
 
 # First Party
+from smdebug.core.collection import DEFAULT_TF_COLLECTIONS
 from smdebug.core.config_constants import DEFAULT_WORKER_NAME
 from smdebug.core.hook import BaseHook
 from smdebug.core.modes import ModeKeys
@@ -179,6 +180,9 @@ def _get_worker_name(self) -> str:
         elif self.distribution_strategy == TFDistributionStrategy.UNSUPPORTED:
             raise NotImplementedError
 
+    def _get_default_collections(self):
+        return DEFAULT_TF_COLLECTIONS
+
     def export_collections(self):
         assert self._prepared_tensors[self.mode]
 

smdebug/tensorflow/collection.py

@@ -9,6 +9,7 @@
 from tensorflow.python.distribute import values
 
 # First Party
+from smdebug.core.collection import DEFAULT_TF_COLLECTIONS
 from smdebug.core.collection import Collection as BaseCollection
 from smdebug.core.collection import CollectionKeys
 from smdebug.core.collection_manager import CollectionManager as BaseCollectionManager
@@ -136,18 +137,7 @@ class CollectionManager(BaseCollectionManager):
     def __init__(self, collections=None, create_default=True):
         super().__init__(collections=collections)
         if create_default:
-            for n in [
-                CollectionKeys.DEFAULT,
-                CollectionKeys.WEIGHTS,
-                CollectionKeys.BIASES,
-                CollectionKeys.GRADIENTS,
-                CollectionKeys.LOSSES,
-                CollectionKeys.METRICS,
-                CollectionKeys.INPUTS,
-                CollectionKeys.OUTPUTS,
-                CollectionKeys.ALL,
-                CollectionKeys.SM_METRICS,
-            ]:
+            for n in DEFAULT_TF_COLLECTIONS:
                 self.create_collection(n)
             self.get(CollectionKeys.BIASES).include("bias")
 

smdebug/xgboost/collection.py

@@ -1,5 +1,5 @@
 # First Party
-from smdebug.core.collection import CollectionKeys
+from smdebug.core.collection import DEFAULT_XGBOOST_COLLECTIONS, CollectionKeys
 from smdebug.core.collection_manager import CollectionManager as BaseCollectionManager
 
 
@@ -10,6 +10,8 @@ def __init__(self, create_default=True):
             self._register_default_collections()
 
     def _register_default_collections(self):
+        for c in DEFAULT_XGBOOST_COLLECTIONS:
+            self.create_collection(c)
         self.get(CollectionKeys.HYPERPARAMETERS).include("^hyperparameters/.*$")
         self.get(CollectionKeys.METRICS).include("^[a-zA-z]+-[a-zA-z0-9]+$")
         self.get(CollectionKeys.PREDICTIONS).include("^predictions$")

smdebug/xgboost/hook.py

@@ -9,7 +9,7 @@
 from xgboost.core import CallbackEnv
 
 # First Party
-from smdebug.core.collection import CollectionKeys
+from smdebug.core.collection import DEFAULT_XGBOOST_COLLECTIONS, CollectionKeys
 from smdebug.core.hook import CallbackHook
 from smdebug.core.json_config import create_hook_from_json_config
 from smdebug.core.save_config import SaveConfig
@@ -144,6 +144,12 @@ def create_from_json_file(cls, json_file_path=None):
     def hook_from_config(cls, json_config_path=None):
         return cls.create_from_json_file(json_file_path=json_config_path)
 
+    def _get_default_collections(self):
+        return DEFAULT_XGBOOST_COLLECTIONS
+
+    def _prepare_collections(self):
+        super()._prepare_collections()
+
     def _is_last_step(self, env: CallbackEnv) -> bool:
         # env.iteration: current boosting round.
         # env.end_iteration: round # when training will end. this is always num_round + 1.  # noqa: E501

tests/core/test_collections.py

@@ -10,6 +10,7 @@
 from smdebug.core.reduction_config import ReductionConfig
 from smdebug.core.save_config import SaveConfig, SaveConfigMode
 from smdebug.core.utils import get_path_to_collections
+from smdebug.exceptions import InvalidCollectionConfiguration
 from smdebug.mxnet.hook import Hook
 
 
@@ -87,6 +88,7 @@ def test_collection_defaults_to_hook_config():
   """
     cm = CollectionManager()
     cm.create_collection("foo")
+    cm.get("foo").include_regex = "*"
     cm.get("foo").save_config = {ModeKeys.EVAL: SaveConfigMode(save_interval=20)}
 
     hook = Hook(
@@ -101,3 +103,28 @@ def test_collection_defaults_to_hook_config():
     hook._prepare_collections()
     assert cm.get("foo").save_config.mode_save_configs[ModeKeys.TRAIN].save_interval == 10
     assert cm.get("foo").reduction_config.save_raw_tensor is True
+
+
+def test_invalid_collection_config_exception():
+    cm = CollectionManager()
+    cm.create_collection("foo")
+
+    hook = Hook(
+        out_dir="/tmp/test_collections/" + str(datetime.datetime.now()),
+        save_config={ModeKeys.TRAIN: SaveConfigMode(save_interval=10)},
+        include_collections=["foo"],
+        reduction_config=ReductionConfig(save_raw_tensor=True),
+    )
+    hook.collection_manager = cm
+    try:
+        hook._prepare_collections()
+    except InvalidCollectionConfiguration:
+        pass
+    else:
+        assert False, "Invalid Collection Name did not raise error"
+
+    cm.get("foo").include_regex = "*"
+    try:
+        hook._prepare_collections()
+    except InvalidCollectionConfiguration:
+        assert False, "Valid Collection Name raised an error"