Repository: pfnet/pytorch-pfn-extras

Issue:
Append mode Chrome Tracing profiler
https://docs.google.com/document/d/1CvAClvFfyA5R-PhYUmn5OOQtYMH4h6I0nSsKchNAySU/preview

> When provided as a string to the importer the ] at the end of the JSON Array Format is optional. The Trace Viewer importer will automatically add a ] if needed to turn the string into valid JSON.. This is to support tracing systems that can not cleanly finish writing the trace. For example, when tracing the exit of a program.

Since there is no need to add a trailing ], it is possible to implement appending a json string without removing it from the file.


Discussion:

Resolving patch:
diff --git a/pytorch_pfn_extras/profiler/__init__.py b/pytorch_pfn_extras/profiler/__init__.py
--- a/pytorch_pfn_extras/profiler/__init__.py
+++ b/pytorch_pfn_extras/profiler/__init__.py
@@ -11,4 +11,5 @@
     enable_global_trace,
     enable_thread_trace,
     get_tracer,
+    load_chrome_trace_as_json,
 )
diff --git a/pytorch_pfn_extras/profiler/_tracing.py b/pytorch_pfn_extras/profiler/_tracing.py
--- a/pytorch_pfn_extras/profiler/_tracing.py
+++ b/pytorch_pfn_extras/profiler/_tracing.py
@@ -12,6 +12,9 @@
 
 
 class Tracer:
+    def initialize_writer(self, filename: str, writer: Writer) -> None:
+        raise NotImplementedError("Tracers must implement initialize")
+
     @contextlib.contextmanager
     def add_event(self, name: str) -> Generator[None, None, None]:
         raise NotImplementedError("Tracers must implement add_event")
@@ -39,6 +42,9 @@ def load_state_dict(self, to_load: Dict[str, Any]) -> None:
 
 
 class DummyTracer(Tracer):
+    def initialize_writer(self, filename: str, writer: Writer) -> None:
+        pass
+
     @contextlib.contextmanager
     def add_event(self, name: str) -> Generator[None, None, None]:
         yield
@@ -51,10 +57,37 @@ def flush(self, filename: str, writer: Writer) -> None:
 
 
 class ChromeTracingSaveFunc:
-    def __call__(self, target: Dict[str, Any], file_o: Any) -> None:
+    def _write(self, target: List[Dict[str, Any]], file_o: Any) -> None:
         log = json.dumps(target, indent=4)
         file_o.write(log.encode("ascii"))
 
