wehub-resource-sync
341 lines
13 KiB
Python
341 lines
13 KiB
Python
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
|