IndexReader exception handling incase the index files are broken (aws#264)

Author: NihalHarish

tornasole/core/index_reader.py

@@ -7,6 +7,7 @@
 from tornasole.core.logger import get_logger
 from tornasole.core.tfrecord.tensor_reader import TensorReader
 from tornasole.core.modes import ModeKeys
+from tornasole.exceptions import IndexReaderException
 
 
 """
@@ -121,9 +122,17 @@ def load_tensor_data_from_index_files(path, start_after_key=None, range_steps=No
             tensor_data = IndexReader._update_tensors_from_json(tensor_data, step, response, path, worker)
         return tensor_data, last_index_token
 
+    @staticmethod
+    def _validate(index_dict):
+        if 'meta' not in index_dict or len(index_dict['meta']) == 0:
+            raise IndexReaderException('meta section is not present')
+        if 'tensor_payload' not in index_dict:
+            raise IndexReaderException('tensor_payload section is not present')
+
     @staticmethod
     def _update_tensors_from_json(index_tensors_dict, step, response, path, worker):
         index_dict = json.loads(response)
+        IndexReader._validate(index_dict)
         index_meta = index_dict['meta']
         mode = index_meta['mode']
         mode = ModeKeys[mode.strip()]

tornasole/exceptions.py

@@ -11,6 +11,14 @@ def __str__(self):
                                                          self.mode.name)
 
 
+class IndexReaderException(Exception):
+  def __init__(self, message):
+    self.message = message
+
+  def __str__(self):
+    return self.message
+
+
 class StepUnavailable(Exception):
   def __init__(self, step, mode):
     self.step = step