+    def _append(self, target: List[Dict[str, Any]], file_o: Any) -> None:
+        log = "".join(f"\n{json.dumps(o, indent=4)}," for o in target)
+        file_o.write(log.encode("ascii"))
+
+    def __call__(
+        self,
+        target: List[Dict[str, Any]],
+        file_o: Any,
+        append_mode: bool,
+    ) -> None:
+        if append_mode:
+            self._append(target, file_o)
+        else:
+            self._write(target, file_o)
+
+    def init(self, target: List[Dict[str, Any]], file_o: Any) -> None:
+        log = "["
+        file_o.write(log.encode("ascii"))
+
+
+def load_chrome_trace_as_json(filename: str) -> List[Dict[str, Any]]:
+    with open(filename) as f:
+        s = f.read()
+    if s[-1] != "]":
+        s = s[:-1] + "]"
+    return cast(List[Dict[str, Any]], json.loads(s))
+
 
 class ChromeTracer(Tracer):
     """Tracer object that outputs a timeline in Chrome format.
@@ -70,6 +103,7 @@ def __init__(
         self,
         max_event_count: Optional[int] = None,
         enable: bool = True,
+        append: bool = True,
     ) -> None:
         self._enable = enable
         self._event_list: List[Dict[str, Union[str, int, float]]] = []
@@ -82,6 +116,8 @@ def __init__(
             self.add_remote_event, 1000
         )
         self._tracer_queue.initialize()
+        self._append = append
+        self._savefun = ChromeTracingSaveFunc()
 
     @contextlib.contextmanager
     def add_event(self, name: str) -> Generator[None, None, None]:
@@ -130,19 +166,34 @@ def add_remote_event(
     ) -> None:
         self._event_list.append(event)
 
+    def initialize_writer(self, filename: str, writer: Writer) -> None:
+        if not self._enable:
+            return
+        if self._append:
+            writer(
+                filename,
+                "",
+                {},
+                savefun=self._savefun.init,
+                append=False,
+            )
+
     def flush(self, filename: str, writer: Writer) -> None:
         if not self._enable:
             return
         self._tracer_queue.synchronize()
         # TODO(ecastill): try to work on some append mode manipulating the
         # file pointer and with json.dumps?
-        savefun = ChromeTracingSaveFunc()
-        writer(
+        writer.save(
             filename,
             "",  # out_dir arg is ignored in the writer, uses the writer attr
             self._event_list,
-            savefun=savefun,
+            savefun=self._savefun,
+            append=self._append,
+            append_mode=self._append,
         )
+        if self._append:
+            self._event_list.clear()
 
     def enable(self, enable_flag: bool) -> None:
         self._enable = enable_flag
diff --git a/pytorch_pfn_extras/training/extensions/timeline_trace.py b/pytorch_pfn_extras/training/extensions/timeline_trace.py
--- a/pytorch_pfn_extras/training/extensions/timeline_trace.py
+++ b/pytorch_pfn_extras/training/extensions/timeline_trace.py
@@ -81,6 +81,10 @@ def _flush_trace(self, manager: ExtensionsManagerProtocol) -> None:
         if self._filename is not None:
             self._tracer.flush(self._filename, writer)
 
+    def initialize(self, manager: ExtensionsManagerProtocol) -> None:
+        writer = manager.writer if self._writer is None else self._writer
+        self._tracer.initialize_writer(self._filename, writer)
+
     def __call__(self, manager: ExtensionsManagerProtocol) -> None:
         if self._enable is not None and self._enable(manager):
             self._tracer.enable(True)


Test patch:
diff --git a/tests/pytorch_pfn_extras_tests/profiler_tests/test_record.py b/tests/pytorch_pfn_extras_tests/profiler_tests/test_record.py
--- a/tests/pytorch_pfn_extras_tests/profiler_tests/test_record.py
+++ b/tests/pytorch_pfn_extras_tests/profiler_tests/test_record.py
@@ -1,4 +1,3 @@
-import json
 import os
 import tempfile
 import threading
@@ -279,16 +278,17 @@ def test_record_iterable_with_multiprocessing(device):
 
     with tempfile.TemporaryDirectory() as t_path:
         w = ppe.writing.SimpleWriter(out_dir=t_path)
+        ppe.profiler.get_tracer().initialize_writer("trace.json", w)
         with torch.profiler.profile():
             for x in dataloader:
                 model(x.to(device))
         ppe.profiler.get_tracer().flush("trace.json", w)
         assert os.path.exists(os.path.join(t_path, "trace.json"))
         pid = os.getpid()
-        with open(os.path.join(t_path, "trace.json")) as f:
-            data = f.read()
-            values = json.loads(data)
-            assert len(values) == 5
-            # Check that the values were written by a dataloader worker
-            for v in values:
-                assert v["pid"] != pid
+        values = ppe.profiler.load_chrome_trace_as_json(
+            os.path.join(t_path, "trace.json")
+        )
+        assert len(values) == 5
+        # Check that the values were written by a dataloader worker
+        for v in values:
+            assert v["pid"] != pid
diff --git a/tests/pytorch_pfn_extras_tests/training_tests/extensions_tests/test_timeline_trace.py b/tests/pytorch_pfn_extras_tests/training_tests/extensions_tests/test_timeline_trace.py
--- a/tests/pytorch_pfn_extras_tests/training_tests/extensions_tests/test_timeline_trace.py
+++ b/tests/pytorch_pfn_extras_tests/training_tests/extensions_tests/test_timeline_trace.py
@@ -1,4 +1,3 @@
-import json
 import os
 import tempfile
 import time
@@ -28,10 +27,10 @@ def _body():
             for _ in range(iters_per_epoch):
                 with manager.run_iteration():
                     _body()
-        with open(os.path.join(tmpdir, "trace.json")) as f:
-            data = f.read()
-            values = json.loads(data)
-            assert len(values) == max_epochs * iters_per_epoch
+        values = ppe.profiler.load_chrome_trace_as_json(
+            os.path.join(tmpdir, "trace.json")
+        )
+        assert len(values) == max_epochs * iters_per_epoch
     ppe.profiler.clear_tracer()
 
 
@@ -63,10 +62,44 @@ def _body():
                     _body()
                 count += 1
                 assert tracer._event_count == count
-        with open(os.path.join(tmpdir, "trace.json")) as f:
-            data = f.read()
-            values = json.loads(data)
-            assert len(values) == max_epochs * iters_per_epoch
+        values = ppe.profiler.load_chrome_trace_as_json(
+            os.path.join(tmpdir, "trace.json")
+        )
+        assert len(values) == max_epochs * iters_per_epoch
+
+
+def test_tracer_no_append_object():
+    max_epochs = 3
+    iters_per_epoch = 5
+    tracer = ppe.profiler.ChromeTracer(append=False)
+    ext = ppe.training.extensions.TimelineTrace(
+        filename="trace.json", tracer=tracer
+    )
+
+    def _body():
+        with ppe.profiler.record("tag", trace=tracer):
+            time.sleep(0.1)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        manager = ppe.training.ExtensionsManager(
+            {},
+            {},
+            max_epochs=max_epochs,
+            iters_per_epoch=iters_per_epoch,
+            out_dir=tmpdir,
+        )
+        manager.extend(ext)
+        count = 0
+        for _epoch_idx in range(max_epochs):
+            for _ in range(iters_per_epoch):
+                with manager.run_iteration():
+                    _body()
+                count += 1
+                assert tracer._event_count == count
+        values = ppe.profiler.load_chrome_trace_as_json(
+            os.path.join(tmpdir, "trace.json")
+        )
+        assert len(values) == max_epochs * iters_per_epoch
 
 
 def test_tracer_enable():
@@ -100,7 +133,7 @@ def _body():
                 with manager.run_iteration():
                     _body()
                 count += 1
-        with open(os.path.join(tmpdir, "trace.json")) as f:
-            data = f.read()
-            values = json.loads(data)
-            assert len(values) == (max_epochs * iters_per_epoch - 10 + 5)
+        values = ppe.profiler.load_chrome_trace_as_json(
+            os.path.join(tmpdir, "trace.json")
+        )
+        assert len(values) == (max_epochs * iters_per_epoch - 10 + 5)