[Core] Turn off color logging prefix (#45218)

This pr adds optionality to turn off color logging from actors. This is useful in the case where the logs are outputted to a location that does not support ansi codes, which leads to messy logs.

Signed-off-by: mliu <[email protected]>

doc/source/ray-observability/user-guides/configure-logging.md

@@ -95,8 +95,10 @@ The resulting output follows:
 
 ### Coloring Actor log prefixes
 By default, Ray prints Actor log prefixes in light blue.
-Activate multi-color prefixes by setting the environment variable ``RAY_COLOR_PREFIX=1``.
-This indexes into an array of colors modulo the PID of each process.
+Turn color logging off by setting the environment variable ``RAY_COLOR_PREFIX=0``
+(for example, when outputting logs to a file or other location that does not support ANSI codes).
+Or activate multi-color prefixes by setting the environment variable ``RAY_COLOR_PREFIX=1``;
+this indexes into an array of colors modulo the PID of each process.
 
 ![coloring-actor-log-prefixes](../images/coloring-actor-log-prefixes.png)
 

python/ray/_private/worker.py

@@ -2045,39 +2045,25 @@ def color_for(data: Dict[str, str], line: str) -> str:
  pid = data.get("pid")
  lines = data.get("lines", [])
 
- if data.get("ip") == data.get("localhost"):
- for line in lines:
- if RAY_TQDM_MAGIC in line:
- process_tqdm(line)
- else:
- hide_tqdm()
- print(
- "{}({}{}){} {}".format(
- color_for(data, line),
- prefix_for(data),
- pid,
- colorama.Style.RESET_ALL,
- message_for(data, line),
- ),
- file=print_file,
- )
- else:
- for line in lines:
- if RAY_TQDM_MAGIC in line:
- process_tqdm(line)
+ ip = data.get("ip")
+ ip_prefix = "" if ip == data.get("localhost") else f", ip={ip}"
+ for line in lines:
+ if RAY_TQDM_MAGIC in line:
+ process_tqdm(line)
+ else:
+ hide_tqdm()
+ # If RAY_COLOR_PREFIX=0, do not wrap with any color codes
+ if os.getenv("RAY_COLOR_PREFIX") == "0":
+ color_pre = ""
+ color_post = ""
  else:
- hide_tqdm()
- print(
- "{}({}{}, ip={}){} {}".format(
- color_for(data, line),
- prefix_for(data),
- pid,
- data.get("ip"),
- colorama.Style.RESET_ALL,
- message_for(data, line),
- ),
- file=print_file,
- )
+ color_pre = color_for(data, line)
+ color_post = colorama.Style.RESET_ALL
+ print(
+ f"{color_pre}({prefix_for(data)}{pid}{ip_prefix}){color_post} "
+ f"{message_for(data, line)}",
+ file=print_file,
+ )
  # Restore once at end of batch to avoid excess hiding/unhiding of tqdm.
  restore_tqdm()
 

python/ray/tests/test_logging.py

@@ -9,8 +9,10 @@
 from collections import Counter, defaultdict
 from contextlib import redirect_stderr, redirect_stdout
 from pathlib import Path
-from unittest.mock import Mock, MagicMock
+from typing import Dict
+from unittest.mock import Mock, MagicMock, patch
 
+import colorama
 import pytest
 
 import ray
@@ -32,6 +34,7 @@
 )
 from ray.cross_language import java_actor_class
 from ray.autoscaler._private.cli_logger import cli_logger
+from ray._private.worker import print_worker_logs
 
 
 def set_logging_config(monkeypatch, max_bytes, backup_count):
@@ -961,6 +964,100 @@ def print_msg():
  assert data[0]["ip"] == "127.0.0.2"
 
 
+def get_print_worker_logs_output(data: Dict[str, str]) -> str:
+ """
+ Helper function that returns the output of `print_worker_logs` as a str.
+ """
+ out = io.StringIO()
+ print_worker_logs(data, out)
+ out.seek(0)
+ return out.readline()
+
+
+def test_print_worker_logs_default_color() -> None:
+ # Test multiple since pid may affect color
+ for pid in (0, 1):
+ data = dict(
+ ip="10.0.0.1",
+ localhost="172.0.0.1",
+ pid=str(pid),
+ task_name="my_task",
+ lines=["is running"],
+ )
+ output = get_print_worker_logs_output(data)
+ assert output == (
+ f"{colorama.Fore.CYAN}(my_task pid={pid}, ip=10.0.0.1)"
+ + f"{colorama.Style.RESET_ALL} is running\n"
+ )
+
+ # Special case
+ raylet = dict(
+ ip="10.0.0.1",
+ localhost="172.0.0.1",
+ pid="raylet",
+ task_name="my_task",
+ lines=["Warning: uh oh"],
+ )
+ output = get_print_worker_logs_output(raylet)
+ assert output == (
+ f"{colorama.Fore.YELLOW}(raylet, ip=10.0.0.1){colorama.Style.RESET_ALL} "
+ + "Warning: uh oh\n"
+ )
+
+
+@patch.dict(os.environ, {"RAY_COLOR_PREFIX": "0"})
+def test_print_worker_logs_no_color() -> None:
+ for pid in (0, 1):
+ data = dict(
+ ip="10.0.0.1",
+ localhost="172.0.0.1",
+ pid=str(pid),
+ task_name="my_task",
+ lines=["is running"],
+ )
+ output = get_print_worker_logs_output(data)
+ assert output == f"(my_task pid={pid}, ip=10.0.0.1) is running\n"
+
+ raylet = dict(
+ ip="10.0.0.1",
+ localhost="172.0.0.1",
+ pid="raylet",
+ task_name="my_task",
+ lines=["Warning: uh oh"],
+ )
+ output = get_print_worker_logs_output(raylet)
+ assert output == "(raylet, ip=10.0.0.1) Warning: uh oh\n"
+
+
+@patch.dict(os.environ, {"RAY_COLOR_PREFIX": "1"})
+def test_print_worker_logs_multi_color() -> None:
+ data_pid_0 = dict(
+ ip="10.0.0.1",
+ localhost="172.0.0.1",
+ pid="0",
+ task_name="my_task",
+ lines=["is running"],
+ )
+ output = get_print_worker_logs_output(data_pid_0)
+ assert output == (
+ f"{colorama.Fore.MAGENTA}(my_task pid=0, ip=10.0.0.1)"
+ + f"{colorama.Style.RESET_ALL} is running\n"
+ )
+
+ data_pid_2 = dict(
+ ip="10.0.0.1",
+ localhost="172.0.0.1",
+ pid="2",
+ task_name="my_task",
+ lines=["is running"],
+ )
+ output = get_print_worker_logs_output(data_pid_2)
+ assert output == (
+ f"{colorama.Fore.GREEN}(my_task pid=2, ip=10.0.0.1){colorama.Style.RESET_ALL} "
+ + "is running\n"
+ )
+
+
 if __name__ == "__main__":
  import sys