mode writer support (aws#144)

Author: NihalHarish

tests/core/test_index.py

@@ -32,9 +32,9 @@ def test_index():
         line_count = 0
         i = 0
         for row in csv_reader:
-            count = int(row[2])
+            count = int(row[-2])
             fo.seek(count, 0)
-            end = int(row[3])
+            end = int(row[-1])
             line = fo.read(end)
             zoo = open("test.txt", "wb")
             zoo.write(line)

tornasole/core/indexutils.py

@@ -1,17 +1,30 @@
 class TensorLocation:
-    def __init__(self, tname, event_file_name, start_idx, length):
+    def __init__(self, tname, mode, mode_step, event_file_name, start_idx, length):
         self.tensorname = tname
+        self.mode = mode
+        self.mode_step = mode_step
         self.event_file_name = event_file_name
         self.start_idx = start_idx
         self.length = length
 
     def serialize(self):
-        return format(f'{self.tensorname}, {self.event_file_name}, {self.start_idx},{self.length}')
+        return format(
+            f'{self.tensorname},'
+            f'{self.mode}, '
+            f'{self.mode_step}, '
+            f'{self.event_file_name}, '
+            f'{self.start_idx}, '
+            f'{self.length}')
 
     @staticmethod
     def deserialize(manifest_line_str, manifest_key_name):
         arr = manifest_line_str.split(",")
-        return TensorLocation(arr[0], arr[1], arr[2], arr[3])
+        return TensorLocation(arr[0],
+                              arr[1],
+                              arr[2],
+                              arr[3],
+                              arr[4],
+                              arr[5])
 
 
 class IndexUtil:

tornasole/core/tfevent/event_file_writer.py

@@ -226,7 +226,7 @@ def write_tensor(self, tdata, tname, write_index,
         s = Summary(value=[Summary.Value(tag=tag, metadata=smd,
                                          tensor=tensor_proto)])
         if write_index:
-            self.write_summary_with_index(s, self.step, tname)
+            self.write_summary_with_index(s, self.step, tname, mode, mode_step)
         else:
             self.write_summary(s, self.step)
 
@@ -236,11 +236,11 @@ def write_summary(self, summary, step):
         event.step = step
         self.write_event(event)
 
-    def write_summary_with_index(self, summary, step, tname):
+    def write_summary_with_index(self, summary, step, tname, mode, mode_step):
         event = Event(summary=summary)
         event.wall_time = time.time()
         event.step = step
-        return self.write_event(IndexArgs(event, tname))
+        return self.write_event(IndexArgs(event, tname, mode, mode_step))
 
     def write_event(self, event):
         """Adds an event to the event file."""
@@ -302,10 +302,12 @@ def run(self):
                 if isinstance(event_in_queue, IndexArgs):
                     tname = event_in_queue.tensorname
                     eventfile = self._ev_writer.name()
+                    mode = event_in_queue.get_mode()
+                    mode_step = event_in_queue.get_mode_step()
                     s3, _, _ = is_s3(eventfile)
                     if not s3:
                         eventfile = os.path.abspath(self._ev_writer.name())
-                    tensorlocation = TensorLocation(tname, eventfile, positions[0], positions[1])
+                    tensorlocation = TensorLocation(tname, mode, mode_step, eventfile, positions[0], positions[1])
                     self._ev_writer.indexwriter.add_index(tensorlocation)
                 # Flush the event writer every so often.
                 now = time.time()

tornasole/core/tfevent/index_file_writer.py

@@ -2,6 +2,7 @@
 from tornasole.core.access_layer.s3 import TSAccessS3
 from tornasole.core.utils import is_s3
 
+
 class IndexWriter(object):
     def __init__(self, file_path):
         self.file_path = file_path
@@ -38,13 +39,22 @@ def close(self):
             self.writer.close()
             self.writer = None
 
+
 class IndexArgs(object):
-    def __init__(self, event, tensorname):
+    def __init__(self, event, tensorname, mode, mode_step):
         self.event = event
         self.tensorname = tensorname
+        self.mode = mode
+        self.mode_step = mode_step
 
     def get_event(self):
         return self.event
 
     def get_tensorname(self):
         return self.tensorname
+
+    def get_mode(self):
+        return str(self.mode).split('.')[-1]
+
+    def get_mode_step(self):
+        return self.mode_step