Source Metadata: Parse source writer information from the first event of each run (#6014)

Motivation:
Googlers, see go/tb-writer-source-metadata for context.

Note that:
- The current assumption is that each run would have at most one source writer.
- The maximum allowed length for the writer name is set to 128. 
- For simplicity, only source writer name rather than the entire `SourceMetadata` proto is parsed (since this is the only field).
- In the future if more fields are added, the `SourceWriter` property can be easily generalized to `SourceMetadata` in the `EventAccumulator`.

Author: yatbear

tensorboard/backend/event_processing/BUILD

@@ -208,6 +208,30 @@ py_test(
     ],
 )
 
+py_library(
+    name = "event_util",
+    srcs = ["event_util.py"],
+    srcs_version = "PY3",
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorboard/compat/proto:protos_all_py_pb2",
+        "//tensorboard/util:tb_logging",
+    ],
+)
+
+py_test(
+    name = "event_util_test",
+    size = "small",
+    srcs = ["event_util_test.py"],
+    srcs_version = "PY3",
+    deps = [
+        ":event_util",
+        "//tensorboard:test",
+        "//tensorboard/compat/proto:protos_all_py_pb2",
+        "//tensorboard/util:tb_logging",
+    ],
+)
+
 py_library(
     name = "tag_types",
     srcs = ["tag_types.py"],
@@ -226,6 +250,7 @@ py_library(
         ":directory_loader",
         ":directory_watcher",
         ":event_file_loader",
+        ":event_util",
         ":io_wrapper",
         ":plugin_asset_util",
         ":reservoir",

tensorboard/backend/event_processing/event_accumulator.py

@@ -17,8 +17,11 @@
 import collections
 import threading
 
+from typing import Optional
+
 from tensorboard.backend.event_processing import directory_watcher
 from tensorboard.backend.event_processing import event_file_loader
+from tensorboard.backend.event_processing import event_util
 from tensorboard.backend.event_processing import io_wrapper
 from tensorboard.backend.event_processing import plugin_asset_util
 from tensorboard.backend.event_processing import reservoir
@@ -224,6 +227,9 @@ def __init__(
         self.most_recent_wall_time = -1
         self.file_version = None
 
+        # Name of the source writer that writes the event.
+        self._source_writer = None
+
         # The attributes that get built up by the accumulator
         self.accumulated_attrs = (
             "scalars",
@@ -302,6 +308,21 @@ def FirstEventTimestamp(self):
             except StopIteration:
                 raise ValueError("No event timestamp could be found")
 
+    @property
+    def SourceWriter(self) -> Optional[str]:
+        """Returns the name of the event writer."""
+        if self._source_writer is not None:
+            return self._source_writer
+        with self._generator_mutex:
+            try:
+                event = next(self._generator.Load())
+                self._ProcessEvent(event)
+                return self._source_writer
+            except StopIteration:
+                logger.info(
+                    "End of file in %s, no source writer was found.", self.path
+                )
+
     def PluginTagToContent(self, plugin_name):
         """Returns a dict mapping tags to content specific to that plugin.
 
@@ -339,8 +360,22 @@ def _ProcessEvent(self, event):
         if self._first_event_timestamp is None:
             self._first_event_timestamp = event.wall_time
 
+        if event.HasField("source_metadata"):
+            new_source_writer = event_util.GetSourceWriter(
+                event.source_metadata
+            )
+            if self._source_writer and self._source_writer != new_source_writer:
+                # This should not happen.
+                logger.warning(
+                    (
+                        "Found new source writer for event.proto. "
+                        "Old: {0}, New: {1}"
+                    ).format(self._source_writer, new_source_writer)
+                )
+            self._source_writer = new_source_writer
+
         if event.HasField("file_version"):
-            new_file_version = _ParseFileVersion(event.file_version)
+            new_file_version = event_util.ParseFileVersion(event.file_version)
             if self.file_version and self.file_version != new_file_version:
                 ## This should not happen.
                 logger.warning(
@@ -824,27 +859,3 @@ def _GeneratorFromPath(path):
             event_file_loader.LegacyEventFileLoader,
             io_wrapper.IsSummaryEventsFile,
         )
-
-
-def _ParseFileVersion(file_version):
-    """Convert the string file_version in event.proto into a float.
-
-    Args:
-      file_version: String file_version from event.proto
-
-    Returns:
-      Version number as a float.
-    """
-    tokens = file_version.split("brain.Event:")
-    try:
-        return float(tokens[-1])
-    except ValueError:
-        ## This should never happen according to the definition of file_version
-        ## specified in event.proto.
-        logger.warning(
-            (
-                "Invalid event.proto file_version. Defaulting to use of "
-                "out-of-order event.step logic for purging expired events."
-            )
-        )
-        return -1

tensorboard/backend/event_processing/event_accumulator_test.py

@@ -709,11 +709,65 @@ def testFirstEventTimestampLoadsEvent(self):
         gen.AddEvent(
             event_pb2.Event(wall_time=1, step=2, file_version="brain.Event:2")
         )
-
         self.assertEqual(acc.FirstEventTimestamp(), 1)
         acc.Reload()
         self.assertEqual(acc.file_version, 2.0)
 
+    def testSourceWriter(self):
+        gen = _EventGenerator(self)
+        acc = ea.EventAccumulator(gen)
+        gen.AddEvent(
+            event_pb2.Event(
+                wall_time=10,
+                step=20,
+                source_metadata=event_pb2.SourceMetadata(
+                    writer="custom_writer"
+                ),
+            )
+        )
+        gen.AddScalar("s1", wall_time=30, step=40, value=20)
+        self.assertEqual(acc.SourceWriter, "custom_writer")
+
+    def testReloadPopulatesSourceWriter(self):
+        """Test that Reload() means SourceWriter won't load events."""
+        gen = _EventGenerator(self)
+        acc = ea.EventAccumulator(gen)
+        gen.AddEvent(
+            event_pb2.Event(
+                wall_time=1,
+                step=2,
+                source_metadata=event_pb2.SourceMetadata(
+                    writer="custom_writer"
+                ),
+            )
+        )
+        acc.Reload()
+
+        def _Die(*args, **kwargs):  # pylint: disable=unused-argument
+            raise RuntimeError("Load() should not be called")
+
+        self.stubs.Set(gen, "Load", _Die)
+        self.assertEqual(acc.SourceWriter, "custom_writer")
+
+    def testSourceWriterLoadsEvent(self):
+        """Test that SourceWriter doesn't discard the loaded event."""
+        gen = _EventGenerator(self)
+        acc = ea.EventAccumulator(gen)
+        gen.AddEvent(
+            event_pb2.Event(
+                wall_time=1,
+                step=2,
+                file_version="brain.Event:2",
+                source_metadata=event_pb2.SourceMetadata(
+                    writer="custom_writer"
+                ),
+            )
+        )
+
+        self.assertEqual(acc.SourceWriter, "custom_writer")
+        acc.Reload()
+        self.assertEqual(acc.file_version, 2.0)
+
     def testTFSummaryScalar(self):
         """Verify processing of tf.summary.scalar."""
         event_sink = _EventGenerator(self, zero_out_timestamps=True)

tensorboard/backend/event_processing/event_multiplexer.py

@@ -18,6 +18,7 @@
 import os
 import threading
 
+from typing import Optional
 
 from tensorboard.backend.event_processing import directory_watcher
 from tensorboard.backend.event_processing import event_accumulator
@@ -258,6 +259,21 @@ def FirstEventTimestamp(self, run):
         accumulator = self.GetAccumulator(run)
         return accumulator.FirstEventTimestamp()
 
+    def GetSourceWriter(self, run) -> Optional[str]:
+        """Returns the source writer name from the first event of the given run.
+
+        Assuming each run has only one source writer.
+
+        Args:
+          run: A string name of the run from which the event source information
+            is retrieved.
+
+        Returns:
+          Name of the writer that wrote the events in the run.
+        """
+        accumulator = self.GetAccumulator(run)
+        return accumulator.SourceWriter
+
     def Scalars(self, run, tag):
         """Retrieve the scalar events associated with a run and tag.
 

tensorboard/backend/event_processing/event_multiplexer_test.py

@@ -67,6 +67,10 @@ def Tags(self):
     def FirstEventTimestamp(self):
         return 0
 
+    @property
+    def SourceWriter(self):
+        return "%s_writer" % self._path
+
     def _TagHelper(self, tag_name, enum):
         if tag_name not in self.Tags()[enum]:
             raise KeyError
@@ -151,6 +155,13 @@ def testReload(self):
         self.assertTrue(x.GetAccumulator("run1").reload_called)
         self.assertTrue(x.GetAccumulator("run2").reload_called)
 
+    def testGetSourceWriter(self):
+        x = event_multiplexer.EventMultiplexer(
+            {"run1": "path1", "run2": "path2"}
+        )
+        self.assertEqual(x.GetSourceWriter("run1"), "path1_writer")
+        self.assertEqual(x.GetSourceWriter("run2"), "path2_writer")
+
     def testScalars(self):
         """Tests Scalars function returns suitable values."""
         x = event_multiplexer.EventMultiplexer(

tensorboard/backend/event_processing/event_util.py

@@ -0,0 +1,68 @@
+# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Functionality for processing events."""
+
+from typing import Optional
+
+from tensorboard.compat.proto import event_pb2
+from tensorboard.util import tb_logging
+
+logger = tb_logging.get_logger()
+
+# Maxmimum length for event writer name.
+_MAX_WRITER_NAME_LEN = 128
+
+
+def ParseFileVersion(file_version: str) -> float:
+    """Convert the string file_version in event.proto into a float.
+
+    Args:
+      file_version: String file_version from event.proto
+
+    Returns:
+      Version number as a float.
+    """
+    tokens = file_version.split("brain.Event:")
+    try:
+        return float(tokens[-1])
+    except ValueError:
+        ## This should never happen according to the definition of file_version
+        ## specified in event.proto.
+        logger.warning(
+            (
+                "Invalid event.proto file_version. Defaulting to use of "
+                "out-of-order event.step logic for purging expired events."
+            )
+        )
+        return -1
+
+
+def GetSourceWriter(
+    source_metadata: event_pb2.SourceMetadata,
+) -> Optional[str]:
+    """Gets the source writer name from the source metadata proto."""
+    writer_name = source_metadata.writer
+    if not writer_name:
+        return None
+    # Checks the length of the writer name.
+    if len(writer_name) > _MAX_WRITER_NAME_LEN:
+        logger.error(
+            "Source writer name `%s` is too long, maximum allowed length is %d.",
+            writer_name,
+            _MAX_WRITER_NAME_LEN,
+        )
+        return None
+    return writer_name

tensorboard/backend/event_processing/event_util_test.py

@@ -0,0 +1,73 @@
+# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Tests for event_util."""
+
+from unittest import mock
+
+from tensorboard import test as tb_test
+from tensorboard.backend.event_processing import event_util
+from tensorboard.compat.proto import event_pb2
+from tensorboard.util import tb_logging
+
+logger = tb_logging.get_logger()
+
+
+class EventUtilTest(tb_test.TestCase):
+    def testParseFileVersion_success(self):
+        self.assertEqual(event_util.ParseFileVersion("brain.Event:1.0"), 1.0)
+
+    def testParseFileVersion_invalidFileVersion(self):
+        with mock.patch.object(
+            logger, "warning", autospec=True, spec_set=True
+        ) as mock_log:
+            version = event_util.ParseFileVersion("invalid")
+            self.assertEqual(version, -1)
+        mock_log.assert_called_once_with(
+            "Invalid event.proto file_version. Defaulting to use of "
+            "out-of-order event.step logic for purging expired events."
+        )
+
+    def testGetSourceWriter_success(self):
+        expected_writer = "tensorboard.summary.writer.event_file_writer"
+        actual_writer = event_util.GetSourceWriter(
+            event_pb2.SourceMetadata(writer=expected_writer)
+        )
+        self.assertEqual(actual_writer, expected_writer)
+
+    def testGetSourceWriter_noWriter(self):
+        actual_writer = event_util.GetSourceWriter(
+            event_pb2.SourceMetadata(writer="")
+        )
+        self.assertIsNone(actual_writer)
+
+    def testGetSourceWriter_writerNameTooLong(self):
+        long_writer_name = "a" * (event_util._MAX_WRITER_NAME_LEN + 1)
+        with mock.patch.object(
+            logger, "error", autospec=True, spec_set=True
+        ) as mock_log:
+            actual_writer = event_util.GetSourceWriter(
+                event_pb2.SourceMetadata(writer=long_writer_name)
+            )
+            self.assertIsNone(actual_writer)
+        mock_log.assert_called_once_with(
+            "Source writer name `%s` is too long, maximum allowed length is %d.",
+            long_writer_name,
+            event_util._MAX_WRITER_NAME_LEN,
+        )
+
+
+if __name__ == "__main__":
+    tb_test.main()