wehub-resource-sync
77 lines
3.1 KiB
Python
77 lines
3.1 KiB
Python
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
|