chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,3 @@
|
||||
from .logging import LoggingManager
|
||||
|
||||
__all__ = ["LoggingManager"]
|
||||
@@ -0,0 +1,340 @@
|
||||
import logging.config
|
||||
import os
|
||||
from enum import Enum
|
||||
from typing import Optional, Union
|
||||
|
||||
import ray
|
||||
from ray._common.filters import CoreContextFilter
|
||||
from ray._common.formatters import JSONFormatter
|
||||
from ray._private.log import PlainRayHandler
|
||||
from ray.train.v2._internal.execution.context import TrainContext, TrainRunContext
|
||||
from ray.train.v2._internal.util import get_module_name
|
||||
|
||||
|
||||
class TrainContextFilter(logging.Filter):
|
||||
"""Add Ray Train metadata to the log records.
|
||||
|
||||
This filter is applied to Ray Train controller and worker processes.
|
||||
"""
|
||||
|
||||
# Log keys for Ray Train controller and worker processes.
|
||||
class LogKey(str, Enum):
|
||||
RUN_NAME = "run_name"
|
||||
COMPONENT = "component"
|
||||
WORLD_RANK = "world_rank"
|
||||
LOCAL_RANK = "local_rank"
|
||||
NODE_RANK = "node_rank"
|
||||
|
||||
# Ray Train Component by process types
|
||||
class TrainComponent(str, Enum):
|
||||
CONTROLLER = "controller"
|
||||
WORKER = "worker"
|
||||
|
||||
def __init__(self, context: Union[TrainRunContext, TrainContext]):
|
||||
self._is_worker: bool = isinstance(context, TrainContext)
|
||||
if self._is_worker:
|
||||
self._run_name: str = context.train_run_context.get_run_config().name
|
||||
self._world_rank: int = context.get_world_rank()
|
||||
self._local_rank: int = context.get_local_rank()
|
||||
self._node_rank: int = context.get_node_rank()
|
||||
self._component: str = TrainContextFilter.TrainComponent.WORKER
|
||||
else:
|
||||
self._run_name: str = context.get_run_config().name
|
||||
self._component: str = TrainContextFilter.TrainComponent.CONTROLLER
|
||||
|
||||
def controller_filter(self, record):
|
||||
# Add the run_id and component to Ray Train controller processes.
|
||||
setattr(record, TrainContextFilter.LogKey.RUN_NAME, self._run_name)
|
||||
setattr(record, TrainContextFilter.LogKey.COMPONENT, self._component)
|
||||
return True
|
||||
|
||||
def worker_filter(self, record):
|
||||
# Add the run_id and component to Ray Train worker processes.
|
||||
setattr(record, TrainContextFilter.LogKey.RUN_NAME, self._run_name)
|
||||
setattr(record, TrainContextFilter.LogKey.COMPONENT, self._component)
|
||||
# Add all the rank related information to the log record for worker processes.
|
||||
setattr(record, TrainContextFilter.LogKey.WORLD_RANK, self._world_rank)
|
||||
setattr(record, TrainContextFilter.LogKey.LOCAL_RANK, self._local_rank)
|
||||
setattr(record, TrainContextFilter.LogKey.NODE_RANK, self._node_rank)
|
||||
return True
|
||||
|
||||
def filter(self, record):
|
||||
if self._is_worker:
|
||||
return self.worker_filter(record)
|
||||
else:
|
||||
return self.controller_filter(record)
|
||||
|
||||
|
||||
class TrainLogLevelFilter(logging.Filter):
|
||||
"""Filter that applies log level filtering only to ray.train log records."""
|
||||
|
||||
def __init__(self, log_level: str = "INFO"):
|
||||
super().__init__()
|
||||
self._log_level = getattr(logging, log_level)
|
||||
|
||||
def filter(self, record):
|
||||
if record.name == "ray.train" or record.name.startswith("ray.train."):
|
||||
return record.levelno >= self._log_level
|
||||
return True
|
||||
|
||||
|
||||
class SessionFileHandler(logging.Handler):
|
||||
"""A handler that writes to a log file in the Ray session directory.
|
||||
|
||||
The Ray session directory isn't available until Ray is initialized, so any logs
|
||||
emitted before Ray is initialized will be lost.
|
||||
This handler will not create the file handler until you emit a log record.
|
||||
|
||||
Args:
|
||||
filename: The name of the log file. The file is created in the 'logs/train'
|
||||
directory of the Ray session directory.
|
||||
"""
|
||||
|
||||
# TODO (hpguo): This handler class is shared by both Ray Train and ray data. We
|
||||
# should move this to ray core and make it available to both libraries.
|
||||
|
||||
def __init__(self, filename: str):
|
||||
super().__init__()
|
||||
self._filename = filename
|
||||
self._handler = None
|
||||
self._formatter = None
|
||||
self._path = None
|
||||
|
||||
def emit(self, record):
|
||||
if self._handler is None:
|
||||
self._try_create_handler()
|
||||
if self._handler is not None:
|
||||
self._handler.emit(record)
|
||||
|
||||
def setFormatter(self, fmt: logging.Formatter) -> None:
|
||||
if self._handler is not None:
|
||||
self._handler.setFormatter(fmt)
|
||||
self._formatter = fmt
|
||||
|
||||
def get_log_file_path(self) -> Optional[str]:
|
||||
if self._handler is None:
|
||||
self._try_create_handler()
|
||||
return self._path
|
||||
|
||||
def _try_create_handler(self):
|
||||
assert self._handler is None
|
||||
|
||||
# Get the Ray Train log directory. If not in a Ray session, return.
|
||||
# This handler will only be created within a Ray session.
|
||||
log_directory = LoggingManager.get_log_directory()
|
||||
if log_directory is None:
|
||||
return
|
||||
|
||||
os.makedirs(log_directory, exist_ok=True)
|
||||
|
||||
# Create the log file.
|
||||
self._path = os.path.join(log_directory, self._filename)
|
||||
self._handler = logging.FileHandler(self._path)
|
||||
if self._formatter is not None:
|
||||
self._handler.setFormatter(self._formatter)
|
||||
|
||||
|
||||
class LoggingManager:
|
||||
"""
|
||||
A utility class for managing the logging configuration of Ray Train.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def _get_base_logger_config_dict(
|
||||
context: Union[TrainRunContext, TrainContext],
|
||||
) -> dict:
|
||||
"""Return the base logging configuration dictionary."""
|
||||
log_level = LoggingManager._resolve_log_level(context)
|
||||
# Using Ray worker ID as the file identifier where logs are written to.
|
||||
file_identifier = ray.get_runtime_context().get_worker_id()
|
||||
# Return the base logging configuration as a Python dictionary.
|
||||
return {
|
||||
"version": 1,
|
||||
"disable_existing_loggers": False,
|
||||
"formatters": {
|
||||
"ray_json": {"class": get_module_name(JSONFormatter)},
|
||||
},
|
||||
"filters": {
|
||||
"core_context_filter": {"()": CoreContextFilter},
|
||||
"train_context_filter": {"()": TrainContextFilter, "context": context},
|
||||
"train_log_level_filter": {
|
||||
"()": TrainLogLevelFilter,
|
||||
"log_level": log_level,
|
||||
},
|
||||
},
|
||||
"handlers": {
|
||||
"console": {
|
||||
"class": get_module_name(PlainRayHandler),
|
||||
"filters": ["train_log_level_filter"],
|
||||
},
|
||||
"file_train_sys_controller": {
|
||||
"class": get_module_name(SessionFileHandler),
|
||||
"formatter": "ray_json",
|
||||
"filename": f"ray-train-sys-controller-{file_identifier}.log",
|
||||
"filters": ["core_context_filter", "train_context_filter"],
|
||||
},
|
||||
"file_train_app_controller": {
|
||||
"class": get_module_name(SessionFileHandler),
|
||||
"formatter": "ray_json",
|
||||
"filename": f"ray-train-app-controller-{file_identifier}.log",
|
||||
"filters": [
|
||||
"core_context_filter",
|
||||
"train_context_filter",
|
||||
"train_log_level_filter",
|
||||
],
|
||||
},
|
||||
"file_train_sys_worker": {
|
||||
"class": get_module_name(SessionFileHandler),
|
||||
"formatter": "ray_json",
|
||||
"filename": f"ray-train-sys-worker-{file_identifier}.log",
|
||||
"filters": ["core_context_filter", "train_context_filter"],
|
||||
},
|
||||
"file_train_app_worker": {
|
||||
"class": get_module_name(SessionFileHandler),
|
||||
"formatter": "ray_json",
|
||||
"filename": f"ray-train-app-worker-{file_identifier}.log",
|
||||
"filters": [
|
||||
"core_context_filter",
|
||||
"train_context_filter",
|
||||
"train_log_level_filter",
|
||||
],
|
||||
},
|
||||
},
|
||||
"loggers": {},
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _resolve_log_level(
|
||||
context: Union[TrainRunContext, TrainContext],
|
||||
) -> str:
|
||||
"""Returns the log level from RunConfig's LoggingConfig."""
|
||||
if isinstance(context, TrainContext):
|
||||
run_config = context.train_run_context.get_run_config()
|
||||
else:
|
||||
run_config = context.get_run_config()
|
||||
|
||||
return run_config.logging_config.log_level
|
||||
|
||||
@staticmethod
|
||||
def _get_controller_logger_config_dict(context: TrainRunContext) -> dict:
|
||||
"""Return the controller logger configuration dictionary.
|
||||
|
||||
On the controller process, only the `ray.train` logger is configured.
|
||||
It is broadly set to level DEBUG, with downstream processing by log handlers.
|
||||
This logger emits logs to the following three locations:
|
||||
- `file_train_sys_controller`: Ray Train system logs.
|
||||
- `file_train_app_controller`: Ray Train application logs.
|
||||
- `console`: Logs to the console.
|
||||
"""
|
||||
|
||||
config_dict = LoggingManager._get_base_logger_config_dict(context)
|
||||
config_dict["loggers"]["ray.train"] = {
|
||||
"level": "DEBUG",
|
||||
"handlers": [
|
||||
"file_train_sys_controller",
|
||||
"file_train_app_controller",
|
||||
"console",
|
||||
],
|
||||
"propagate": False,
|
||||
}
|
||||
return config_dict
|
||||
|
||||
@staticmethod
|
||||
def _get_worker_logger_config_dict(context: TrainContext) -> dict:
|
||||
"""Return the worker loggers configuration dictionary.
|
||||
|
||||
On the worker process, there are two loggers being configured:
|
||||
|
||||
First, the `ray.train` logger is configured and emits logs to the
|
||||
following three locations:
|
||||
- `file_train_sys_worker`: Ray Train system logs.
|
||||
- `file_train_app_worker`: Ray Train application logs.
|
||||
- `console`: Logs to the console.
|
||||
It is broadly set to level DEBUG, with downstream processing by log handlers.
|
||||
|
||||
Second, the root logger is configured and emits logs to the following
|
||||
two locations:
|
||||
- `console`: Logs to the console.
|
||||
- `file_train_app_worker`: Ray Train application logs.
|
||||
The root logger will not emit Ray Train system logs and thus not writing to
|
||||
`file_train_sys_worker` file handler.
|
||||
"""
|
||||
|
||||
config_dict = LoggingManager._get_base_logger_config_dict(context)
|
||||
config_dict["loggers"]["ray.train"] = {
|
||||
"level": "DEBUG",
|
||||
"handlers": ["file_train_sys_worker", "file_train_app_worker", "console"],
|
||||
"propagate": False,
|
||||
}
|
||||
config_dict["root"] = {
|
||||
"level": "INFO",
|
||||
"handlers": ["file_train_app_worker", "console"],
|
||||
}
|
||||
return config_dict
|
||||
|
||||
@staticmethod
|
||||
def configure_controller_logger(context: TrainRunContext) -> None:
|
||||
"""
|
||||
Configure the logger on the controller process, which is the `ray.train` logger.
|
||||
"""
|
||||
config = LoggingManager._get_controller_logger_config_dict(context)
|
||||
logging.config.dictConfig(config)
|
||||
# TODO: Return the controller log file path.
|
||||
|
||||
@staticmethod
|
||||
def configure_worker_logger(context: TrainContext) -> None:
|
||||
"""
|
||||
Configure the loggers on the worker process, which contains the
|
||||
`ray.train` logger and the root logger.
|
||||
"""
|
||||
config = LoggingManager._get_worker_logger_config_dict(context)
|
||||
logging.config.dictConfig(config)
|
||||
# TODO: Return the worker log file path.
|
||||
|
||||
@staticmethod
|
||||
def get_log_directory() -> Optional[str]:
|
||||
"""Return the directory where Ray Train writes log files.
|
||||
|
||||
If not in a Ray session, return None.
|
||||
|
||||
This path looks like: "/tmp/ray/session_xxx/logs/train/"
|
||||
"""
|
||||
global_node = ray._private.worker._global_node
|
||||
|
||||
if global_node is None:
|
||||
return None
|
||||
|
||||
root_dir = global_node.get_session_dir_path()
|
||||
return os.path.join(root_dir, "logs", "train")
|
||||
|
||||
|
||||
def get_train_application_controller_log_path() -> Optional[str]:
|
||||
"""
|
||||
Return the path to the file train application controller log file.
|
||||
"""
|
||||
# TODO: This is a temporary solution. We should return the log file path in
|
||||
# the `configure_controller_logger` function.
|
||||
logger = logging.getLogger("ray.train")
|
||||
for handler in logger.handlers:
|
||||
if (
|
||||
isinstance(handler, SessionFileHandler)
|
||||
and "ray-train-app-controller" in handler._filename
|
||||
):
|
||||
return handler.get_log_file_path()
|
||||
return None
|
||||
|
||||
|
||||
def get_train_application_worker_log_path() -> Optional[str]:
|
||||
"""
|
||||
Return the path to the file train application worker log file.
|
||||
"""
|
||||
# TODO: This is a temporary solution. We should return the log file path in
|
||||
# the `configure_worker_logger` function.
|
||||
logger = logging.getLogger("ray.train")
|
||||
for handler in logger.handlers:
|
||||
if (
|
||||
isinstance(handler, SessionFileHandler)
|
||||
and "ray-train-app-worker" in handler._filename
|
||||
):
|
||||
return handler.get_log_file_path()
|
||||
return None
|
||||
@@ -0,0 +1,76 @@
|
||||
import builtins
|
||||
import contextlib
|
||||
import logging
|
||||
import sys
|
||||
from typing import Callable
|
||||
|
||||
from ray._private.ray_constants import env_bool
|
||||
from ray.train.v2._internal.constants import (
|
||||
DEFAULT_ENABLE_PRINT_PATCH,
|
||||
ENABLE_PRINT_PATCH_ENV_VAR,
|
||||
)
|
||||
|
||||
# Save the original print function
|
||||
_original_print = builtins.print
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def print_context_manager(print_fn: Callable):
|
||||
"""Context manager to set the builtin print function as print_fn."""
|
||||
current_print = builtins.print
|
||||
builtins.print = print_fn
|
||||
yield
|
||||
builtins.print = current_print
|
||||
|
||||
|
||||
def redirected_print(*objects, sep=" ", end="\n", file=None, flush=False):
|
||||
"""Implement python's print function to redirect logs to Train's logger.
|
||||
|
||||
If the file is set to anything other than stdout, stderr, or None, call the
|
||||
builtin print. Else, construct the message and redirect to Train's logger.
|
||||
|
||||
This makes sure that print to customized file in user defined function will not
|
||||
be overwritten by the redirected print function.
|
||||
|
||||
See https://docs.python.org/3/library/functions.html#print
|
||||
"""
|
||||
# TODO (hpguo): This handler class is shared by both ray train and ray serve. We
|
||||
# should move this to ray core and make it available to both libraries.
|
||||
|
||||
if file not in [sys.stdout, sys.stderr, None]:
|
||||
_original_print(*objects, sep=sep, end=end, file=file, flush=flush)
|
||||
return
|
||||
|
||||
# If sys.stdout/stderr has been redirected (e.g. contextlib.redirect_stdout(),
|
||||
# or wrapping by libraries like wandb / MLflow / colorama / IPython), tee to
|
||||
# the original print so the redirect target also receives the output. The
|
||||
# logger still gets the message below, so structured logs aren't silently
|
||||
# dropped when a third-party library wraps the stream.
|
||||
if (file in (sys.stdout, None) and sys.stdout is not sys.__stdout__) or (
|
||||
file is sys.stderr and sys.stderr is not sys.__stderr__
|
||||
):
|
||||
_original_print(*objects, sep=sep, end=end, file=file, flush=flush)
|
||||
|
||||
root_logger = logging.getLogger()
|
||||
message = sep.join(map(str, objects))
|
||||
# Use the original `print` method for the scope of the logger call, in order to
|
||||
# avoid infinite recursion errors if any exceptions get raised (since exception
|
||||
# handling involves another `print(..., file=sys.stderr)`.
|
||||
# Note that an exception being raised here is not expected (e.g. it would be a
|
||||
# bug in our own logging code), so this is just to keep the error logs sane
|
||||
# during development.
|
||||
with print_context_manager(_original_print):
|
||||
# We want this log to be associated with the line of code where user calls
|
||||
# `print`, which is stacklevel 2.
|
||||
# Frame [stacklevel]:
|
||||
# User's call to print [2] -> `redirected_print` [1] -> root_logger.log [0]
|
||||
root_logger.log(logging.INFO, message, stacklevel=2)
|
||||
|
||||
|
||||
def patch_print_function() -> None:
|
||||
"""
|
||||
Patch the print function to redirect logs to Train's logger.
|
||||
Only patch the print function if the environment variable is set to "1"
|
||||
"""
|
||||
if env_bool(ENABLE_PRINT_PATCH_ENV_VAR, DEFAULT_ENABLE_PRINT_PATCH):
|
||||
builtins.print = redirected_print
|
||||
Reference in New Issue
Block a user