135 lines
5.0 KiB
Python
135 lines
5.0 KiB
Python
![]() |
# Copyright 2020 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.
|
||
|
# ==============================================================================
|
||
|
"""Memory leak detection utility."""
|
||
|
|
||
|
from tensorflow.python.framework.python_memory_checker import _PythonMemoryChecker
|
||
|
from tensorflow.python.profiler import trace
|
||
|
from tensorflow.python.util import tf_inspect
|
||
|
|
||
|
try:
|
||
|
from tensorflow.python.platform.cpp_memory_checker import _CppMemoryChecker as CppMemoryChecker # pylint:disable=g-import-not-at-top
|
||
|
except ImportError:
|
||
|
CppMemoryChecker = None
|
||
|
|
||
|
|
||
|
def _get_test_name_best_effort():
|
||
|
"""If available, return the current test name. Otherwise, `None`."""
|
||
|
for stack in tf_inspect.stack():
|
||
|
function_name = stack[3]
|
||
|
if function_name.startswith('test'):
|
||
|
try:
|
||
|
class_name = stack[0].f_locals['self'].__class__.__name__
|
||
|
return class_name + '.' + function_name
|
||
|
except: # pylint:disable=bare-except
|
||
|
pass
|
||
|
|
||
|
return None
|
||
|
|
||
|
|
||
|
# TODO(kkb): Also create decorator versions for convenience.
|
||
|
class MemoryChecker(object):
|
||
|
"""Memory leak detection class.
|
||
|
|
||
|
This is a utility class to detect Python and C++ memory leaks. It's intended
|
||
|
for both testing and debugging. Basic usage:
|
||
|
|
||
|
>>> # MemoryChecker() context manager tracks memory status inside its scope.
|
||
|
>>> with MemoryChecker() as memory_checker:
|
||
|
>>> tensors = []
|
||
|
>>> for _ in range(10):
|
||
|
>>> # Simulating `tf.constant(1)` object leak every iteration.
|
||
|
>>> tensors.append(tf.constant(1))
|
||
|
>>>
|
||
|
>>> # Take a memory snapshot for later analysis.
|
||
|
>>> memory_checker.record_snapshot()
|
||
|
>>>
|
||
|
>>> # `report()` generates a html graph file showing allocations over
|
||
|
>>> # snapshots per every stack trace.
|
||
|
>>> memory_checker.report()
|
||
|
>>>
|
||
|
>>> # This assertion will detect `tf.constant(1)` object leak.
|
||
|
>>> memory_checker.assert_no_leak_if_all_possibly_except_one()
|
||
|
|
||
|
`record_snapshot()` must be called once every iteration at the same location.
|
||
|
This is because the detection algorithm relies on the assumption that if there
|
||
|
is a leak, it's happening similarly on every snapshot.
|
||
|
"""
|
||
|
|
||
|
@trace.trace_wrapper
|
||
|
def __enter__(self):
|
||
|
self._python_memory_checker = _PythonMemoryChecker()
|
||
|
if CppMemoryChecker:
|
||
|
self._cpp_memory_checker = CppMemoryChecker(_get_test_name_best_effort())
|
||
|
return self
|
||
|
|
||
|
@trace.trace_wrapper
|
||
|
def __exit__(self, exc_type, exc_value, traceback):
|
||
|
if CppMemoryChecker:
|
||
|
self._cpp_memory_checker.stop()
|
||
|
|
||
|
# We do not enable trace_wrapper on this function to avoid contaminating
|
||
|
# the snapshot.
|
||
|
def record_snapshot(self):
|
||
|
"""Take a memory snapshot for later analysis.
|
||
|
|
||
|
`record_snapshot()` must be called once every iteration at the same
|
||
|
location. This is because the detection algorithm relies on the assumption
|
||
|
that if there is a leak, it's happening similarly on every snapshot.
|
||
|
|
||
|
The recommended number of `record_snapshot()` call depends on the testing
|
||
|
code complexity and the allcoation pattern.
|
||
|
"""
|
||
|
self._python_memory_checker.record_snapshot()
|
||
|
if CppMemoryChecker:
|
||
|
self._cpp_memory_checker.record_snapshot()
|
||
|
|
||
|
@trace.trace_wrapper
|
||
|
def report(self):
|
||
|
"""Generates a html graph file showing allocations over snapshots.
|
||
|
|
||
|
It create a temporary directory and put all the output files there.
|
||
|
If this is running under Google internal testing infra, it will use the
|
||
|
directory provided the infra instead.
|
||
|
"""
|
||
|
self._python_memory_checker.report()
|
||
|
if CppMemoryChecker:
|
||
|
self._cpp_memory_checker.report()
|
||
|
|
||
|
@trace.trace_wrapper
|
||
|
def assert_no_leak_if_all_possibly_except_one(self):
|
||
|
"""Raises an exception if a leak is detected.
|
||
|
|
||
|
This algorithm classifies a series of allocations as a leak if it's the same
|
||
|
type(Python) or it happens at the same stack trace(C++) at every snapshot,
|
||
|
but possibly except one snapshot.
|
||
|
"""
|
||
|
|
||
|
self._python_memory_checker.assert_no_leak_if_all_possibly_except_one()
|
||
|
if CppMemoryChecker:
|
||
|
self._cpp_memory_checker.assert_no_leak_if_all_possibly_except_one()
|
||
|
|
||
|
@trace.trace_wrapper
|
||
|
def assert_no_new_python_objects(self, threshold=None):
|
||
|
"""Raises an exception if there are new Python objects created.
|
||
|
|
||
|
It computes the number of new Python objects per type using the first and
|
||
|
the last snapshots.
|
||
|
|
||
|
Args:
|
||
|
threshold: A dictionary of [Type name string], [count] pair. It won't
|
||
|
raise an exception if the new Python objects are under this threshold.
|
||
|
"""
|
||
|
self._python_memory_checker.assert_no_new_objects(threshold=threshold)
|