Add support to save mode and mode step in the summary (#38)

* core mode changes

* update with changes to writer and reader

* update tensor reader

* simplified remove tensor as tf is using its own method now

* implemented except when_nan

* tensor can match multiple collections

* save manager working, and tensor reduction change

* delete save manager, move to new PR

* Revert "delete save manager, move to new PR"

This reverts commit 475122de1c5613c800fdbfc9abdb1c2d24a64add.

* update to support mxnet as well

* Changes for mxnet

* remove default

* reorg and fix tests

* Added string constant

* add comment and an allowed_modes list

Author: rahul003

tests/test_collections.py renamed to tests/core/test_collections.py

@@ -15,6 +15,7 @@ def test_load_empty():
 
 def test_manager_export_load():
   cm = CollectionManager()
+  cm.create_collection('default')
   cm.get('default').include('loss')
   cm.add(Collection('trial1'))
   cm.add('trial2')
@@ -25,6 +26,7 @@ def test_manager_export_load():
 
 def test_manager():
   cm = CollectionManager()
+  cm.create_collection('default')
   cm.get('default').include('loss')
   cm.add(Collection('trial1'))
   cm.add('trial2')

tests/test_handler.py renamed to tests/core/test_handler.py

@@ -0,0 +1,32 @@
+from tornasole_core.writer import FileWriter
+from tornasole_core.reader import FileReader
+import numpy as np
+from tornasole_core.modes import ModeKeys
+from datetime import datetime
+import glob
+import shutil
+
+def test_mode_writing():
+  run_id = 'trial_' + datetime.now().strftime('%Y%m%d-%H%M%S%f')
+  for s in range(0, 10):
+
+    fw = FileWriter(logdir='ts_outputs', trial=run_id, step=s)
+    if s % 2 == 0:
+      fw.write_tensor(tdata=np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32),
+                      tname='arr', mode=ModeKeys.TRAIN, mode_step=s//2)
+    else:
+      fw.write_tensor(tdata=np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32),
+                      tname='arr', mode=ModeKeys.EVAL, mode_step=s // 2)
+  fw.close()
+  files = glob.glob('ts_outputs/' + run_id + '/**/*.tfevents',
+                    recursive=True)
+  for f in files:
+    fr = FileReader(fname=f)
+    for tu in fr.read_tensors():
+      tensor_name, step, tensor_data, mode, mode_step = tu
+      if step % 2 == 0:
+        assert mode == ModeKeys.TRAIN
+      else:
+        assert mode == ModeKeys.EVAL
+      assert mode_step == step // 2
+  shutil.rmtree('ts_outputs/' + run_id)

tests/test_index.py renamed to tests/core/test_index.py

@@ -34,8 +34,8 @@ def test_s3():
     my_session = boto3.session.Session()
     my_region = my_session.region_name
     my_account = boto3.client('sts').get_caller_identity().get('Account')
-    bucket_name = 'sagemaker-{}-{}'.format(my_region,my_account)
-    key_name = 'tornasole/{}'.format(str(uuid.uuid4()))
+    bucket_name = 'tornasole-testing'
+    key_name = 'core-tests/tornasole/{}'.format(str(uuid.uuid4()))
     #sagemaker-us-east-1-722321484884
     location = 's3://{}/{}'.format(bucket_name,key_name)
     print("Saving to Location")

tests/core/test_modes.py

@@ -87,16 +87,10 @@ def add_tensor(self, t):
       self.tensors.append(t)
 
   def remove_tensor(self, t):
-    # have to compare names because tensors can have variables, \
-    # we don't want to end up comparing tensors and variables
     if t.name in self.tensor_names:
-      found_index = None
-      for i, lt in enumerate(self.tensors):
-        if lt.name == t.name:
-          found_index = i
-      assert found_index is not None
-      self.tensors.pop(found_index)
       self.tensor_names.remove(t.name)
+    if t in self.tensors:
+      self.tensors.remove(t)
 
   def add_reduction_tensor(self, s):
     self.reduction_tensor_names.append(s.name)

tests/test_numpy.py renamed to tests/core/test_numpy.py

@@ -8,21 +8,19 @@ class CollectionManager:
   It contains a default collection into which tensors are inserted
   without specifying collection name
   """
-  def __init__(self, create_default=True):
+  def __init__(self):
     self.collections = {}
-    if create_default:
-      self.collections['default'] = self.get_new_collection('default')
 
-  def get_new_collection(self, name):
-    return Collection(name)
+  def create_collection(self, name):
+    self.collections[name] = Collection(name)
 
   def get_collections(self):
     return self.collections
 
   def add(self, arg):
     if isinstance(arg, str):
       if arg not in self.collections:
-        self.collections[arg] = self.get_new_collection(arg)
+        self.create_collection(arg)
     elif isinstance(arg, Collection):
       if arg.name not in self.collections:
         self.collections[arg.name] = arg
@@ -43,7 +41,7 @@ def export(self, filename):
 
   @staticmethod
   def load(filename):
-    cm = CollectionManager(create_default=False)
+    cm = CollectionManager()
     with open(filename, 'r') as f:
       line = f.readline()
       while line:
@@ -54,7 +52,7 @@ def load(filename):
 
   @staticmethod
   def load_from_string(s):
-    cm = CollectionManager(create_default=False)
+    cm = CollectionManager()
     lines = s.split('\n')
     for line in lines:
       c = Collection.load(line.rstrip())

tests/test_reduction_config.py renamed to tests/core/test_reduction_config.py

@@ -0,0 +1,12 @@
+from enum import Enum
+
+# Note that Keras has similar concept of ModeKeys
+class ModeKeys(Enum):
+    TRAIN = 1 #training/fitting mode
+    EVAL = 2  # testing/evaluation mode
+    PREDICT = 3 # prediction/inference mode
+    GLOBAL = 4
+
+ALLOWED_MODES = [ModeKeys.TRAIN, ModeKeys.EVAL, ModeKeys.PREDICT]
+MODE_STEP_PLUGIN_NAME = "mode_step"
+MODE_PLUGIN_NAME = "mode"

tests/test_save_config.py renamed to tests/core/test_save_config.py

@@ -1,4 +1,36 @@
 SAVE_CONFIG_VERSION_NUM = 'v0'
+from .modes import ModeKeys as modes
+
+
+class SaveConfigModes:
+  def __init__(self, mode_save_configs=None):
+    if mode_save_configs is None:
+      mode_save_configs = {}
+    self.mode_save_configs = mode_save_configs
+
+  def add_for_all_modes(self, save_config):
+    for mode in modes:
+      self.mode_save_configs[mode] = save_config
+
+  def add(self, mode, save_config):
+    self.mode_save_configs[mode] = save_config
+
+  def should_save_step(self, mode, step_num):
+    return self.mode_save_configs[mode].should_save_step(step_num)
+
+  def add_when_nan_tensor(self, tensor):
+    for mode in modes:
+      self.mode_save_configs[mode].when_nan_tensors.append(tensor)
+
+  def get_save_config(self, mode):
+    return self.mode_save_configs[mode]
+
+  @staticmethod
+  def create_simple_save_mode(save_config):
+    sm = SaveConfigModes()
+    sm.add_for_all_modes(save_config)
+    return sm
+
 
 class SaveConfig:
   """
@@ -31,6 +63,9 @@ def __init__(self, save_interval=100, skip_num_steps=0, save_steps=None, when_na
     self.skip_num_steps = skip_num_steps
     self.when_nan = when_nan if when_nan is not None else []
 
+    # will be populated by hook
+    self.when_nan_tensors = []
+
   def export(self):
     separator = '%'
     list_separator = ','

tests/test_utils.py renamed to tests/core/test_utils.py

@@ -0,0 +1,138 @@
+from .save_config import SaveConfig, SaveConfigModes
+from .utils import match_inc
+
+class SaveManager:
+  def __init__(self, collection_manager, include_collections_names,
+               default_reduction_config,
+               default_save_config):
+    self.configs_for_collections = {}
+    if isinstance(default_save_config, SaveConfig):
+      sm = SaveConfigModes.create_simple_save_mode(default_save_config)
+      self.default_save_modes = sm
+    elif isinstance(default_save_config, SaveConfigModes):
+      self.default_save_modes = default_save_config
+    elif isinstance(default_save_config, dict):
+      self.default_save_modes = SaveConfigModes(default_save_config)
+    else:
+      raise TypeError('save_config can only be a SaveConfig instance, or '
+                      'a dictionary mapping from mode '
+                      'to SaveConfig instance.')
+    self.default_reduction_config = default_reduction_config
+    self.collection_manager = collection_manager
+    self.include_collections_names = include_collections_names
+    self.save_collections = []
+    self.save_states_cache = {}
+    # todo clear cache for old steps
+    self.tensor_to_collection = {}
+    self.when_nan_tensors = {}
+
+  def prepare(self):
+    # below is to control the order
+    # in which these collections appear in save_collections
+    for cname in ['weights', 'gradients', 'bias', 'optimizer_variables']:
+      if self._should_collection_be_saved(cname):
+        self.save_collections.append(self.collection_manager.get(cname))
+
+    # adding other collections to save_collections
+    for c_name, c in self.collection_manager.get_collections().items():
+      if self._should_collection_be_saved(c_name) \
+        and c not in self.save_collections:
+        self.save_collections.append(c)
+
+    for c_name, c in self.collection_manager.get_collections().items():
+      if c.save_config is not None:
+        if isinstance(c.save_config, dict):
+          self.configs_for_collections[c_name] = c.save_config
+        elif isinstance(c.save_config, SaveConfig):
+          sm = SaveConfigModes.create_simple_save_mode(c.save_config)
+          self.configs_for_collections[c_name] = sm
+        else:
+          raise TypeError('collection {} has save config of wrong type {}'
+                          .format(c_name, type(c.save_config)))
+      else:
+        self.configs_for_collections[c_name] = self.default_save_modes
+
+      if c.reduction_config is None and self.default_reduction_config is not None:
+        c.reduction_config = self.default_reduction_config
+
+  def _should_collection_be_saved(self, coll_name):
+    return coll_name in self.include_collections_names
+
+  def get_all_collections_to_save(self):
+    return self.save_collections
+
+  def collections_to_save(self, mode, step):
+    if (mode, step) not in self.save_states_cache:
+      collection_save_state = {}
+      for coll in self.save_collections:
+        sm = self.configs_for_collections[coll.name]
+        rv = sm.should_save_step(mode, step)
+        if any(rv.values()):
+          collection_save_state[coll.name] = rv
+      self.save_states_cache[(mode, step)] = collection_save_state
+    return self.save_states_cache[(mode, step)]
+
+  def get_save_config(self, collection, mode):
+    return self.configs_for_collections[collection.name].get_save_config(mode)
+
+  def from_collections(self, tensor_name):
+    # for tf this will be prepopulated because of prepare_tensors
+    if not tensor_name in self.tensor_to_collection:
+      # for mxnet it is computed and then cached
+      matched_colls = []
+      for coll in self.save_collections:
+        if tensor_name in coll.tensor_names:
+          matched_colls.append(coll)
+        elif match_inc(tensor_name, coll.get_include_regex()):
+          matched_colls.append(coll)
+      self.tensor_to_collection[tensor_name] = matched_colls
+    return self.tensor_to_collection[tensor_name]
+
+  def should_save_tensor(self, tensorname, mode, step):
+    # returns dictionary with two keys:
+    # if value for step is true in the dict, then we are saving this tensor
+    # because we have hit the step to save this
+    # if value for when_nan is true, we are considering saving this tensor
+    # because this tensor might be saved if some other tensor is nan
+    colls = self.from_collections(tensorname)
+    final_ss = {'step': False, 'when_nan': False}
+    ss_colls = self.collections_to_save(mode, step)
+    for c in colls:
+      if c.name in ss_colls:
+        ss = ss_colls[c.name]
+        final_ss['step'] = final_ss['step'] or ss['step']
+        final_ss['when_nan'] = final_ss['when_nan'] or ss['when_nan']
+    return final_ss
+
+  # below are used only by TF
+  def prepare_tensors(self):
+    for c_name, c in self.collection_manager.get_collections().items():
+      if c_name == 'when_nan':
+        continue
+      if c not in self.save_collections:
+        continue
+      for t in c.tensors + c.reduction_tensors:
+        self._add_tensor_to_collection(t, c)
+
+  def _add_tensor_to_collection(self, t, c):
+    if t.name not in self.tensor_to_collection:
+      self.tensor_to_collection[t.name] = [c]
+    else:
+      self.tensor_to_collection[t.name].append(c)
+
+  def add_when_nan_tensor(self, collection, tensor):
+    self.configs_for_collections[collection.name].add_when_nan_tensor(tensor)
+    if tensor.name not in self.when_nan_tensors:
+      self.when_nan_tensors[tensor.name] = []
+    self.when_nan_tensors[tensor.name].append(collection)
+    self._add_tensor_to_collection(tensor, collection)
+
+    if 'when_nan' not in self.collection_manager.collections:
+      self.collection_manager.create_collection('when_nan')
+    self.collection_manager.get('when_nan').add_tensor(tensor)
+
+  def is_when_nan_tensor(self, tensor_name):
+    return tensor_name in self.when_nan_tensors
+
+  def when_nan_collections(self, tensor_name):
+    return self.when_nan_tensors[tensor_name]

tornasole_core/collection.py

@@ -28,6 +28,8 @@
 from .summary_pb2 import Summary, SummaryMetadata
 
 from tornasole_core.tfrecord.record_reader import RecordReader
+from tornasole_core.modes import ModeKeys, MODE_STEP_PLUGIN_NAME, MODE_PLUGIN_NAME
+
 
 #todo: remove this logger perhaps
 logging.basicConfig()
@@ -133,7 +135,7 @@ def __exit__(self,exc_type, exc_value, traceback):
         self._ev_reader.__exit__(exc_type, exc_value, traceback)
 
     def read_tensors(self, read_data=False, check=False):
-        for (step,summ) in self.read_summaries(check=check):
+        for step, summ in self.read_summaries(check=check):
             for v in summ.value:
                 assert v.WhichOneof('value') == 'tensor'
                 tensor_name = v.tag
@@ -142,7 +144,17 @@ def read_tensors(self, read_data=False, check=False):
                     tensor_data = get_tensor_data(v.tensor)
                 else:
                     tensor_data = None
-                yield (tensor_name, step, tensor_data)
+
+                # default values
+                # todo: validate the logic extensively
+                mode_step = step
+                mode = ModeKeys.GLOBAL
+                for metadata in v.metadata.plugin_data:
+                    if metadata.plugin_name == MODE_STEP_PLUGIN_NAME:
+                        mode_step = int(metadata.content)
+                    if metadata.plugin_name == MODE_PLUGIN_NAME:
+                        mode = ModeKeys(int(metadata.content))
+                yield (tensor_name, step, tensor_data, mode, mode_step)
 
     def read_summaries(self, check=True):
         for ev in self.read_events(check=check):