1# Copyright 2015 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"""Functions used to extract and analyze stacks.  Faster than Python libs."""
16# pylint: disable=g-bad-name
17from __future__ import absolute_import
18from __future__ import division
19from __future__ import print_function
20
21import collections
22import inspect
23import threading
24
25import six
26
27# TODO(b/138203821): change to from ...util import ... once the bug is fixed.
28from tensorflow.python.util import _tf_stack
29
30# Generally such lookups should be done using `threading.local()`. See
31# https://blogs.gnome.org/jamesh/2008/06/11/tls-python/ for a detailed
32# explanation of why. However the transform stacks are expected to be empty
33# when a thread is joined, so reusing the key does not introduce a correctness
34# issue. Moreover, get_ident is faster than storing and retrieving a unique
35# key in a thread local store.
36if six.PY2:
37  import thread  # pylint: disable=g-import-not-at-top
38  _get_thread_key = thread.get_ident
39else:
40  _get_thread_key = threading.get_ident
41
42
43# TODO(mdan): Move these to C++ as well.
44# Moving to C++ can further avoid extra copies made by get_effective_map.
45_source_mapper_stacks = collections.defaultdict(lambda: [SentinelMapper()])
46_source_filter_stacks = collections.defaultdict(lambda: [SentinelFilter()])
47
48
49class StackTraceTransform(object):
50  """Base class for stack trace transformation functions."""
51
52  _stack_dict = None  # Subclasses should override
53  _thread_key = None
54
55  def __enter__(self):
56    # Any given instance is assumed to be used by a single thread, which reduces
57    # expensive thread local lookups.
58    if self._thread_key is None:
59      self._thread_key = _get_thread_key()
60    else:
61      assert self._thread_key == _get_thread_key(), 'Shared across threads?'
62
63    stack = self._stack_dict[self._thread_key]
64    self.parent = stack[-1]
65    stack.append(self)
66    self.update()
67    return self
68
69  def __exit__(self, unused_type, unused_value, unused_traceback):
70    top = self._stack_dict[self._thread_key].pop()
71    assert top is self, 'Concurrent access?'
72
73  def update(self):
74    raise NotImplementedError('subclasses need to override this')
75
76
77class StackTraceMapper(StackTraceTransform):
78  """Allows remapping traceback information to different source code."""
79  _stack_dict = _source_mapper_stacks
80
81  def __init__(self):
82    self.internal_map = _tf_stack.PyBindSourceMap()
83
84  def update(self):
85    self.internal_map.update_to(tuple(self.get_effective_source_map().items()))
86
87  def get_effective_source_map(self):
88    """Returns a map (filename, lineno) -> (filename, lineno, function_name)."""
89    raise NotImplementedError('subclasses need to override this')
90
91
92EMPTY_DICT = {}
93
94
95class SentinelMapper(StackTraceMapper):
96
97  def get_effective_source_map(self):
98    return EMPTY_DICT
99
100
101class StackTraceFilter(StackTraceTransform):
102  """Allows filtering traceback information by removing superfluous frames."""
103  _stack_dict = _source_filter_stacks
104
105  def __init__(self):
106    self.internal_set = _tf_stack.PyBindFileSet()
107
108  def update(self):
109    self.internal_set.update_to(set(self.get_filtered_filenames()))
110
111  def get_filtered_filenames(self):
112    raise NotImplementedError('subclasses need to override this')
113
114
115EMPTY_SET = frozenset()
116
117
118class SentinelFilter(StackTraceFilter):
119
120  def get_filtered_filenames(self):
121    return EMPTY_SET
122
123
124class CurrentModuleFilter(StackTraceFilter):
125  """Filters stack frames from the module where this is used (best effort)."""
126
127  def __init__(self):
128    super().__init__()
129    filter_filename = None
130    outer_f = None
131    f = inspect.currentframe()
132    try:
133      if f is not None:
134        # The current frame is __init__. The first outer frame should be the
135        # caller.
136        outer_f = f.f_back
137        if outer_f is not None:
138          filter_filename = inspect.getsourcefile(outer_f)
139      self._filename = filter_filename
140      # This may be called repeatedly: once on entry by the superclass, then by
141      # each child context manager.
142      self._cached_set = None
143    finally:
144      # Avoid reference cycles, see:
145      # https://docs.python.org/3.7/library/inspect.html#the-interpreter-stack
146      del f
147      del outer_f
148
149  def get_filtered_filenames(self):
150    if self._cached_set is not None:
151      return self._cached_set
152
153    filtered_filenames = frozenset((self._filename,))
154    if self.parent is not None:
155      filtered_filenames |= self.parent.get_filtered_filenames()
156    self._cached_set = filtered_filenames
157    return filtered_filenames
158
159
160def extract_stack():
161  """An eager-friendly alternative to traceback.extract_stack.
162
163  Returns:
164    A list-like FrameSummary containing StackFrame-like objects, which are
165    namedtuple-like objects with the following fields: filename, lineno, name,
166    line, meant to masquerade as traceback.FrameSummary objects.
167  """
168  # N.B ExtractStack in tf_stack.cc will drop this frame prior to
169  # traversing the stack.
170  # TODO(cheshire): Remove this function, use extract_stack_for_node or Python
171  # traceback module.
172  thread_key = _get_thread_key()
173  return _tf_stack.extract_stack(
174      _source_mapper_stacks[thread_key][-1].internal_map,
175      _source_filter_stacks[thread_key][-1].internal_set)
176
177
178# TODO(mdan): Revisit these - a single location is almost always sufficient.
179def extract_stack_for_node(node):
180  """Attaches the current stack trace to `node`.
181
182  Args:
183    node: a Node object.
184
185  Returns:
186    A list-like FrameSummary containing StackFrame-like objects, which are
187    namedtuple-like objects with the following fields: filename, lineno, name,
188    line, meant to masquerade as traceback.FrameSummary objects.
189  """
190  # N.B ExtractStack in tf_stack.cc will drop this frame prior to
191  # traversing the stack.
192  thread_key = _get_thread_key()
193  return _tf_stack.extract_stack_for_node(
194      _source_mapper_stacks[thread_key][-1].internal_map,
195      _source_filter_stacks[thread_key][-1].internal_set, node)
196
197
198StackSummary = _tf_stack.StackTraceWrapper
199FrameSummary = _tf_stack.StackFrame
200