wehub-resource-sync
289 lines
9.0 KiB
Python
289 lines
9.0 KiB
Python
import gc
|
|
import os
|
|
import re
|
|
import time
|
|
import tracemalloc
|
|
from collections import defaultdict, namedtuple
|
|
from typing import Callable, List, Optional
|
|
|
|
from ray.util.annotations import DeveloperAPI
|
|
|
|
_logged = set()
|
|
_disabled = False
|
|
_periodic_log = False
|
|
_last_logged = 0.0
|
|
|
|
|
|
@DeveloperAPI
|
|
def log_once(key: str) -> bool:
|
|
"""Returns True if this is the "first" call for a given key.
|
|
|
|
Various logging settings can adjust the definition of "first".
|
|
|
|
Args:
|
|
key: A unique identifier for the call site.
|
|
|
|
Returns:
|
|
True if this is the first call for ``key`` (subject to the current
|
|
``log_once`` settings), False otherwise.
|
|
|
|
Example:
|
|
|
|
.. testcode::
|
|
|
|
import logging
|
|
from ray.util.debug import log_once
|
|
|
|
logger = logging.getLogger(__name__)
|
|
if log_once("some_key"):
|
|
logger.info("Some verbose logging statement")
|
|
"""
|
|
|
|
global _last_logged
|
|
|
|
if _disabled:
|
|
return False
|
|
elif key not in _logged:
|
|
_logged.add(key)
|
|
_last_logged = time.time()
|
|
return True
|
|
elif _periodic_log and time.time() - _last_logged > 60.0:
|
|
_logged.clear()
|
|
_last_logged = time.time()
|
|
return False
|
|
else:
|
|
return False
|
|
|
|
|
|
@DeveloperAPI
|
|
def disable_log_once_globally():
|
|
"""Make log_once() return False in this process."""
|
|
|
|
global _disabled
|
|
_disabled = True
|
|
|
|
|
|
@DeveloperAPI
|
|
def enable_periodic_logging():
|
|
"""Make log_once() periodically return True in this process."""
|
|
|
|
global _periodic_log
|
|
_periodic_log = True
|
|
|
|
|
|
@DeveloperAPI
|
|
def reset_log_once(key: Optional[str] = None):
|
|
"""Resets log_once for the provided key.
|
|
|
|
If you don't provide a key, resets log_once for all keys.
|
|
"""
|
|
if key is None:
|
|
_logged.clear()
|
|
else:
|
|
_logged.discard(key)
|
|
|
|
|
|
# A suspicious memory-allocating stack-trace that we should re-test
|
|
# to make sure it's not a false positive.
|
|
Suspect = DeveloperAPI(
|
|
namedtuple(
|
|
"Suspect",
|
|
[
|
|
# The stack trace of the allocation, going back n frames, depending
|
|
# on the tracemalloc.start(n) call.
|
|
"traceback",
|
|
# The amount of memory taken by this particular stack trace
|
|
# over the course of the experiment.
|
|
"memory_increase",
|
|
# The slope of the scipy linear regression (x=iteration; y=memory size).
|
|
"slope",
|
|
# The rvalue of the scipy linear regression.
|
|
"rvalue",
|
|
# The memory size history (list of all memory sizes over all iterations).
|
|
"hist",
|
|
],
|
|
)
|
|
)
|
|
|
|
|
|
def _test_some_code_for_memory_leaks(
|
|
desc: str,
|
|
init: Optional[Callable[[], None]],
|
|
code: Callable[[], None],
|
|
repeats: int,
|
|
max_num_trials: int = 1,
|
|
min_memory_increase: int = 0,
|
|
) -> List[Suspect]:
|
|
"""Runs given code (and init code) n times and checks for memory leaks.
|
|
|
|
Args:
|
|
desc: A descriptor of the test.
|
|
init: Optional code to be executed initially.
|
|
code: The actual code to be checked for producing memory leaks.
|
|
repeats: How many times to repeatedly execute `code`.
|
|
max_num_trials: The maximum number of trials to run. A new trial is only
|
|
run, if the previous one produced a memory leak. For all non-1st trials,
|
|
`repeats` calculates as: actual_repeats = `repeats` * (trial + 1), where
|
|
the first trial is 0.
|
|
min_memory_increase: Minimum total memory increase in bytes for a suspect
|
|
to be reported. Use this to ignore small allocations from internal
|
|
caches and allocator artifacts that aren't real leaks.
|
|
|
|
Returns:
|
|
A list of Suspect objects, describing possible memory leaks. If list
|
|
is empty, no leaks have been found.
|
|
"""
|
|
|
|
def _i_print(i):
|
|
if (i + 1) % 10 == 0:
|
|
print(".", end="" if (i + 1) % 100 else f" {i + 1}\n", flush=True)
|
|
|
|
# Do n trials to make sure a found leak is really one.
|
|
suspicious = set()
|
|
suspicious_stats = []
|
|
for trial in range(max_num_trials):
|
|
# Store up to n frames of each call stack.
|
|
tracemalloc.start(20)
|
|
|
|
table = defaultdict(list)
|
|
|
|
# Repeat running code for n times.
|
|
# Increase repeat value with each trial to make sure stats are more
|
|
# solid each time (avoiding false positives).
|
|
actual_repeats = repeats * (trial + 1)
|
|
|
|
print(f"{desc} {actual_repeats} times.")
|
|
|
|
# Initialize if necessary.
|
|
if init is not None:
|
|
init()
|
|
# Run `code` n times, each time taking a memory snapshot.
|
|
for i in range(actual_repeats):
|
|
_i_print(i)
|
|
# Manually trigger garbage collection before and after code runs in order to
|
|
# make tracemalloc snapshots as accurate as possible.
|
|
gc.collect()
|
|
code()
|
|
gc.collect()
|
|
_take_snapshot(table, suspicious)
|
|
print("\n")
|
|
|
|
# Check, which traces have moved up in their memory consumption
|
|
# constantly over time.
|
|
suspicious.clear()
|
|
suspicious_stats.clear()
|
|
# Suspicious memory allocation found?
|
|
suspects = _find_memory_leaks_in_table(table)
|
|
if min_memory_increase > 0:
|
|
suspects = [s for s in suspects if s.memory_increase >= min_memory_increase]
|
|
for suspect in sorted(suspects, key=lambda s: s.memory_increase, reverse=True):
|
|
# Only print out the biggest offender:
|
|
if len(suspicious) == 0:
|
|
_pprint_suspect(suspect)
|
|
print("-> added to retry list")
|
|
suspicious.add(suspect.traceback)
|
|
suspicious_stats.append(suspect)
|
|
|
|
tracemalloc.stop()
|
|
|
|
# Some suspicious memory allocations found.
|
|
if len(suspicious) > 0:
|
|
print(f"{len(suspicious)} suspects found. Top-ten:")
|
|
for i, s in enumerate(suspicious_stats):
|
|
if i > 10:
|
|
break
|
|
print(
|
|
f"{i}) line={s.traceback[-1]} mem-increase={s.memory_increase}B "
|
|
f"slope={s.slope}B/detection rval={s.rvalue}"
|
|
)
|
|
# Nothing suspicious found -> Exit trial loop and return.
|
|
else:
|
|
print("No remaining suspects found -> returning")
|
|
break
|
|
|
|
# Print out final top offender.
|
|
if len(suspicious_stats) > 0:
|
|
_pprint_suspect(suspicious_stats[0])
|
|
|
|
return suspicious_stats
|
|
|
|
|
|
def _take_snapshot(table, suspicious=None):
|
|
# Take a memory snapshot.
|
|
snapshot = tracemalloc.take_snapshot()
|
|
# Group all memory allocations by their stacktrace (going n frames
|
|
# deep as defined above in tracemalloc.start(n)).
|
|
# Then sort groups by size, then count, then trace.
|
|
top_stats = snapshot.statistics("traceback")
|
|
|
|
# For the first m largest increases, keep only, if a) first trial or b) those
|
|
# that are already in the `suspicious` set.
|
|
for stat in top_stats[:100]:
|
|
if not suspicious or stat.traceback in suspicious:
|
|
table[stat.traceback].append(stat.size)
|
|
|
|
|
|
def _find_memory_leaks_in_table(table):
|
|
import numpy as np
|
|
import scipy.stats
|
|
|
|
suspects = []
|
|
|
|
for traceback, hist in table.items():
|
|
# Do a quick mem increase check.
|
|
memory_increase = hist[-1] - hist[0]
|
|
|
|
# Only if memory increased, do we check further.
|
|
if memory_increase <= 0.0:
|
|
continue
|
|
|
|
# Ignore this very module here (we are collecting lots of data
|
|
# so an increase is expected).
|
|
top_stack = str(traceback[-1])
|
|
drive_separator = "\\\\" if os.name == "nt" else "/"
|
|
if any(
|
|
s in top_stack
|
|
for s in [
|
|
"tracemalloc",
|
|
"pycharm",
|
|
"thirdparty_files/psutil",
|
|
re.sub("\\.", drive_separator, __name__) + ".py",
|
|
]
|
|
):
|
|
continue
|
|
|
|
# Do a linear regression to get the slope and R-value.
|
|
line = scipy.stats.linregress(x=np.arange(len(hist)), y=np.array(hist))
|
|
|
|
# - If weak positive slope and some confidence and
|
|
# increase > n bytes -> error.
|
|
# - If stronger positive slope -> error.
|
|
if memory_increase > 1000 and (
|
|
(line.slope > 60.0 and line.rvalue > 0.875)
|
|
or (line.slope > 20.0 and line.rvalue > 0.9)
|
|
or (line.slope > 10.0 and line.rvalue > 0.95)
|
|
):
|
|
suspects.append(
|
|
Suspect(
|
|
traceback=traceback,
|
|
memory_increase=memory_increase,
|
|
slope=line.slope,
|
|
rvalue=line.rvalue,
|
|
hist=hist,
|
|
)
|
|
)
|
|
|
|
return suspects
|
|
|
|
|
|
def _pprint_suspect(suspect):
|
|
print(
|
|
"Most suspicious memory allocation in traceback "
|
|
"(only printing out this one, but all (less suspicious)"
|
|
" suspects will be investigated as well):"
|
|
)
|
|
print("\n".join(suspect.traceback.format()))
|
|
print(f"Increase total={suspect.memory_increase}B")
|
|
print(f"Slope={suspect.slope} B/detection")
|
|
print(f"Rval={suspect.rvalue}")
|