Allows summary writer to take an optional suffix= argument, which

is appended to every summary file name.
Change: 152736272

3 files changed, 34 insertions, 11 deletions

diff --git a/tensorflow/python/summary/writer/event_file_writer.py b/tensorflow/python/summary/writer/event_file_writer.py
index 7142998ce7..8940d9b72e 100644
--- a/tensorflow/python/summary/writer/event_file_writer.py
+++ b/tensorflow/python/summary/writer/event_file_writer.py
@@ -37,7 +37,8 @@ class EventFileWriter(object):
   is encoded using the tfrecord format, which is similar to RecordIO.
   """
 
-  def __init__(self, logdir, max_queue=10, flush_secs=120):
+  def __init__(self, logdir, max_queue=10, flush_secs=120,
+               filename_suffix=None):
     """Creates a `EventFileWriter` and an event file to write to.
 
     On construction the summary writer creates a new event file in `logdir`.
@@ -57,6 +58,8 @@ class EventFileWriter(object):
       max_queue: Integer. Size of the queue for pending events and summaries.
       flush_secs: Number. How often, in seconds, to flush the
         pending events and summaries to disk.
+      filename_suffix: A string. Every event file's name is suffixed with
+        `filename_suffix`.
     """
     self._logdir = logdir
     if not gfile.IsDirectory(self._logdir):
@@ -64,6 +67,8 @@ class EventFileWriter(object):
     self._event_queue = six.moves.queue.Queue(max_queue)
     self._ev_writer = pywrap_tensorflow.EventsWriter(
         compat.as_bytes(os.path.join(self._logdir, "events")))
+    if filename_suffix:
+      self._ev_writer.InitWithSuffix(compat.as_bytes(filename_suffix))
     self._closed = False
     self._worker = _EventLoggerThread(self._event_queue, self._ev_writer,
                                       flush_secs)
diff --git a/tensorflow/python/summary/writer/writer.py b/tensorflow/python/summary/writer/writer.py
index 2dad7b88d3..05f97fb284 100644
--- a/tensorflow/python/summary/writer/writer.py
+++ b/tensorflow/python/summary/writer/writer.py
@@ -32,7 +32,6 @@ from tensorflow.python.platform import tf_logging as logging
 from tensorflow.python.summary import plugin_asset
 from tensorflow.python.summary.writer.event_file_writer import EventFileWriter
 
-
 _PLUGINS_DIR = "plugins"
 
 
@@ -81,12 +80,11 @@ class SummaryToEventTransformer(object):
       self.add_graph(graph=graph, graph_def=graph_def)
       # Also export the meta_graph_def in this case.
       # graph may itself be a graph_def due to positional arguments
-      maybe_graph_as_def = (
-          graph.as_graph_def(add_shapes=True) if isinstance(graph, ops.Graph)
-          else graph)
+      maybe_graph_as_def = (graph.as_graph_def(add_shapes=True)
+                            if isinstance(graph, ops.Graph) else graph)
       self.add_meta_graph(
-          meta_graph.create_meta_graph_def(
-              graph_def=graph_def or maybe_graph_as_def))
+          meta_graph.create_meta_graph_def(graph_def=graph_def or
+                                           maybe_graph_as_def))
 
   def add_summary(self, summary, global_step=None):
     """Adds a `Summary` protocol buffer to the event file.
@@ -214,8 +212,8 @@ class SummaryToEventTransformer(object):
       TypeError: If both `meta_graph_def` is not an instance of `MetaGraphDef`.
     """
     if not isinstance(meta_graph_def, meta_graph_pb2.MetaGraphDef):
-      raise TypeError("meta_graph_def must be type MetaGraphDef, saw type: %s"
-                      % type(meta_graph_def))
+      raise TypeError("meta_graph_def must be type MetaGraphDef, saw type: %s" %
+                      type(meta_graph_def))
     meta_graph_bytes = meta_graph_def.SerializeToString()
     event = event_pb2.Event(meta_graph_def=meta_graph_bytes)
     self._add_event(event, global_step)
@@ -266,7 +264,8 @@ class FileWriter(SummaryToEventTransformer):
                graph=None,
                max_queue=10,
                flush_secs=120,
-               graph_def=None):
+               graph_def=None,
+               filename_suffix=None):
     """Creates a `FileWriter` and an event file.
 
     On construction the summary writer creates a new event file in `logdir`.
@@ -304,8 +303,11 @@ class FileWriter(SummaryToEventTransformer):
       flush_secs: Number. How often, in seconds, to flush the
         pending events and summaries to disk.
       graph_def: DEPRECATED: Use the `graph` argument instead.
+      filename_suffix: A string. Every event file's name is suffixed with
+        `suffix`.
     """
-    event_writer = EventFileWriter(logdir, max_queue, flush_secs)
+    event_writer = EventFileWriter(logdir, max_queue, flush_secs,
+                                   filename_suffix)
     super(FileWriter, self).__init__(event_writer, graph, graph_def)
 
   def get_logdir(self):
diff --git a/tensorflow/python/summary/writer/writer_test.py b/tensorflow/python/summary/writer/writer_test.py
index 50797483ee..b31c41d112 100644
--- a/tensorflow/python/summary/writer/writer_test.py
+++ b/tensorflow/python/summary/writer/writer_test.py
@@ -308,6 +308,22 @@ class SummaryWriterTestCase(test.TestCase):
     # We should be done.
     self.assertRaises(StopIteration, lambda: next(rr))
 
+  def testFileWriterWithSuffix(self):
+    test_dir = self._CleanTestDir("test_suffix")
+    sw = writer.FileWriter(test_dir, filename_suffix="_test_suffix")
+    for _ in range(10):
+      sw.add_summary(
+          summary_pb2.Summary(value=[
+              summary_pb2.Summary.Value(tag="float_ten", simple_value=10.0)
+          ]),
+          10)
+      sw.close()
+      sw.reopen()
+    sw.close()
+    event_filenames = glob.glob(os.path.join(test_dir, "event*"))
+    for filename in event_filenames:
+      self.assertTrue(filename.endswith("_test_suffix"))
+
 
 class SummaryWriterCacheTest(test.TestCase):
   """SummaryWriterCache tests."""