chore: import upstream snapshot with attribution

python/ray/train/v2/_internal/logging/patch_print.py

@@ -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