123 lines
3.6 KiB
Python
123 lines
3.6 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.
|
||
|
|
# ==============================================================================
|
||
|
|
"""Utilities for measuring elapsed time."""
|
||
|
|
|
||
|
|
import contextlib
|
||
|
|
import logging
|
||
|
|
import threading
|
||
|
|
import time
|
||
|
|
|
||
|
|
from tensorboard.util import tb_logging
|
||
|
|
|
||
|
|
logger = tb_logging.get_logger()
|
||
|
|
|
||
|
|
|
||
|
|
def log_latency(region_name_or_function_to_decorate, log_level=None):
|
||
|
|
"""Log latency in a function or region.
|
||
|
|
|
||
|
|
Three usages are supported. As a decorator:
|
||
|
|
|
||
|
|
>>> @log_latency
|
||
|
|
... def function_1():
|
||
|
|
... pass
|
||
|
|
...
|
||
|
|
|
||
|
|
|
||
|
|
As a decorator with a custom label for the region:
|
||
|
|
|
||
|
|
>>> @log_latency("custom_label")
|
||
|
|
... def function_2():
|
||
|
|
... pass
|
||
|
|
...
|
||
|
|
|
||
|
|
As a context manager:
|
||
|
|
|
||
|
|
>>> def function_3():
|
||
|
|
... with log_latency("region_within_function"):
|
||
|
|
... pass
|
||
|
|
...
|
||
|
|
|
||
|
|
Args:
|
||
|
|
region_name_or_function_to_decorate: Either: a `str`, in which
|
||
|
|
case the result of this function may be used as either a
|
||
|
|
decorator or a context manager; or a callable, in which case
|
||
|
|
the result of this function is a decorated version of that
|
||
|
|
callable.
|
||
|
|
log_level: Optional integer logging level constant. Defaults to
|
||
|
|
`logging.INFO`.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
A decorated version of the input callable, or a dual
|
||
|
|
decorator/context manager with the input region name.
|
||
|
|
"""
|
||
|
|
|
||
|
|
if log_level is None:
|
||
|
|
log_level = logging.INFO
|
||
|
|
|
||
|
|
if isinstance(region_name_or_function_to_decorate, str):
|
||
|
|
region_name = region_name_or_function_to_decorate
|
||
|
|
return _log_latency(region_name, log_level)
|
||
|
|
else:
|
||
|
|
function_to_decorate = region_name_or_function_to_decorate
|
||
|
|
qualname = getattr(function_to_decorate, "__qualname__", None)
|
||
|
|
if qualname is None:
|
||
|
|
qualname = str(function_to_decorate)
|
||
|
|
decorator = _log_latency(qualname, log_level)
|
||
|
|
return decorator(function_to_decorate)
|
||
|
|
|
||
|
|
|
||
|
|
class _ThreadLocalStore(threading.local):
|
||
|
|
def __init__(self):
|
||
|
|
self.nesting_level = 0
|
||
|
|
|
||
|
|
|
||
|
|
_store = _ThreadLocalStore()
|
||
|
|
|
||
|
|
|
||
|
|
@contextlib.contextmanager
|
||
|
|
def _log_latency(name, log_level):
|
||
|
|
if not logger.isEnabledFor(log_level):
|
||
|
|
yield
|
||
|
|
return
|
||
|
|
|
||
|
|
start_level = _store.nesting_level
|
||
|
|
try:
|
||
|
|
started = time.time()
|
||
|
|
_store.nesting_level = start_level + 1
|
||
|
|
indent = (" " * 2) * start_level
|
||
|
|
thread = threading.current_thread()
|
||
|
|
prefix = "%s[%x]%s" % (thread.name, thread.ident, indent)
|
||
|
|
_log(log_level, "%s ENTER %s", prefix, name)
|
||
|
|
yield
|
||
|
|
finally:
|
||
|
|
_store.nesting_level = start_level
|
||
|
|
elapsed = time.time() - started
|
||
|
|
_log(
|
||
|
|
log_level,
|
||
|
|
"%s LEAVE %s - %0.6fs elapsed",
|
||
|
|
prefix,
|
||
|
|
name,
|
||
|
|
elapsed,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _log(log_level, msg, *args):
|
||
|
|
# Forwarding method to ensure that all logging statements
|
||
|
|
# originating in this module have the same line number; if the
|
||
|
|
# "ENTER" log is on a line with 2-digit number and the "LEAVE" log
|
||
|
|
# is on a line with 3-digit number, the logs are misaligned and
|
||
|
|
# harder to read.
|
||
|
|
logger.log(log_level, msg, *args)
|