Refactor tf.Operation.traceback implementation in to methods of tf.Graph.
Adds an `_extract_frame_info` method to allow derived classes to extend the information available in each op traceback, if desired. The default result of `tf.Operation.traceback` is unchanged. Also fixes a poorly scoped `pylint disable=line-too-long`, so adds the necessary enable/disable blocks to silence pylint for the offending docstrings. PiperOrigin-RevId: 157466174
This commit is contained in:
Tim Harley
TensorFlower Gardener
parent
f7ca8db7d8
commit
b73fea6e27
@ -88,70 +88,6 @@ def _override_helper(clazz_object, operator, func):
|
||||
setattr(clazz_object, operator, func)
|
||||
|
||||
|
||||
def _convert_stack(stack, include_func_start_lineno=False):
|
||||
"""Converts a stack extracted using _extract_stack() to a traceback stack.
|
||||
|
||||
Args:
|
||||
stack: A list of n 5-tuples,
|
||||
(filename, lineno, name, frame_globals, func_start_lineno).
|
||||
include_func_start_lineno: True if function start line number should be
|
||||
included as the 5th entry in return tuples.
|
||||
|
||||
Returns:
|
||||
A list of n 4-tuples or 5-tuples
|
||||
(filename, lineno, name, code, [optional: func_start_lineno]), where the
|
||||
code tuple element is calculated from the corresponding elements of the
|
||||
input tuple.
|
||||
"""
|
||||
ret = []
|
||||
for filename, lineno, name, frame_globals, func_start_lineno in stack:
|
||||
linecache.checkcache(filename)
|
||||
line = linecache.getline(filename, lineno, frame_globals)
|
||||
if line:
|
||||
line = line.strip()
|
||||
else:
|
||||
line = None
|
||||
if include_func_start_lineno:
|
||||
ret.append((filename, lineno, name, line, func_start_lineno))
|
||||
else:
|
||||
ret.append((filename, lineno, name, line))
|
||||
return ret
|
||||
|
||||
|
||||
# pylint: disable=line-too-long
|
||||
def _extract_stack():
|
||||
"""A lightweight re-implementation of traceback.extract_stack.
|
||||
|
||||
NOTE(mrry): traceback.extract_stack eagerly retrieves the line of code for
|
||||
each stack frame using linecache, which results in an abundance of stat()
|
||||
calls. This implementation does not retrieve the code, and any consumer
|
||||
should apply _convert_stack to the result to obtain a traceback that can
|
||||
be formatted etc. using traceback methods.
|
||||
|
||||
Returns:
|
||||
A list of 5-tuples
|
||||
(filename, lineno, name, frame_globals, func_start_lineno) corresponding to
|
||||
the call stack of the current thread.
|
||||
"""
|
||||
# pylint: enable=line-too-long
|
||||
try:
|
||||
raise ZeroDivisionError
|
||||
except ZeroDivisionError:
|
||||
f = sys.exc_info()[2].tb_frame.f_back
|
||||
ret = []
|
||||
while f is not None:
|
||||
lineno = f.f_lineno
|
||||
co = f.f_code
|
||||
filename = co.co_filename
|
||||
name = co.co_name
|
||||
frame_globals = f.f_globals
|
||||
func_start_lineno = co.co_firstlineno
|
||||
ret.append((filename, lineno, name, frame_globals, func_start_lineno))
|
||||
f = f.f_back
|
||||
ret.reverse()
|
||||
return ret
|
||||
|
||||
|
||||
def _as_graph_element(obj):
|
||||
"""Convert `obj` to a graph element if possible, otherwise return `None`.
|
||||
|
||||
@ -1264,7 +1200,7 @@ class Operation(object):
|
||||
|
||||
self._original_op = original_op
|
||||
self._op_def = op_def
|
||||
self._traceback = _extract_stack()
|
||||
self._traceback = self._graph._extract_stack() # pylint: disable=protected-access
|
||||
# Add this op to the current control flow context:
|
||||
self._control_flow_context = g._get_control_flow_context()
|
||||
if self._control_flow_context is not None:
|
||||
@ -1613,6 +1549,7 @@ class Operation(object):
|
||||
|
||||
@property
|
||||
def node_def(self):
|
||||
# pylint: disable=line-too-long
|
||||
"""Returns a serialized `NodeDef` representation of this operation.
|
||||
|
||||
Returns:
|
||||
@ -1620,10 +1557,12 @@ class Operation(object):
|
||||
[`NodeDef`](https://www.tensorflow.org/code/tensorflow/core/framework/node_def.proto)
|
||||
protocol buffer.
|
||||
"""
|
||||
# pylint: enable=line-too-long
|
||||
return self._node_def
|
||||
|
||||
@property
|
||||
def op_def(self):
|
||||
# pylint: disable=line-too-long
|
||||
"""Returns the `OpDef` proto that represents the type of this op.
|
||||
|
||||
Returns:
|
||||
@ -1631,12 +1570,13 @@ class Operation(object):
|
||||
[`OpDef`](https://www.tensorflow.org/code/tensorflow/core/framework/op_def.proto)
|
||||
protocol buffer.
|
||||
"""
|
||||
# pylint: enable=line-too-long
|
||||
return self._op_def
|
||||
|
||||
@property
|
||||
def traceback(self):
|
||||
"""Returns the call stack from when this operation was constructed."""
|
||||
return _convert_stack(self._traceback)
|
||||
return self._graph._convert_stack(self._traceback) # pylint: disable=protected-access
|
||||
|
||||
@property
|
||||
def traceback_with_start_lines(self):
|
||||
@ -1645,7 +1585,8 @@ class Operation(object):
|
||||
Returns:
|
||||
A list of 5-tuples (filename, lineno, name, code, func_start_lineno).
|
||||
"""
|
||||
return _convert_stack(self._traceback, include_func_start_lineno=True)
|
||||
return self._graph._convert_stack( # pylint: disable=protected-access
|
||||
self._traceback, include_func_start_lineno=True)
|
||||
|
||||
def get_attr(self, name):
|
||||
"""Returns the value of the attr of this op with the given `name`.
|
||||
@ -2170,6 +2111,76 @@ class Graph(object):
|
||||
else:
|
||||
self._scoped_c_graph = None
|
||||
|
||||
def _convert_stack(self, stack, include_func_start_lineno=False):
|
||||
"""Converts a stack extracted using _extract_stack() to a traceback stack.
|
||||
|
||||
Args:
|
||||
stack: A list of n 5-tuples,
|
||||
(filename, lineno, name, frame_globals, func_start_lineno).
|
||||
include_func_start_lineno: True if function start line number should be
|
||||
included as the 5th entry in return tuples.
|
||||
|
||||
Returns:
|
||||
A list of n 4-tuples or 5-tuples
|
||||
(filename, lineno, name, code, [optional: func_start_lineno]), where the
|
||||
code tuple element is calculated from the corresponding elements of the
|
||||
input tuple.
|
||||
"""
|
||||
ret = []
|
||||
for (filename, lineno, name, frame_globals, func_start_lineno,
|
||||
unused_frame_info) in stack:
|
||||
linecache.checkcache(filename)
|
||||
line = linecache.getline(filename, lineno, frame_globals)
|
||||
if line:
|
||||
line = line.strip()
|
||||
else:
|
||||
line = None
|
||||
if include_func_start_lineno:
|
||||
ret.append((filename, lineno, name, line, func_start_lineno))
|
||||
else:
|
||||
ret.append((filename, lineno, name, line))
|
||||
return ret
|
||||
|
||||
def _extract_stack(self):
|
||||
"""A lightweight, extensible re-implementation of traceback.extract_stack.
|
||||
|
||||
NOTE(mrry): traceback.extract_stack eagerly retrieves the line of code for
|
||||
each stack frame using linecache, which results in an abundance of stat()
|
||||
calls. This implementation does not retrieve the code, and any consumer
|
||||
should apply _convert_stack to the result to obtain a traceback that can
|
||||
be formatted etc. using traceback methods.
|
||||
|
||||
Derived classes can implement _extract_frame_info() to add extra information
|
||||
to the traceback.
|
||||
|
||||
Returns:
|
||||
A list of 6-tuples
|
||||
(filename, lineno, name, frame_globals, func_start_lineno, custom_info)
|
||||
corresponding to the call stack of the current thread.
|
||||
"""
|
||||
try:
|
||||
raise ZeroDivisionError
|
||||
except ZeroDivisionError:
|
||||
f = sys.exc_info()[2].tb_frame.f_back
|
||||
ret = []
|
||||
while f is not None:
|
||||
lineno = f.f_lineno
|
||||
co = f.f_code
|
||||
filename = co.co_filename
|
||||
name = co.co_name
|
||||
frame_globals = f.f_globals
|
||||
func_start_lineno = co.co_firstlineno
|
||||
frame_info = self._extract_frame_info(f)
|
||||
ret.append((filename, lineno, name, frame_globals, func_start_lineno,
|
||||
frame_info))
|
||||
f = f.f_back
|
||||
ret.reverse()
|
||||
return ret
|
||||
|
||||
def _extract_frame_info(self, frame): # pylint: disable=unused-argument
|
||||
"""Extracts custom information from a frame in an op traceback."""
|
||||
return None
|
||||
|
||||
def _check_not_finalized(self):
|
||||
"""Check if the graph is finalized.
|
||||
|
||||
@ -2226,6 +2237,7 @@ class Graph(object):
|
||||
|
||||
@property
|
||||
def graph_def_versions(self):
|
||||
# pylint: disable=line-too-long
|
||||
"""The GraphDef version information of this graph.
|
||||
|
||||
For details on the meaning of each version, see
|
||||
@ -2234,6 +2246,7 @@ class Graph(object):
|
||||
Returns:
|
||||
A `VersionDef`.
|
||||
"""
|
||||
# pylint: enable=line-too-long
|
||||
return self._graph_def_versions
|
||||
|
||||
@property
|
||||
@ -2287,6 +2300,7 @@ class Graph(object):
|
||||
self._control_flow_context = context
|
||||
|
||||
def _as_graph_def(self, from_version=None, add_shapes=False):
|
||||
# pylint: disable=line-too-long
|
||||
"""Returns a serialized `GraphDef` representation of this graph.
|
||||
|
||||
The serialized `GraphDef` can be imported into another `Graph`
|
||||
@ -2312,6 +2326,7 @@ class Graph(object):
|
||||
ValueError: If the `graph_def` would be too large.
|
||||
|
||||
"""
|
||||
# pylint: enable=line-too-long
|
||||
with self._lock:
|
||||
graph = graph_pb2.GraphDef()
|
||||
graph.versions.CopyFrom(self._graph_def_versions)
|
||||
@ -2341,6 +2356,7 @@ class Graph(object):
|
||||
return graph, self._version
|
||||
|
||||
def as_graph_def(self, from_version=None, add_shapes=False):
|
||||
# pylint: disable=line-too-long
|
||||
"""Returns a serialized `GraphDef` representation of this graph.
|
||||
|
||||
The serialized `GraphDef` can be imported into another `Graph`
|
||||
@ -2363,6 +2379,7 @@ class Graph(object):
|
||||
Raises:
|
||||
ValueError: If the `graph_def` would be too large.
|
||||
"""
|
||||
# pylint: enable=line-too-long
|
||||
result, _ = self._as_graph_def(from_version, add_shapes)
|
||||
return result
|
||||
|
||||
@ -2930,7 +2947,7 @@ class Graph(object):
|
||||
finally:
|
||||
self._default_original_op = old_original_op
|
||||
|
||||
# pylint: disable=g-doc-return-or-yield
|
||||
# pylint: disable=g-doc-return-or-yield,line-too-long
|
||||
@tf_contextlib.contextmanager
|
||||
def name_scope(self, name):
|
||||
r"""Returns a context manager that creates hierarchical names for operations.
|
||||
@ -3040,7 +3057,7 @@ class Graph(object):
|
||||
yield "" if new_stack is None else new_stack + "/"
|
||||
finally:
|
||||
self._name_stack = old_stack
|
||||
# pylint: enable=g-doc-return-or-yield
|
||||
# pylint: enable=g-doc-return-or-yield,line-too-long
|
||||
|
||||
def unique_name(self, name, mark_as_used=True):
|
||||
"""Return a unique operation name for `name`.
|
||||
@ -3181,6 +3198,7 @@ class Graph(object):
|
||||
|
||||
@tf_contextlib.contextmanager
|
||||
def device(self, device_name_or_function):
|
||||
# pylint: disable=line-too-long
|
||||
"""Returns a context manager that specifies the default device to use.
|
||||
|
||||
The `device_name_or_function` argument may either be a device name
|
||||
@ -3237,6 +3255,7 @@ class Graph(object):
|
||||
created ops.
|
||||
|
||||
"""
|
||||
# pylint: enable=line-too-long
|
||||
if (device_name_or_function is not None
|
||||
and not callable(device_name_or_function)):
|
||||
device_function = pydev.merge_device(device_name_or_function)
|
||||
|
||||
Loading…
Reference in New Issue
Block a user