DT Trial APIs (aws#239)

* test

* added missing import

* get worker name from func

* random seed:

* str error

* error handling

* hvd-size

* worker name to collection

* removed unecessary comment

* try catch for modulenotfound

* simplified mkdir error handling

* handle hvd.size ValueError

* modify try except block

* moving worker name to utils

* worker to collections.ts

* list collection files

* checkpoint_1

* Mirrored Variable

* mirrored var

* revert accidental change

* print statement cleanup

* save json file

* update num workers

* s3 mode trial api

* tests starting

* fixed test trial modes

* bug fix

* removed randint generator

* pytorch test fix

* fixed bug in test

* merge conflicts and tf tests

* merge alpha to fix breaking tests

* code cleanup

* fixed get correct worker bug

* fixed None not 0 bug

* made worker name for hvd more verbose

* fixed mode step bug

* changed worker name parsing logic

* missing line

* removed hard coded worker_num values

* paramterized collections name

* changes for PR comments

* PR changes

* step variable rename and refactor

* reverting 4eeb4d4

* refactored changes

* read test

* refactor changes

* handle import error for pytorch

* Handle Import Error

* removed worker from hook, set default value to num_workers, addressed PR comments

* doc string update

* replace StepIsUnavailable with ValueError and fix return with raise

Author: NihalHarish

tests/analysis/rules/test_confusion.py

@@ -1,5 +1,6 @@
 from tornasole.core.writer import FileWriter
-from tornasole.core.collection_manager import CollectionManager, COLLECTIONS_FILE_NAME
+from tornasole.core.collection_manager import CollectionManager
+from tornasole.core.config_constants import TORASOLE_DEFAULT_COLLECTIONS_FILE_NAME
 
 from tornasole.rules.generic import Confusion
 from tornasole.trials import create_trial
@@ -10,7 +11,8 @@
 from tornasole.exceptions import *
 from tornasole.rules.rule_invoker import invoke_rule
 
-def gen_y_and_y_hat( path, trial, step, y, y_name, y_hat, y_hat_name, colls = {} ):
+
+def gen_y_and_y_hat( path, trial, step, y, y_name, y_hat, y_hat_name, colls={}):
     trial_dir=os.path.join(path, trial)
     with FileWriter(trial_dir=trial_dir,
                     step=step, worker='algo-1') as fw:
@@ -20,7 +22,7 @@ def gen_y_and_y_hat( path, trial, step, y, y_name, y_hat, y_hat_name, colls = {}
     for coll in colls:
         c.add(coll)
         c.get(coll).tensor_names = colls[coll]
-    c.export(os.path.join(trial_dir, COLLECTIONS_FILE_NAME))
+    c.export(os.path.join(trial_dir, TORASOLE_DEFAULT_COLLECTIONS_FILE_NAME))
 
 
 def test_confusion():
@@ -32,19 +34,19 @@ def test_confusion():
     run_id = str(uuid.uuid4())
     y = np.random.randint(cat_no,size=(100,))
     y_hat = y
-    gen_y_and_y_hat( path, run_id, 0, y, 'y', y_hat, 'y_hat', colls = { 'labels' : [ 'y'], 'preds': ['y_hat'] } )
+    gen_y_and_y_hat( path, run_id, 0, y, 'y', y_hat, 'y_hat', colls = { 'labels' : [ 'y'], 'preds': ['y_hat']})
     tr = create_trial(os.path.join(path, run_id))
-    r = Confusion(tr, labels_collection='labels', predictions_collection='preds' )
+    r = Confusion(tr, labels_collection='labels', predictions_collection='preds')
     invoke_rule(r, start_step=0, end_step=1, raise_eval_cond=True)
 
     # Test 2: should fail on row 4 because the
     run_id = str(uuid.uuid4())
     y = np.arange(cat_no)
     y_hat = np.copy(y)
     y_hat[4] = 7
-    gen_y_and_y_hat( path, run_id, 1, y, 'foo', y_hat, 'bar' )
+    gen_y_and_y_hat( path, run_id, 1, y, 'foo', y_hat, 'bar')
     tr = create_trial(os.path.join(path, run_id))
-    r = Confusion(tr, cat_no, 'foo', 'bar' )
+    r = Confusion(tr, cat_no, 'foo', 'bar')
     try:
         invoke_rule(r, start_step=1, end_step=2, raise_eval_cond=True)
         assert False
@@ -64,7 +66,7 @@ def test_confusion():
     y = np.arange(10)
     y_hat = y
     # 'label' and 'pred' are magic names
-    gen_y_and_y_hat( path, run_id, 1, y, 'labels', y_hat, 'predictions' )
+    gen_y_and_y_hat( path, run_id, 1, y, 'labels', y_hat, 'predictions')
     tr = create_trial(os.path.join(path, run_id))
     r = Confusion(tr)
     invoke_rule(r, start_step=1, end_step=2, raise_eval_cond=True)

tests/analysis/trials/test_modes.py

@@ -5,7 +5,9 @@
 from tornasole.core.tensor import StepState
 from datetime import datetime
 from tornasole.core.writer import FileWriter
-from tornasole.core.collection_manager import CollectionManager, COLLECTIONS_FILE_NAME
+from tornasole.core.collection_manager import CollectionManager
+from tornasole.core.config_constants import TORASOLE_DEFAULT_COLLECTIONS_FILE_NAME
+
 
 def test_modes_on_global_data():
   pass # other tests in create, local, s3 do this
@@ -18,7 +20,7 @@ def test_mode_data():
   c = CollectionManager()
   c.add("default")
   c.get("default").tensor_names = ["arr"]
-  c.export(os.path.join(trial_dir, COLLECTIONS_FILE_NAME))
+  c.export(os.path.join(trial_dir, TORASOLE_DEFAULT_COLLECTIONS_FILE_NAME))
   tr = create_trial(trial_dir)
   for s in range(0, 10):
     fw = FileWriter(trial_dir=trial_dir, step=s)

tests/analysis/trials/test_s3.py

@@ -1,6 +1,6 @@
 from tornasole.trials import S3Trial
-from tornasole.core.collection_manager import CollectionManager, \
-    COLLECTIONS_FILE_NAME
+from tornasole.core.collection_manager import CollectionManager
+from tornasole.core.config_constants import TORASOLE_DEFAULT_COLLECTIONS_FILE_NAME
 import uuid
 import os
 import pytest
@@ -33,8 +33,8 @@ def help_test_multiple_trials(num_steps = 20, num_tensors = 10):
     c = CollectionManager()
     c.add("default")
     c.get("default").tensor_names = ["foo_" + str(i) for i in range(num_tensors)]
-    c.export(path + trial_name + "/" + COLLECTIONS_FILE_NAME)
-    c.export(path + trial_name + "/" + COLLECTIONS_FILE_NAME)
+    c.export(path + trial_name + "/" + TORASOLE_DEFAULT_COLLECTIONS_FILE_NAME)
+    c.export(path + trial_name + "/" + TORASOLE_DEFAULT_COLLECTIONS_FILE_NAME)
     for i in range(num_steps):
         generate_data(path=path, trial=trial_name, num_tensors=num_tensors,
                       step=i, tname_prefix='foo', worker='algo-1', shape=(3, 3, 3), rank=0)

tests/analysis/utils.py

@@ -1,12 +1,13 @@
 from tornasole.core.writer import FileWriter
 import numpy as np
-from tornasole.core.collection_manager import CollectionManager, \
-    COLLECTIONS_FILE_NAME
+from tornasole.core.collection_manager import CollectionManager
+from tornasole.core.config_constants import TORASOLE_DEFAULT_COLLECTIONS_FILE_NAME
 import os
 import aioboto3
 import asyncio
 from tornasole.core.access_layer.s3handler import S3Handler, ListRequest
 
+
 def generate_data(path, trial, step, tname_prefix,
                   num_tensors, worker, shape, dtype=np.float32,
                   rank=None, mode=None, mode_step=None, export_colls=True,
@@ -24,7 +25,7 @@ def generate_data(path, trial, step, tname_prefix,
         c.get("default").tensor_names = [f'{tname_prefix}_{i}' for i in range(num_tensors)]
         c.add('gradients')
         c.get("gradients").tensor_names = [f'{tname_prefix}_{i}' for i in range(num_tensors)]
-        c.export(os.path.join(path, trial, COLLECTIONS_FILE_NAME))
+        c.export(os.path.join(path, trial, TORASOLE_DEFAULT_COLLECTIONS_FILE_NAME))
 
 
 def check_trial(trial_obj, num_steps, num_tensors):

tests/core/test_collections.py

@@ -1,6 +1,6 @@
 from tornasole.core.collection import Collection
-from tornasole.core.collection_manager import CollectionManager, \
-  COLLECTIONS_FILE_NAME
+from tornasole.core.collection_manager import CollectionManager
+from tornasole.core.config_constants import TORASOLE_DEFAULT_COLLECTIONS_FILE_NAME
 from tornasole.core.reduction_config import ReductionConfig
 from tornasole.core.save_config import SaveConfig, SaveConfigMode
 from tornasole.core.modes import ModeKeys
@@ -42,8 +42,8 @@ def test_manager_export_load():
   cm.add(Collection('trial1'))
   cm.add('trial2')
   cm.get('trial2').include('total_loss')
-  cm.export(COLLECTIONS_FILE_NAME)
-  cm2 = CollectionManager.load(COLLECTIONS_FILE_NAME)
+  cm.export(TORASOLE_DEFAULT_COLLECTIONS_FILE_NAME)
+  cm2 = CollectionManager.load(TORASOLE_DEFAULT_COLLECTIONS_FILE_NAME)
   assert cm == cm2
 
 def test_manager():

tests/core/test_modes.py

@@ -16,7 +16,7 @@ def test_mode_writing():
     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()
+    fw.close()
   files = glob.glob('ts_outputs/' + run_id + '/**/*.tfevents',
                     recursive=True)
   for f in files:

tests/pytorch/test_simple_write.py

@@ -178,7 +178,8 @@ def helper_test_weights_bias_gradients(hook=None):
 
 
 def saveall_test_helper(hook=None):
-    reset_collections()
+    if hook is None:
+        reset_collections()
     prefix = str(uuid.uuid4())
     hook_type = 'saveall'
     device = torch.device("cpu")

tests/tensorflow/hooks/test_dist_horovod.py

@@ -0,0 +1,15 @@
+from tornasole.trials import create_trial
+
+
+def test_s3_read():
+    path = "s3://tornasole-testing/dist-logs-10/"
+    trial = create_trial(path)
+    tensors = trial.tensors()
+    assert len(tensors) == 17
+    t = trial.tensor('gradients/dense_1/MatMul_grad/tuple/control_dependency_1:0')
+    steps = t.steps()
+    assert steps == [0]
+    workers = t.workers_for_step(0)
+    assert len(workers) == 16
+    truth_table = t.value(0, worker='worker_10') == t.value(0, worker='worker_1')
+    assert truth_table.all() == False

tests/tensorflow/hooks/test_reductions.py

@@ -26,7 +26,7 @@ def helper_test_reductions(trial_dir, hook):
         assert len(t.reduction_values(0)) == 18
         for r in ALLOWED_REDUCTIONS + ALLOWED_NORMS:
             for b in [False, True]:
-                assert t.reduction_value(0, reduction_name=r, abs=b) is not None
+                assert t.reduction_value(0, reduction_name=r, abs=b, worker=None) is not None
 
 
 def test_reductions():

tests/tensorflow/hooks/test_save_all_full.py

@@ -3,7 +3,7 @@
 import shutil, glob
 from tornasole.core.reader import FileReader
 from tornasole.core.json_config import TORNASOLE_CONFIG_FILE_PATH_ENV_STR
-from tornasole.core.collection_manager import COLLECTIONS_FILE_NAME
+from tornasole.core.config_constants import TORASOLE_DEFAULT_COLLECTIONS_FILE_NAME
 
 
 def test_save_all_full(hook=None, trial_dir=None):
@@ -29,8 +29,8 @@ def test_save_all_full(hook=None, trial_dir=None):
     assert len(coll['gradients'].tensor_names) == 1
     assert len(coll['losses'].tensor_names) == 1
 
-    assert COLLECTIONS_FILE_NAME in files
-    cm = CollectionManager.load(join(trial_dir, COLLECTIONS_FILE_NAME))
+    assert TORASOLE_DEFAULT_COLLECTIONS_FILE_NAME in files
+    cm = CollectionManager.load(join(trial_dir, TORASOLE_DEFAULT_COLLECTIONS_FILE_NAME))
 
     assert len(cm.collections) == 6
     assert len(cm.collections['weights'].tensor_names) == 1

tests/tensorflow/hooks/test_save_config.py

@@ -1,5 +1,4 @@
 from tornasole.trials import create_trial
-from tornasole import SaveConfig
 from .utils import *
 from tests.tensorflow.hooks.test_estimator_modes import help_test_mnist
 from tornasole.tensorflow import reset_collections, get_collection, TornasoleHook, modes

tests/tensorflow/hooks/test_save_reductions.py

@@ -2,9 +2,9 @@
 from tornasole.tensorflow import reset_collections, get_collections, CollectionManager
 import shutil
 import glob
-from tornasole.core.collection_manager import COLLECTIONS_FILE_NAME
 from tornasole.core.reader import FileReader
 from tornasole.core.json_config import TORNASOLE_CONFIG_FILE_PATH_ENV_STR
+from tornasole.core.config_constants import TORASOLE_DEFAULT_COLLECTIONS_FILE_NAME
 
 def helper_save_reductions(trial_dir, hook):
   simple_model(hook)
@@ -14,8 +14,8 @@ def helper_save_reductions(trial_dir, hook):
   assert len(coll['weights'].tensor_names) == 1
   assert len(coll['gradients'].tensor_names) == 1
 
-  assert COLLECTIONS_FILE_NAME in files
-  cm = CollectionManager.load(join(trial_dir, COLLECTIONS_FILE_NAME))
+  assert TORASOLE_DEFAULT_COLLECTIONS_FILE_NAME in files
+  cm = CollectionManager.load(join(trial_dir, TORASOLE_DEFAULT_COLLECTIONS_FILE_NAME))
   assert len(cm.collections) == 5
   assert len(cm.collections['weights'].tensor_names) == 1
   assert len(cm.collections['gradients'].tensor_names) == 1

tests/tensorflow/hooks/test_simple_include.py

@@ -1,10 +1,11 @@
 from .utils import *
-from tornasole.tensorflow import reset_collections, get_collection
+from tornasole.tensorflow import get_collection
 import tornasole.tensorflow as ts
 import glob, shutil
 from tornasole.core.reader import FileReader
 from tornasole.core.json_config import TORNASOLE_CONFIG_FILE_PATH_ENV_STR
-from tornasole.core.collection_manager import COLLECTIONS_FILE_NAME
+from tornasole.core.config_constants import TORASOLE_DEFAULT_COLLECTIONS_FILE_NAME
+
 
 
 def helper_test_simple_include(trial_dir, hook):
@@ -13,7 +14,7 @@ def helper_test_simple_include(trial_dir, hook):
     _, files = get_dirs_files(trial_dir)
     steps, _ = get_dirs_files(os.path.join(trial_dir, 'events'))
 
-    cm = CollectionManager.load(join(trial_dir, COLLECTIONS_FILE_NAME))
+    cm = CollectionManager.load(join(trial_dir, TORASOLE_DEFAULT_COLLECTIONS_FILE_NAME))
     assert len(cm.collections['default'].tensor_names) == 1
     assert len(steps) == 5
     for step in steps:
@@ -57,7 +58,7 @@ def helper_test_simple_include_regex(trial_dir, hook):
     _, files = get_dirs_files(trial_dir)
     steps, _ = get_dirs_files(os.path.join(trial_dir, 'events'))
 
-    cm = CollectionManager.load(join(trial_dir, COLLECTIONS_FILE_NAME))
+    cm = CollectionManager.load(join(trial_dir, TORASOLE_DEFAULT_COLLECTIONS_FILE_NAME))
     assert len(cm.collections['default'].tensor_names) == 1
     assert len(steps) == 5
 
@@ -104,7 +105,7 @@ def helper_test_multi_collection_match(trial_dir, hook):
     _, files = get_dirs_files(trial_dir)
     steps, _ = get_dirs_files(os.path.join(trial_dir, 'events'))
 
-    cm = CollectionManager.load(join(trial_dir, COLLECTIONS_FILE_NAME))
+    cm = CollectionManager.load(join(trial_dir, TORASOLE_DEFAULT_COLLECTIONS_FILE_NAME))
     assert len(cm.collections['default'].tensor_names) == 1
     assert len(cm.collections['trial'].tensor_names) == 1
     assert len(steps) == 5

tests/tensorflow/hooks/test_weights_gradients.py

@@ -1,7 +1,7 @@
 from .utils import *
 from tornasole.tensorflow import reset_collections
 import tensorflow.compat.v1 as tf
-from tornasole.core.collection_manager import COLLECTIONS_FILE_NAME
+from tornasole.core.config_constants import TORASOLE_DEFAULT_COLLECTIONS_FILE_NAME
 from tornasole.core.json_config import TORNASOLE_CONFIG_FILE_PATH_ENV_STR
 import tornasole.tensorflow as ts
 import shutil
@@ -12,8 +12,8 @@ def helper_test_only_w_g(trial_dir, hook):
     steps, _ = get_dirs_files(os.path.join(trial_dir, 'events'))
     _, files = get_dirs_files(trial_dir)
 
-    assert COLLECTIONS_FILE_NAME in files
-    cm = CollectionManager.load(join(trial_dir, COLLECTIONS_FILE_NAME))
+    assert TORASOLE_DEFAULT_COLLECTIONS_FILE_NAME in files
+    cm = CollectionManager.load(join(trial_dir, TORASOLE_DEFAULT_COLLECTIONS_FILE_NAME))
     num_tensors_loaded_collection = len(cm.collections['weights'].tensor_names) + \
                                     len(cm.collections['gradients'].tensor_names) + \
                                     len(cm.collections['default'].tensor_names)

tests/xgboost/test_hook.py

@@ -1,5 +1,4 @@
 import os
-import json
 import uuid
 import numpy as np
 import pytest

tornasole/core/access_layer/file.py

@@ -12,9 +12,8 @@ def ensure_dir(file_path, is_file=True):
         directory = os.path.dirname(file_path)
     else:
         directory = file_path
-
     if directory and not os.path.exists(directory):
-        os.makedirs(directory)
+        os.makedirs(directory, exist_ok=True)
 
 
 def get_temp_path(file_path):

tornasole/core/collection_manager.py

@@ -3,8 +3,8 @@
 from .utils import is_s3, load_json_as_dict
 import json
 
-ALLOWED_PARAMS = ['collections']
-COLLECTIONS_FILE_NAME = 'collections.json'
+ALLOWED_PARAMS = ['collections', '_meta']
+
 
 class CollectionManager:
   """
@@ -16,6 +16,8 @@ def __init__(self, collections=None, create_default=False):
     if collections is None:
       collections = {}
     self.collections = collections
+    self._meta = {}
+    self._meta['num_workers'] = 1
 
   def create_collection(self, name, cls=Collection):
     if name not in self.collections:
@@ -40,6 +42,18 @@ def get(self, name, create=True):
         raise KeyError(f'Collection {name} has not been created')
     return self.collections[name]
 
+  def update_meta(self, meta):
+    assert isinstance(meta, dict)
+    self._meta.update(meta)
+
+
+  def get_num_workers(self):
+    return int(self._meta['num_workers'])
+
+
+  def set_num_workers(self, num_workers):
+    self._meta['num_workers'] = int(num_workers)
+
   def to_json_dict(self):
     d = dict()
     for a, v in self.__dict__.items():
@@ -81,6 +95,7 @@ def load_from_string(cls, s, collection_class=Collection):
       for c_name, c_dict in params['collections'].items():
         coll = collection_class.from_dict(c_dict)
         cm.add(coll)
+      cm.update_meta(params['_meta'])
       return cm
 
   def __eq__(self, other):
@@ -90,4 +105,4 @@ def __eq__(self, other):
 
   def __repr__(self):
     return f"<class CollectionManager: " \
-      f"collection_names={list(self.collections.keys())}>"
+      f"collection_names={list(self.collections.keys())}>"

tornasole/core/config_constants.py

@@ -1,4 +1,4 @@
-TORNASOLE_CONFIG_DEFAULT_WORKER_NAME = "worker0"
+TORNASOLE_CONFIG_DEFAULT_WORKER_NAME = "worker_0"
 TORNASOLE_CONFIG_FILE_PATH_ENV_STR = "TORNASOLE_CONFIG_FILE_PATH"
 DEFAULT_CONFIG_FILE_PATH = "/opt/ml/input/data/tornasole-config/tornasole-hook-config.json"
 TORNASOLE_CONFIG_REDUCTION_CONFIGS_KEY = "reduction_configs"
@@ -9,3 +9,4 @@
 TORNASOLE_CONFIG_INCLUDE_REGEX_KEY = "include_regex"
 TORNASOLE_CONFIG_SAVE_ALL_KEY = "save_all"
 DEFAULT_SAGEMAKER_TORNASOLE_PATH = "/opt/ml/output/tensors"
+TORASOLE_DEFAULT_COLLECTIONS_FILE_NAME = 'worker_0_collections.json'