Enable file rotation for Horovod trace file (#33)

* Hvd file reader and rotation of files

Co-authored-by: Anirudh <[email protected]>

Author: vandanavk

examples/tensorflow2/scripts/tf_keras_fit_hvd.py

@@ -0,0 +1,146 @@
+"""
+This script is a simple MNIST training script which uses Horovod and Tensorflow 2.x Keras interface.
+It is designed to be used with SageMaker Debugger in an official SageMaker Framework container (i.e. AWS Deep Learning Container).
+You will notice that this script looks exactly like a normal TensorFlow training script.
+The hook needed by SageMaker Debugger to save tensors during training will be automatically added in those environments.
+The hook will load configuration from json configuration that SageMaker will put in the training container from the
+configuration provided using the SageMaker python SDK when creating a job.
+For more information, please refer to https://github.com/awslabs/sagemaker-debugger/blob/master/docs/sagemaker.md
+
+This script has been adapted from an example in Horovod repository https://github.com/uber/horovod
+"""
+
+# Standard Library
+import argparse
+from datetime import datetime
+
+# Third Party
+import horovod.tensorflow.keras as hvd
+import tensorflow.compat.v2 as tf
+
+
+def str2bool(v):
+    if isinstance(v, bool):
+        return v
+    if v.lower() in ("yes", "true", "t", "y", "1"):
+        return True
+    elif v.lower() in ("no", "false", "f", "n", "0"):
+        return False
+    else:
+        raise argparse.ArgumentTypeError("Boolean value expected.")
+
+
+def get_data(batch_size):
+    (mnist_images, mnist_labels), _ = tf.keras.datasets.mnist.load_data(
+        path="mnist-%d.npz" % hvd.rank()
+    )
+
+    dataset = tf.data.Dataset.from_tensor_slices(
+        (
+            tf.cast(mnist_images[..., tf.newaxis] / 255.0, tf.float32),
+            tf.cast(mnist_labels, tf.int64),
+        )
+    )
+    dataset = dataset.repeat().shuffle(10000).batch(batch_size)
+    return dataset
+
+
+def get_model():
+    mnist_model = tf.keras.Sequential(
+        [
+            tf.keras.layers.Conv2D(32, [3, 3], activation="relu", input_shape=(28, 28, 1)),
+            tf.keras.layers.Conv2D(64, [3, 3], activation="relu"),
+            tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
+            tf.keras.layers.Dropout(0.25),
+            tf.keras.layers.Flatten(),
+            tf.keras.layers.Dense(128, activation="relu"),
+            tf.keras.layers.Dropout(0.5),
+            tf.keras.layers.Dense(10, activation="softmax"),
+        ]
+    )
+    return mnist_model
+
+
+def train(model, dataset, epoch):
+    # Horovod: Specify `experimental_run_tf_function=False` to ensure TensorFlow
+    # uses hvd.DistributedOptimizer() to compute gradients.
+    model.compile(
+        loss=tf.losses.SparseCategoricalCrossentropy(),
+        optimizer=opt,
+        metrics=["accuracy"],
+        experimental_run_tf_function=False,
+    )
+
+    # Create a TensorBoard callback
+    logs = "logs/" + datetime.now().strftime("%Y%m%d-%H%M%S")
+
+    tboard_callback = tf.keras.callbacks.TensorBoard(
+        log_dir=logs, histogram_freq=1, profile_batch=2
+    )
+    callbacks = [
+        # Horovod: broadcast initial variable states from rank 0 to all other processes.
+        # This is necessary to ensure consistent initialization of all workers when
+        # training is started with random weights or restored from a checkpoint.
+        hvd.callbacks.BroadcastGlobalVariablesCallback(0),
+        # Horovod: average metrics among workers at the end of every epoch.
+        #
+        # Note: This callback must be in the list before the ReduceLROnPlateau,
+        # TensorBoard or other metrics-based callbacks.
+        hvd.callbacks.MetricAverageCallback(),
+        # Horovod: using `lr = 1.0 * hvd.size()` from the very beginning leads to worse final
+        # accuracy. Scale the learning rate `lr = 1.0` ---> `lr = 1.0 * hvd.size()` during
+        # the first three epochs. See https://arxiv.org/abs/1706.02677 for details.
+        hvd.callbacks.LearningRateWarmupCallback(warmup_epochs=3, verbose=1),
+        tboard_callback,
+    ]
+    # Horovod: save checkpoints only on worker 0 to prevent other workers from corrupting them.
+    if hvd.rank() == 0:
+        callbacks.append(tf.keras.callbacks.ModelCheckpoint("checkpoint-{epoch}.h5"))
+
+    # Horovod: write logs on worker 0.
+    verbose = 1 if hvd.rank() == 0 else 0
+
+    # Train the model.
+    # Horovod: adjust number of steps based on number of GPUs.
+    model.fit(
+        dataset,
+        steps_per_epoch=500 // hvd.size(),
+        callbacks=callbacks,
+        epochs=epoch,
+        verbose=verbose,
+    )
+
+
+if __name__ == "__main__":
+    # Training settings
+    parser = argparse.ArgumentParser(description="Tensorflow2 MNIST Example")
+    parser.add_argument("--use_only_cpu", type=str2bool, default=False)
+    parser.add_argument("--num_epochs", type=int, default=1, help="Number of epochs to train for")
+    parser.add_argument("--model_dir", type=str, default="/tmp/mnist_model")
+
+    args = parser.parse_args()
+
+    # constants
+    lr = 0.001
+    batch_size = 64
+
+    # Horovod: initialize library.
+    hvd.init()
+
+    # Horovod: pin GPU to be used to process local rank (one GPU per process)
+    gpus = tf.config.experimental.list_physical_devices("GPU")
+    for gpu in gpus:
+        tf.config.experimental.set_memory_growth(gpu, True)
+    if gpus:
+        tf.config.experimental.set_visible_devices(gpus[hvd.local_rank()], "GPU")
+
+    # Horovod: adjust learning rate based on number of GPUs.
+    opt = tf.optimizers.Adam(lr * hvd.size())
+
+    # Horovod: add Horovod DistributedOptimizer.
+    opt = hvd.DistributedOptimizer(opt)
+
+    dataset = get_data(batch_size)
+    mnist_model = get_model()
+
+    train(model=mnist_model, dataset=dataset, epoch=args.num_epochs)

smdebug/core/hook.py

@@ -228,6 +228,7 @@ def __init__(
         self.timeline_writer = TimelineFileWriter(
             profiler_config_parser=self.profiler_config_parser
         )
+        self.hvd_reader = None
 
         if is_sagemaker_job() and SageMakerFileMetricsWriter is not None:
             self.metrics_writer = SageMakerFileMetricsWriter()
@@ -515,6 +516,10 @@ def _cleanup(self):
 
         self.timeline_writer.close()
 
+        # close the Horovod file reader thread if it has been enabled
+        if self.hvd_reader and self.hvd_reader.enabled:
+            self.hvd_reader.close()
+
         training_has_ended(self.out_dir)
         if self.first_process is True:
             remove_claim_file(self.out_dir)

smdebug/core/locations.py

@@ -120,7 +120,7 @@ class TraceFileLocation:
     # $ENV_BASE_FOLDER/framework/pevents/$START_TIME_YYYYMMDDHR/
     # $FILEEVENTENDTIMEUTCINEPOCH_{$ENV_NODE_ID}_model_timeline.json
     @staticmethod
-    def get_file_location(timestamp, base_dir):
+    def get_file_location(timestamp, base_dir, suffix=PYTHONTIMELINE_SUFFIX):
         env_base_location = base_dir
         date_hour = time.strftime(
             TRACE_DIRECTORY_FORMAT, time.gmtime(timestamp / CONVERT_TO_MICROSECS)
@@ -137,7 +137,7 @@ def get_file_location(timestamp, base_dir):
             + "_"
             + worker_id
             + "_"
-            + PYTHONTIMELINE_SUFFIX,
+            + suffix,
         )
         return file_path
 

smdebug/core/tfevent/timeline_file_writer.py

@@ -33,7 +33,7 @@
 from smdebug.core.locations import TraceFileLocation
 from smdebug.core.logger import get_logger
 from smdebug.core.utils import ensure_dir, get_node_id
-from smdebug.profiler.profiler_constants import CONVERT_TO_MICROSECS
+from smdebug.profiler.profiler_constants import CONVERT_TO_MICROSECS, PYTHONTIMELINE_SUFFIX
 
 logger = get_logger()
 
@@ -88,7 +88,18 @@ def to_json(self):
             "ph": self.phase,
             "ts": self.rel_ts_micros,
         }
-        if self.phase == "X":
+
+        # handle Instant event
+        if self.phase == "i":
+            if self.args:
+                # Instant events have a field unique to them called scope.
+                # scope can be "g" - global, "p" - process, "t" - thread.
+                # parsing this value that is being passed as args.
+                s = self.args["s"] if "s" in self.args else "t"
+                json_dict.update({"s": s})
+                if "s" in self.args:
+                    self.args.pop("s")
+        elif self.phase == "X":
             json_dict.update({"dur": self.duration})
 
         if self.args:
@@ -105,7 +116,7 @@ class TimelineFileWriter:
     and asynchronously writes TimelineRecord to the file.
     """
 
-    def __init__(self, profiler_config_parser, max_queue=100):
+    def __init__(self, profiler_config_parser, max_queue=100, suffix=PYTHONTIMELINE_SUFFIX):
         """Creates a `TimelineFileWriter` and a trace event file to write to.
         This event file will contain TimelineRecord as JSON strings, which are written to
         disk via the write_record method.
@@ -120,14 +131,34 @@ def __init__(self, profiler_config_parser, max_queue=100):
         self._worker = _TimelineLoggerThread(
             queue=self._event_queue,
             sentinel_event=self._sentinel_event,
-            base_start_time=self.start_time_since_epoch_in_micros,
+            base_start_time_in_us=self.start_time_since_epoch_in_micros,
             profiler_config_parser=self._profiler_config_parser,
+            suffix=suffix,
         )
         self._worker.start()
 
+    def _update_base_start_time(self, base_start_time_in_us):
+        """
+        Some trace files such as the Horovod trace file may start before this timeline
+        writer is initialized. In such case, use this function to update the start time
+        since epoch in micros.
+        """
+        if base_start_time_in_us != self.start_time_since_epoch_in_micros:
+            self.start_time_since_epoch_in_micros = base_start_time_in_us
+            self._worker._update_base_start_time(base_start_time_in_us)
+
     def write_trace_events(
         self, timestamp, training_phase="", op_name="", phase="X", duration=0, **kwargs
     ):
+        """
+        Creates TimelineRecord from the details passed as parameters, and enqueues an event for write.
+        :param timestamp:start_time for the event (in seconds)
+        :param training_phase: strings like, data_iteration, forward, backward, operations etc
+        :param op_name: more details about phase like whether dataset or iterator
+        :param phase: phase of trace event. default is 'X'
+        :param duration: any duration manually computed (in seconds)
+        :param kwargs: other params. can be process id and thread id
+        """
         if not self._worker._healthy or not self._profiler_config_parser.profiling_enabled:
             return
         duration_in_us = int(duration * CONVERT_TO_MICROSECS)  # convert to micro seconds
@@ -167,7 +198,13 @@ class _TimelineLoggerThread(threading.Thread):
     https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/summary/writer/event_file_writer.py#L133"""
 
     def __init__(
-        self, queue, sentinel_event, base_start_time, profiler_config_parser, verbose=False
+        self,
+        queue,
+        sentinel_event,
+        base_start_time_in_us,
+        profiler_config_parser,
+        verbose=False,
+        suffix=PYTHONTIMELINE_SUFFIX,
     ):
         """Creates a _TimelineLoggerThread."""
         threading.Thread.__init__(self)
@@ -180,14 +217,23 @@ def __init__(
         self.tensor_table = collections.defaultdict(int)
         self.continuous_fail_count = 0
         self.is_first = True
-        self.last_event_end_time_in_us = int(round(base_start_time))
+        self._update_base_start_time(base_start_time_in_us)
+        self._healthy = True
+        self._profiler_config_parser = profiler_config_parser
+        self.node_id = get_node_id()
+        self.suffix = suffix
+
+    def _update_base_start_time(self, base_start_time_in_us):
+        """
+        Some trace files such as the Horovod trace file may start before this timeline
+        writer is initialized. In such case, use this function to update the start time
+        since epoch in micros.
+        """
+        self.last_event_end_time_in_us = int(round(base_start_time_in_us))
         self.last_file_close_time_in_us = self.last_event_end_time_in_us
         self.cur_hour = datetime.utcfromtimestamp(
             self.last_file_close_time_in_us / CONVERT_TO_MICROSECS
         ).hour
-        self._healthy = True
-        self._profiler_config_parser = profiler_config_parser
-        self.node_id = get_node_id()
 
     def run(self):
         while True:
@@ -315,13 +361,20 @@ def write_event(self, record):
                 json_dict = {"name": "process_sort_index", "ph": "M", "pid": 0, "args": args}
                 self._writer.write(json.dumps(json_dict) + ",\n")
 
-            args = {"name": record.training_phase}
-            json_dict = {"name": "process_name", "ph": "M", "pid": tensor_idx, "args": args}
-            self._writer.write(json.dumps(json_dict) + ",\n")
+            # Instant events don't have a training phase
+            if record.phase != "i":
+                args = {"name": record.training_phase}
+                json_dict = {"name": "process_name", "ph": "M", "pid": tensor_idx, "args": args}
+                self._writer.write(json.dumps(json_dict) + ",\n")
 
-            args = {"sort_index": tensor_idx}
-            json_dict = {"name": "process_sort_index", "ph": "M", "pid": tensor_idx, "args": args}
-            self._writer.write(json.dumps(json_dict) + ",\n")
+                args = {"sort_index": tensor_idx}
+                json_dict = {
+                    "name": "process_sort_index",
+                    "ph": "M",
+                    "pid": tensor_idx,
+                    "args": args,
+                }
+                self._writer.write(json.dumps(json_dict) + ",\n")
 
             self.is_first = False
 
@@ -366,6 +419,7 @@ def close(self):
                 new_file_name = TraceFileLocation().get_file_location(
                     base_dir=self._profiler_config_parser.config.local_path,
                     timestamp=self.last_event_end_time_in_us,
+                    suffix=self.suffix,
                 )
                 ensure_dir(new_file_name)
                 os.rename(self.name(), new_file_name)
@@ -378,6 +432,8 @@ def name(self):
             self._profiler_config_parser.config.local_path
             + "/framework/"
             + self.node_id
+            + "_"
+            + self.suffix
             + SMDEBUG_TEMP_PATH_SUFFIX
         )
 

smdebug/profiler/algorithm_metrics_reader.py

@@ -10,7 +10,7 @@
 from smdebug.profiler.profiler_constants import (
     DEFAULT_PREFIX,
     ENV_TIME_BUFFER,
-    HOROVODTIMELINE_PREFIX,
+    HOROVODTIMELINE_SUFFIX,
     MODELTIMELINE_SUFFIX,
     PYTHONTIMELINE_SUFFIX,
     TENSORBOARDTIMELINE_SUFFIX,
@@ -122,7 +122,7 @@ def _get_event_parser(self, filename):
             return self._SMEventsParser
         if TENSORBOARDTIMELINE_SUFFIX in filename:
             return self._TBEventsParser
-        if HOROVODTIMELINE_PREFIX in filename:
+        if HOROVODTIMELINE_SUFFIX in filename:
             return self._HorovordEventsParser
 
     def _get_timestamp_from_filename(self, event_file):