Collections only handle tensor_names now, Cleanup (aws#60)

* make core use only names

* update branches

* try to fix boto errors

* accept yes prompt

* trigger ci

* make set

* trigger CI

* fix test

* trigger ci

* fix branch

* trigger ci

Author: rahul003

config/buildspec.yml

@@ -12,6 +12,7 @@ phases:
         - cd $CODEBUILD_SRC_DIR  && chmod +x config/protoc_downloader.sh && ./config/protoc_downloader.sh
         - pip install pytest
         - pip install wheel
+        - pip uninstall -y boto3 && pip uninstall -y aiobotocore && pip uninstall -y botocore
 
   pre_build:
     commands:

config/configure_branch_for_test.txt

@@ -1,5 +1,5 @@
-tornasole_core:default
-tornasole_tf:default
-tornasole_mxnet:default
-tornasole_rules:default
+tornasole_core:redn_fix
+tornasole_tf:redn_fix
+tornasole_mxnet:redn_fix
+tornasole_rules:fix_mode_available
 tornasole_pytorch:default

tests/core/test_collections.py

@@ -28,10 +28,12 @@ def test_manager():
   cm = CollectionManager()
   cm.create_collection('default')
   cm.get('default').include('loss')
+  cm.get('default').add_tensor_name('assaas')
   cm.add(Collection('trial1'))
   cm.add('trial2')
   cm.get('trial2').include('total_loss')
   assert len(cm.collections) == 3
   assert cm.get('default') == cm.collections['default']
   assert 'loss' in cm.get('default').include_regex
+  assert len(cm.get('default').get_tensor_names()) > 0
   assert 'total_loss' in cm.collections['trial2'].include_regex

tornasole_core/collection.py

@@ -42,19 +42,18 @@ def __init__(self, name, include_regex=None,
     self.reduction_config = reduction_config
     self.save_config = save_config
 
-    self.tensor_names = []
-    self.reduction_tensor_names = []
-
-    # these two are internal fields only, used by TF
-    self.tensors = []
-    self.reduction_tensors = []
+    self.tensor_names = set()
+    self.reduction_tensor_names = set()
 
   def get_include_regex(self):
     return self.include_regex
 
   def get_tensor_names(self):
     return self.tensor_names
 
+  def get_reduction_tensor_names(self):
+    return self.reduction_tensor_names
+
   def include(self, t):
     if isinstance(t, list):
       for i in t:
@@ -80,20 +79,17 @@ def set_save_config(self, save_cfg):
       raise TypeError('Can only take an instance of SaveConfig')
     self.save_config = save_cfg
 
-  def add_tensor(self, t):
-    if t.name not in self.tensor_names:
-      self.tensor_names.append(t.name)
-      self.tensors.append(t)
+  def add_tensor_name(self, tname):
+    if tname not in self.tensor_names:
+      self.tensor_names.add(tname)
 
-  def remove_tensor(self, t):
-    if t.name in self.tensor_names:
-      self.tensor_names.remove(t.name)
-    if t in self.tensors:
-      self.tensors.remove(t)
+  def remove_tensor_name(self, tname):
+    if tname in self.tensor_names:
+      self.tensor_names.remove(tname)
 
-  def add_reduction_tensor(self, s):
-    self.reduction_tensor_names.append(s.name)
-    self.reduction_tensors.append(s)
+  def add_reduction_tensor_name(self, sname):
+    if sname not in self.reduction_tensor_names:
+      self.reduction_tensor_names.add(sname)
 
   def export(self):
     # v0 export
@@ -135,8 +131,8 @@ def load(s):
       list_separator = ','
       name = parts[1]
       include = [x for x in parts[2].split(list_separator) if x]
-      tensor_names = [x for x in parts[3].split(list_separator) if x]
-      reduction_tensor_names = [x for x in parts[4].split(list_separator) if x]
+      tensor_names = set([x for x in parts[3].split(list_separator) if x])
+      reduction_tensor_names = set([x for x in parts[4].split(list_separator) if x])
       reduction_config = ReductionConfig.load(parts[5])
       save_config = SaveConfig.load(parts[6])
       c = Collection(name, include_regex=include,

tornasole_core/save_manager.py

@@ -24,7 +24,7 @@ def __init__(self, collection_manager, include_collections_names,
     self.save_states_cache = {}
     # todo clear cache for old steps
     self.tensor_to_collection = {}
-    self.when_nan_tensors = {}
+    self.prepared = False
 
   def prepare(self):
     for c_name, c in self.collection_manager.get_collections().items():
@@ -47,14 +47,24 @@ def prepare(self):
 
       if c.reduction_config is None and self.default_reduction_config is not None:
         c.reduction_config = self.default_reduction_config
+    self.prepared = True
 
   def _should_collection_be_saved(self, coll_name):
     return coll_name in self.include_collections_names
 
+  def _raise_error(self):
+    raise ValueError('Save Manager is not ready. '
+                     'Please call prepare() method '
+                     'before calling this method.')
+
   def get_all_collections_to_save(self):
+    if not self.prepared:
+      self._raise_error()
     return self.save_collections
 
   def collections_to_save(self, mode, step):
+    if not self.prepared:
+      self._raise_error()
     if (mode, step) not in self.save_states_cache:
       collection_save_state = {}
       for coll in self.save_collections:
@@ -66,21 +76,31 @@ def collections_to_save(self, mode, step):
     return self.save_states_cache[(mode, step)]
 
   def get_save_config(self, collection, mode):
+    if not self.prepared:
+      self._raise_error()
     return self.configs_for_collections[collection.name].get_save_config(mode)
 
   def get_reduction_config(self, collection):
+    if not self.prepared:
+      self._raise_error()
     return collection.get_reduction_config()
 
   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:
+      for coll in self.get_all_collections_to_save():
+        if tensor_name in coll.tensor_names or \
+         tensor_name in coll.reduction_tensor_names:
+          # if being matched as reduction,
+          # it must be in reduction_tensor_name, not with regex
           matched_colls.append(coll)
         elif match_inc(tensor_name, coll.get_include_regex()):
-          coll.tensor_names.append(tensor_name)
+          if self.get_reduction_config(coll):
+            coll.add_reduction_tensor_name(tensor_name)
+          else:
+            coll.add_tensor_name(tensor_name)
           matched_colls.append(coll)
       self.tensor_to_collection[tensor_name] = matched_colls
     return self.tensor_to_collection[tensor_name]
@@ -100,36 +120,3 @@ def should_save_tensor(self, tensorname, mode, step):
         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]