Skip to content

Commit

Permalink
Merge branch 'staging' into warmfix/add-thresholds-for-unstaking
Browse files Browse the repository at this point in the history
  • Loading branch information
thewhaleking authored May 1, 2024
2 parents 2fe9cb0 + 5412af0 commit 5ca6534
Show file tree
Hide file tree
Showing 2 changed files with 60 additions and 12 deletions.
23 changes: 20 additions & 3 deletions bittensor/btlogging/format.py
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,18 @@ def success(self, message, *args, **kws):
for level, color in log_level_color_prefix.items()
}

DEFAULT_LOG_FORMAT = (
f"{Fore.BLUE}%(asctime)s{Fore.RESET} | "
f"{Style.BRIGHT}{Fore.WHITE}%(levelname)s{Style.RESET_ALL} | "
f"%(name)s:%(filename)s:%(lineno)s | %(message)s"
)

DEFAULT_TRACE_FORMAT = (
f"{Fore.BLUE}%(asctime)s{Fore.RESET} | "
f"{Style.BRIGHT}{Fore.WHITE}%(levelname)s{Style.RESET_ALL} | "
f"%(name)s:%(filename)s:%(lineno)s | %(message)s"
)


class BtStreamFormatter(logging.Formatter):
def __init__(self, *args, **kwargs):
Expand All @@ -93,10 +105,15 @@ def format(self, record):
format_orig = self._style._fmt
record.levelname = f"{record.levelname:^16}"

if self.trace is True:
self._style._fmt = LOG_TRACE_FORMATS[record.levelno]
if record.levelno not in LOG_FORMATS:
self._style._fmt = (
DEFAULT_TRACE_FORMAT if self.trace else DEFAULT_LOG_FORMAT
)
else:
self._style._fmt = LOG_FORMATS[record.levelno]
if self.trace is True:
self._style._fmt = LOG_TRACE_FORMATS[record.levelno]
else:
self._style._fmt = LOG_FORMATS[record.levelno]

for text, emoji in emoji_map.items():
record.msg = record.msg.replace(text, emoji)
Expand Down
49 changes: 40 additions & 9 deletions bittensor/btlogging/loggingmachine.py
Original file line number Diff line number Diff line change
Expand Up @@ -69,7 +69,7 @@ def __init__(self, config: bittensor.config, name: str = BITTENSOR_LOGGER_NAME):
# basics
super(LoggingMachine, self).__init__()
self._queue = mp.Queue(-1)
self._name = name
self._primary_loggers = {name}
self._config = config

# Formatters
Expand All @@ -89,7 +89,7 @@ def __init__(self, config: bittensor.config, name: str = BITTENSOR_LOGGER_NAME):
self._listener = self._create_and_start_listener(self._handlers)

# set up all the loggers
self._logger = self._initialize_bt_logger(name, config)
self._logger = self._initialize_bt_logger(name)
self.disable_third_party_loggers()

def _configure_handlers(self, config) -> list[stdlogging.Handler]:
Expand Down Expand Up @@ -118,9 +118,8 @@ def set_config(self, config):
"""
self._config = config
if config.logging_dir and config.record_log:
logfile = os.path.abspath(
os.path.join(config.logging_dir, DEFAULT_LOG_FILE_NAME)
)
expanded_dir = os.path.expanduser(config.logging_dir)
logfile = os.path.abspath(os.path.join(expanded_dir, DEFAULT_LOG_FILE_NAME))
self._enable_file_logging(logfile)
if config.trace:
self.enable_trace()
Expand Down Expand Up @@ -149,7 +148,7 @@ def get_queue(self):
"""
return self._queue

def _initialize_bt_logger(self, name, config):
def _initialize_bt_logger(self, name):
"""
Initialize logging for bittensor.
Expand All @@ -162,6 +161,17 @@ def _initialize_bt_logger(self, name, config):
logger.addHandler(queue_handler)
return logger

def _deinitialize_bt_logger(self, name):
"""
Find the logger by name and remove the
queue handler associated with it.
"""
logger = stdlogging.getLogger(name)
for handler in list(logger.handlers):
if isinstance(handler, QueueHandler):
logger.removeHandler(handler)
return logger

def _create_file_handler(self, logfile: str):
file_handler = RotatingFileHandler(
logfile,
Expand All @@ -172,9 +182,30 @@ def _create_file_handler(self, logfile: str):
file_handler.setLevel(stdlogging.TRACE)
return file_handler

def register_primary_logger(self, name: str):
"""
Register a logger as primary logger
This adds a logger to the _primary_loggers set to ensure
it doesn't get disabled when disabling third-party loggers.
A queue handler is also associated with it.
"""
self._primary_loggers.add(name)
self._initialize_bt_logger(name)

def deregister_primary_logger(self, name: str):
"""
De-registers a primary logger
This function removes the logger from the _primary_loggers
set and deinitializes its queue handler
"""
self._primary_loggers.remove(name)
self._deinitialize_bt_logger(name)

def enable_third_party_loggers(self):
for logger in all_loggers():
if logger.name == self._name:
if logger.name in self._primary_loggers:
continue
queue_handler = QueueHandler(self._queue)
logger.addHandler(queue_handler)
Expand All @@ -183,7 +214,7 @@ def enable_third_party_loggers(self):
def disable_third_party_loggers(self):
# remove all handlers
for logger in all_loggers():
if logger.name == self._name:
if logger.name in self._primary_loggers:
continue
for handler in logger.handlers:
logger.removeHandler(handler)
Expand Down Expand Up @@ -211,7 +242,7 @@ def before_enable_default(self):
self._logger.info(f"Enabling default logging.")
self._logger.setLevel(stdlogging.INFO)
for logger in all_loggers():
if logger.name == self._name:
if logger.name in self._primary_loggers:
continue
logger.setLevel(stdlogging.CRITICAL)

Expand Down

0 comments on commit 5ca6534

Please sign in to comment.