xref: /aosp_15_r20/external/tensorflow/tensorflow/python/debug/lib/debug_events_writer.py (revision b6fb3261f9314811a0f4371741dbb8839866f948) 
1# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ==============================================================================
15"""Writer class for `DebugEvent` protos in tfdbg v2."""
16
17import time
18
19from tensorflow.core.protobuf import debug_event_pb2
20from tensorflow.python.client import _pywrap_debug_events_writer
21
22# Default size of each circular buffer (unit: number of DebugEvent protos).
23DEFAULT_CIRCULAR_BUFFER_SIZE = 1000
24
25
26class DebugEventsWriter(object):
27  """A writer for TF debugging events. Used by tfdbg v2."""
28
29  def __init__(self,
30               dump_root,
31               tfdbg_run_id,
32               circular_buffer_size=DEFAULT_CIRCULAR_BUFFER_SIZE):
33    """Construct a DebugEventsWriter object.
34
35    NOTE: Given the same `dump_root`, all objects from this constructor
36      will point to the same underlying set of writers. In other words, they
37      will write to the same set of debug events files in the `dump_root`
38      folder.
39
40    Args:
41      dump_root: The root directory for dumping debug data. If `dump_root` does
42        not exist as a directory, it will be created.
43      tfdbg_run_id: Debugger Run ID.
44      circular_buffer_size: Size of the circular buffer for each of the two
45        execution-related debug events files: with the following suffixes: -
46          .execution - .graph_execution_traces If <= 0, the circular-buffer
47          behavior will be abolished in the constructed object.
48    """
49    if not dump_root:
50      raise ValueError("Empty or None dump root")
51    self._dump_root = dump_root
52    self._tfdbg_run_id = tfdbg_run_id
53    _pywrap_debug_events_writer.Init(self._dump_root, self._tfdbg_run_id,
54                                     circular_buffer_size)
55
56  def WriteSourceFile(self, source_file):
57    """Write a SourceFile proto with the writer.
58
59    Args:
60      source_file: A SourceFile proto, describing the content of a source file
61        involved in the execution of the debugged TensorFlow program.
62    """
63    # TODO(cais): Explore performance optimization that avoids memcpy.
64    debug_event = debug_event_pb2.DebugEvent(source_file=source_file)
65    self._EnsureTimestampAdded(debug_event)
66    _pywrap_debug_events_writer.WriteSourceFile(self._dump_root, debug_event)
67
68  def WriteStackFrameWithId(self, stack_frame_with_id):
69    """Write a StackFrameWithId proto with the writer.
70
71    Args:
72      stack_frame_with_id: A StackFrameWithId proto, describing the content a
73        stack frame involved in the execution of the debugged TensorFlow
74        program.
75    """
76    debug_event = debug_event_pb2.DebugEvent(
77        stack_frame_with_id=stack_frame_with_id)
78    self._EnsureTimestampAdded(debug_event)
79    _pywrap_debug_events_writer.WriteStackFrameWithId(self._dump_root,
80                                                      debug_event)
81
82  def WriteGraphOpCreation(self, graph_op_creation):
83    """Write a GraphOpCreation proto with the writer.
84
85    Args:
86      graph_op_creation: A GraphOpCreation proto, describing the details of the
87        creation of an op inside a TensorFlow Graph.
88    """
89    debug_event = debug_event_pb2.DebugEvent(
90        graph_op_creation=graph_op_creation)
91    self._EnsureTimestampAdded(debug_event)
92    _pywrap_debug_events_writer.WriteGraphOpCreation(self._dump_root,
93                                                     debug_event)
94
95  def WriteDebuggedGraph(self, debugged_graph):
96    """Write a DebuggedGraph proto with the writer.
97
98    Args:
99      debugged_graph: A DebuggedGraph proto, describing the details of a
100        TensorFlow Graph that has completed its construction.
101    """
102    debug_event = debug_event_pb2.DebugEvent(debugged_graph=debugged_graph)
103    self._EnsureTimestampAdded(debug_event)
104    _pywrap_debug_events_writer.WriteDebuggedGraph(self._dump_root, debug_event)
105
106  def WriteExecution(self, execution):
107    """Write a Execution proto with the writer.
108
109    Args:
110      execution: An Execution proto, describing a TensorFlow op or graph
111        execution event.
112    """
113    debug_event = debug_event_pb2.DebugEvent(execution=execution)
114    self._EnsureTimestampAdded(debug_event)
115    _pywrap_debug_events_writer.WriteExecution(self._dump_root, debug_event)
116
117  def WriteGraphExecutionTrace(self, graph_execution_trace):
118    """Write a GraphExecutionTrace proto with the writer.
119
120    Args:
121      graph_execution_trace: A GraphExecutionTrace proto, concerning the value
122        of an intermediate tensor or a list of intermediate tensors that are
123        computed during the graph's execution.
124    """
125    debug_event = debug_event_pb2.DebugEvent(
126        graph_execution_trace=graph_execution_trace)
127    self._EnsureTimestampAdded(debug_event)
128    _pywrap_debug_events_writer.WriteGraphExecutionTrace(
129        self._dump_root, debug_event)
130
131  def RegisterDeviceAndGetId(self, device_name):
132    return _pywrap_debug_events_writer.RegisterDeviceAndGetId(
133        self._dump_root, device_name)
134
135  def FlushNonExecutionFiles(self):
136    """Flush the non-execution debug event files."""
137    _pywrap_debug_events_writer.FlushNonExecutionFiles(self._dump_root)
138
139  def FlushExecutionFiles(self):
140    """Flush the execution debug event files.
141
142    Causes the current content of the cyclic buffers to be written to
143    the .execution and .graph_execution_traces debug events files.
144    Also clears those cyclic buffers.
145    """
146    _pywrap_debug_events_writer.FlushExecutionFiles(self._dump_root)
147
148  def Close(self):
149    """Close the writer."""
150    _pywrap_debug_events_writer.Close(self._dump_root)
151
152  @property
153  def dump_root(self):
154    return self._dump_root
155
156  def _EnsureTimestampAdded(self, debug_event):
157    if debug_event.wall_time == 0:
158      debug_event.wall_time = time.time()
159