chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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
|
||||
Reference in New Issue
Block a user