Save scalar (aws#352)

* Write metrics to a file

* Write contents to metric file

* Enable the code to save_scalar

* cache save_scalar before prepare_collections

* Modify Minerva log file format

* Test save_scalar

* Remove TF save_scalar. Modify tests

* Update test

* Move up an assert

* Remove xgboost changes for now

* Fix CodeBuild

* Change function names in test

* Minerva file format change, Eureka SDK integration

* Write loss to Minerva

* Added some comments

* Log scalars before closing write

* Address review comments, add searchable scalars

* Remove redundant code

* Minor changes

* Enable tensorboard

* Add test for TF searchable scalar

* Add some comments to the test

* Fix CodeBuild

* Use wrap_optimizer

* Fix CodeBuild

* Move metrics file writer out of tfevents

* Keras TF save_scalar

* Fix regression with Eureka integration

* Keras TF save_scalar

* Fix regression with Eureka integration

* Flush out scalars before closing the file

* Fix build error

* close metrics writer after other writers

* Address review comments

* Correct path in the test

* Correct file path

* Fix regression

* combine initialize_writer

Author: vandanavk

smdebug/core/collection.py

@@ -27,6 +27,7 @@ class CollectionKeys:
     LOSSES = "losses"
     BIASES = "biases"
     SCALARS = "scalars"
+    SEARCHABLE_SCALARS = "searchable_scalars"
 
     OPTIMIZER_VARIABLES = "optimizer_variables"
     TENSORFLOW_SUMMARIES = "tensorflow_summaries"
@@ -45,7 +46,12 @@ class CollectionKeys:
 # so we don't create summaries or reductions of these
 SUMMARIES_COLLECTIONS = {CollectionKeys.TENSORFLOW_SUMMARIES}
 
-SCALAR_COLLECTIONS = {CollectionKeys.LOSSES, CollectionKeys.METRICS, CollectionKeys.SCALARS}
+SCALAR_COLLECTIONS = {
+    CollectionKeys.LOSSES,
+    CollectionKeys.METRICS,
+    CollectionKeys.SCALARS,
+    CollectionKeys.SEARCHABLE_SCALARS,
+}
 
 # used by pt, mx, keras
 NON_REDUCTION_COLLECTIONS = SCALAR_COLLECTIONS.union(SUMMARIES_COLLECTIONS)

smdebug/core/config_constants.py

@@ -15,6 +15,7 @@
 CONFIG_INCLUDE_REGEX_KEY = "include_regex"
 CONFIG_SAVE_ALL_KEY = "save_all"
 TENSORBOARD_CONFIG_FILE_PATH_ENV_STR = "TENSORBOARD_CONFIG_FILE_PATH"
+DEFAULT_SAGEMAKER_METRICS_PATH = "SAGEMAKER_METRICS_DIRECTORY"
 DEFAULT_SAGEMAKER_OUTDIR = "/opt/ml/output/tensors"
 DEFAULT_SAGEMAKER_TENSORBOARD_PATH = "/opt/ml/input/config/tensorboardoutputconfig.json"
 DEFAULT_COLLECTIONS_FILE_NAME = "worker_0_collections.json"

smdebug/core/hook.py

@@ -32,6 +32,11 @@
 from smdebug.core.utils import flatten, get_tb_worker, match_inc, size_and_shape
 from smdebug.core.writer import FileWriter
 
+try:
+    from smexperiments.metrics import SageMakerFileMetricsWriter
+except ImportError:
+    from smdebug.core.metrics_file_writer import SageMakerFileMetricsWriter
+
 logger = get_logger()
 
 
@@ -153,8 +158,15 @@ def __init__(
         self.mode = ModeKeys.GLOBAL
         self.mode_steps = {ModeKeys.GLOBAL: init_step}
         self.writer = None
+
+        self.metrics_writer = None
+
         # Maps ModeKeys to FileWriter objects
         self.tb_writers = {}
+
+        # Cache scalars that are being saved through save_scalar() calls
+        self.scalar_cache = []
+
         self.logger.info("Saving to {}".format(self.out_dir))
         atexit.register(self._cleanup)
 
@@ -309,6 +321,12 @@ def _close_writers(self) -> None:
         if self.dry_run:
             return
 
+        # flush out searchable scalars to metrics file
+        if self.metrics_writer is not None:
+            self._write_scalars()
+            self.metrics_writer.close()
+            self.metrics_writer = None
+
         self._close_writer()
         to_delete_writers = []
 
@@ -321,10 +339,11 @@ def _close_writers(self) -> None:
         for mode in to_delete_writers:
             del self.tb_writers[mode]
 
-    def _initialize_writer(self) -> None:
+    def _initialize_writers(self) -> None:
         if self.dry_run:
             return
         self.writer = FileWriter(trial_dir=self.out_dir, step=self.step, worker=self.worker)
+        self.metrics_writer = SageMakerFileMetricsWriter()
 
     def get_writers(self, tensor_name, tensor_ref=None) -> List[FileWriter]:
         """
@@ -470,6 +489,16 @@ def _write_scalar_summary(self, tensor_name, tensor_value, save_colls):
                             f"so scalar summary could not be created"
                         )
                     break
+        for s_col in save_colls:
+            if s_col.name in [
+                CollectionKeys.LOSSES,
+                CollectionKeys.SEARCHABLE_SCALARS,
+                CollectionKeys.METRICS,
+            ]:
+                np_val = self._make_numpy_array(tensor_value)
+                # Always log loss to Minerva
+                tensor_val = np.mean(np_val)
+                self.scalar_cache.append((tensor_name, tensor_val, True))
 
     def _write_histogram_summary(self, tensor_name, tensor_value, save_collections):
         """ Maybe write to TensorBoard. """
@@ -490,18 +519,45 @@ def _write_histogram_summary(self, tensor_name, tensor_value, save_collections):
                     )
                     break
 
+    def _write_scalars(self):
+        """
+        This function writes all the scalar values saved in the scalar_cache to file.
+        If searchable is set to True for certain scalars, then that scalar is written to
+        Minerva as well. By default, loss values are searchable.
+        """
+        if self.writer is None:
+            self._initialize_writers()
+        tb_writer = self._maybe_get_tb_writer()
+        for scalar_name, scalar_val, searchable in self.scalar_cache:
+            save_collections = self._get_collections_with_tensor(scalar_name)
+            logger.debug(
+                f"Saving scalar {scalar_name} {scalar_val} for step {self.step} {self.mode} "
+                f"{self.mode_steps[self.mode]}"
+            )
+            if searchable:
+                self.metrics_writer.log_metric(scalar_name, scalar_val, self.mode_steps[self.mode])
+            if tb_writer:
+                self._write_raw_tensor(scalar_name, scalar_val, save_collections)
+            self.scalar_cache = []
+
     # Fix step number for saving scalar and tensor
-    # def save_scalar(self, name, value):
-    #     get_collection(CollectionKeys.SCALARS).add_tensor_name(name)
-    #     if self.writer is None:
-    #         self._init_writer()
-    #     val = make_numpy_array(value)
-    #     if val.size != 1:
-    #         raise TypeError(
-    #             f'{name} has non scalar value of type: {type(value)}')
-    #     self._save_scalar_summary(name, val)
-    #     logger.debug(f'Saving scalar {name} {val} for step {self.step} {self.mode} {self.mode_steps[self.mode]}')
-    #     self._save_raw_tensor(name, val)
+    def save_scalar(self, name, value, searchable=False):
+        """
+        Call save_scalar at any point in the training script to log a scalar value,
+        such as a metric or any other value.
+        :param name: Name of the scalar. A prefix 'scalar/' will be added to it
+        :param value: Scalar value
+        :param searchable: True/False. If set to True, the scalar value will be written to
+        SageMaker Minerva
+        """
+        name = CallbackHook.SCALAR_PREFIX + name
+        val = self._make_numpy_array(value)
+        if val.size != 1:
+            raise TypeError(f"{name} has non scalar value of type: {type(value)}")
+        self.collection_manager.get(CollectionKeys.SCALARS).add_tensor_name(name)
+        self.scalar_cache.append((name, val, searchable))
+        if self.prepared_collections:
+            self._write_scalars()
 
     # def save_tensor(self, name, value):
     #     # todo: support to add these tensors to any collection.
@@ -627,6 +683,7 @@ class CallbackHook(BaseHook):
     INPUT_TENSOR_SUFFIX = "_input_"
     OUTPUT_TENSOR_SUFFIX = "_output_"
     GRADIENT_PREFIX = "gradient/"
+    SCALAR_PREFIX = "scalar/"
 
     def __init__(
         self,

smdebug/core/metrics_file_writer.py

@@ -0,0 +1,75 @@
+# Standard Library
+import json
+import os
+import time
+
+# First Party
+from smdebug.core.config_constants import DEFAULT_SAGEMAKER_METRICS_PATH
+
+METRICS_DIR = os.environ.get(DEFAULT_SAGEMAKER_METRICS_PATH, ".")
+
+
+class _RawMetricData(object):
+    def __init__(self, metric_name, value, iteration_number, timestamp):
+        self.MetricName = metric_name
+        self.Value = value
+        self.Timestamp = timestamp
+        self.IterationNumber = iteration_number
+
+
+class SageMakerFileMetricsWriter(object):
+    def __init__(self, filename=None):
+        self._file = open(filename or self._metrics_file_name(), "a")
+        self._indexes = {}
+        self._closed = False
+
+    def _metrics_file_name(self):
+        return "{}/{}.json".format(METRICS_DIR, str(os.getpid()))
+
+    def _write_metric_value(self, file, raw_metric_data):
+        try:
+            self._file.write(json.dumps(raw_metric_data.__dict__))
+            self._file.write("\n")
+        except AttributeError:
+            if self._closed:
+                raise ValueError("log_metric called on a closed writer")
+            elif not self._file:
+                self._file = open(self._metrics_file_name(), "a")
+                self._file.write(json.dumps(raw_metric_data.__dict__))
+                self._file.write("\n")
+            else:
+                raise
+
+    def log_metric(self, metric_name, value, iteration_number=None, timestamp=None):
+        timestamp = int(round(time.time())) if timestamp is None else int(timestamp)
+        resolved_index = int(
+            self._indexes.get(metric_name, 0) if iteration_number is None else iteration_number
+        )
+
+        value = float(value)
+        assert isinstance(resolved_index, int)
+        assert isinstance(timestamp, int)
+
+        self._write_metric_value(
+            self._file, _RawMetricData(metric_name, value, iteration_number, timestamp)
+        )
+        if not iteration_number:
+            self._indexes[metric_name] = resolved_index + 1
+
+    def close(self):
+        if not self._closed and self._file:
+            self._file.close()
+            self._file = None
+        self._closed = True
+
+    def __enter__(self):
+        """Return self"""
+        return self
+
+    def __exit__(self, type, value, traceback):
+        """Execute self.close()"""
+        self.close()
+
+    def __del__(self):
+        """Execute self.close()"""
+        self.close()

smdebug/mxnet/collection.py

@@ -25,6 +25,7 @@ def _register_default_collections(self):
         self.get(CollectionKeys.BIASES).include("^(?!gradient).*bias")
         self.get(CollectionKeys.GRADIENTS).include("^gradient")
         self.get(CollectionKeys.LOSSES).include(".*loss")
+        self.get(CollectionKeys.SCALARS).include("^scalar")
 
     def create_collection(self, name):
         super().create_collection(name, cls=Collection)

smdebug/mxnet/hook.py

@@ -17,6 +17,7 @@
     CollectionKeys.BIASES,
     CollectionKeys.GRADIENTS,
     CollectionKeys.LOSSES,
+    CollectionKeys.SCALARS,
 ]
 
 
@@ -140,7 +141,7 @@ def forward_pre_hook(self, block, inputs):
         self._increment_step()
 
         if self._get_collections_to_save_for_step():
-            self._initialize_writer()
+            self._initialize_writers()
 
         if self.exported_model is False:
             self._export_model()

smdebug/pytorch/collection.py

@@ -40,6 +40,7 @@ def _register_default_collections(self):
         self.get(CollectionKeys.BIASES).include("^(?!gradient).*bias")
         self.get(CollectionKeys.GRADIENTS).include("^gradient")
         self.get(CollectionKeys.LOSSES).include("[Ll]oss")
+        self.get(CollectionKeys.SCALARS).include("^scalar")
 
     def create_collection(self, name):
         super().create_collection(name, cls=Collection)

smdebug/pytorch/hook.py

@@ -12,7 +12,7 @@
 from smdebug.pytorch.singleton_utils import set_hook
 from smdebug.pytorch.utils import get_reduction_of_data, make_numpy_array
 
-DEFAULT_INCLUDE_COLLECTIONS = [CollectionKeys.LOSSES]
+DEFAULT_INCLUDE_COLLECTIONS = [CollectionKeys.LOSSES, CollectionKeys.SCALARS]
 
 
 class Hook(CallbackHook):
@@ -43,6 +43,10 @@ def __init__(
             include_collections=include_collections,
             save_all=save_all,
         )
+        # We would like to collect loss collection
+        # even if user does not specify any collections
+        if CollectionKeys.LOSSES not in self.include_collections:
+            self.include_collections.append(CollectionKeys.LOSSES)
         # mapping of module objects to their names,
         # useful in forward hook for logging input/output of modules
         self.module_maps = dict()
@@ -143,7 +147,7 @@ def forward_pre_hook(self, module, inputs):
         self._increment_step()
 
         if self._get_collections_to_save_for_step():
-            self._initialize_writer()
+            self._initialize_writers()
             self.log_params(module)
 
         if self.last_saved_step is not None and not self.exported_collections:

smdebug/tensorflow/base_hook.py

@@ -27,10 +27,17 @@
     is_parameter_server_strategy,
 )
 
+try:
+    from smexperiments.metrics import SageMakerFileMetricsWriter
+except ImportError:
+    from smdebug.core.metrics_file_writer import SageMakerFileMetricsWriter
+
+
 DEFAULT_INCLUDE_COLLECTIONS = [
     CollectionKeys.METRICS,
     CollectionKeys.LOSSES,
     CollectionKeys.SCALARS,
+    CollectionKeys.SEARCHABLE_SCALARS,
 ]
 
 
@@ -183,7 +190,7 @@ def get_writers(self, tensor_name, tensor_ref) -> List[FileWriter]:
         else:
             return [self.writer]
 
-    def _initialize_writer(self, only_initialize_if_missing=False) -> None:
+    def _initialize_writers(self, only_initialize_if_missing=False) -> None:
         # In keras, sometimes we are not sure if writer is initialized
         # (such as metrics at end of epoch), that's why it passes the flag only_init_if_missing
 
@@ -203,11 +210,19 @@ def _initialize_writer(self, only_initialize_if_missing=False) -> None:
         else:
             if self.writer is None or only_initialize_if_missing is False:
                 self.writer = FileWriter(trial_dir=self.out_dir, step=self.step, worker=self.worker)
+            if self.metrics_writer is None or only_initialize_if_missing is False:
+                self.metrics_writer = SageMakerFileMetricsWriter()
 
     def _close_writer(self) -> None:
         if self.dry_run:
             return
 
+        # flush out searchable scalars to metrics file
+        if self.metrics_writer is not None:
+            self._write_scalars()
+            self.metrics_writer.close()
+            self.metrics_writer = None
+
         if self.writer is not None:
             self.writer.flush()
             self.writer.close()

smdebug/tensorflow/collection.py

@@ -138,6 +138,7 @@ def __init__(self, collections=None, create_default=True):
                 CollectionKeys.INPUTS,
                 CollectionKeys.OUTPUTS,
                 CollectionKeys.ALL,
+                CollectionKeys.SEARCHABLE_SCALARS,
             ]:
                 self.create_collection(n)
             self.get(CollectionKeys.BIASES).include("bias")

smdebug/tensorflow/keras.py

@@ -316,7 +316,7 @@ def _save_metrics(self, batch, logs, force_save=False):
             return
 
         if force_save or self._is_collection_being_saved_for_step(CollectionKeys.METRICS):
-            self._initialize_writer(only_initialize_if_missing=True)
+            self._initialize_writers(only_initialize_if_missing=True)
             logs["batch"] = batch
             for key in logs:
                 if key in ["loss", "val_loss", "outputs"]:
@@ -326,7 +326,7 @@ def _save_metrics(self, batch, logs, force_save=False):
                 self._save_for_tensor(key, logs[key], check_before_write=False)
 
         if force_save or self._is_collection_being_saved_for_step(CollectionKeys.LOSSES):
-            self._initialize_writer(only_initialize_if_missing=True)
+            self._initialize_writers(only_initialize_if_missing=True)
             for key in ["loss", "val_loss"]:
                 if key in logs:
                     self._add_metric(metric_name=key)
@@ -442,7 +442,7 @@ def _on_any_batch_begin(self, batch, mode, logs=None):
 
         if self.tensor_refs_to_save_this_step:
             # if saving metric, writer may not be initialized as a result
-            self._initialize_writer()
+            self._initialize_writers()
 
         self._add_callbacks(mode)
 

smdebug/tensorflow/session.py

@@ -306,7 +306,7 @@ def _get_all_tensors_values(self, results):
 
     def after_run(self, run_context, run_values):
         if self.tensors_to_save_this_step:
-            self._initialize_writer()
+            self._initialize_writers()
             for (tensor, value) in self._get_all_tensors_values(run_values.results):
                 if tensor.dtype == tf.string:
                     self._write_tf_summary(tensor, value)