From b359c97b4751f202ab511b810844221c9de5861c Mon Sep 17 00:00:00 2001
From: Gene Der Su <e870252314@gmail.com>
Date: Fri, 28 Jun 2024 21:31:04 -0700
Subject: [PATCH] [Core] Allow env setup logger encoding (#46242)

Signed-off-by: Gene Su <e870252314@gmail.com>
---
 python/ray/_private/ray_constants.py |  2 +
 python/ray/_private/worker.py        | 17 +++++--
 python/ray/tests/test_logging_2.py   | 76 ++++++++++++++++++++++++++++
 3 files changed, 91 insertions(+), 4 deletions(-)

diff --git a/python/ray/_private/ray_constants.py b/python/ray/_private/ray_constants.py
index fe145496978e..1ecfdd752019 100644
--- a/python/ray/_private/ray_constants.py
+++ b/python/ray/_private/ray_constants.py
@@ -483,3 +483,5 @@ def gcs_actor_scheduling_enabled():
 RAY_NODE_IP_FILENAME = "node_ip_address.json"
 
 PLACEMENT_GROUP_BUNDLE_RESOURCE_NAME = "bundle"
+
+RAY_LOGGING_CONFIG_ENCODING = os.environ.get("RAY_LOGGING_CONFIG_ENCODING")
diff --git a/python/ray/_private/worker.py b/python/ray/_private/worker.py
index b612a87ea0a3..e3b336b599b6 100644
--- a/python/ray/_private/worker.py
+++ b/python/ray/_private/worker.py
@@ -1406,8 +1406,14 @@ def init(
         logging.getLogger("ray").handlers.clear()
 
     # Configure the logging settings for the driver process.
-    if logging_config:
-        dict_config = logging_config._get_dict_config()
+    if logging_config or ray_constants.RAY_LOGGING_CONFIG_ENCODING:
+        dict_config = (
+            logging_config._get_dict_config()
+            if logging_config
+            else LoggingConfig(
+                encoding=ray_constants.RAY_LOGGING_CONFIG_ENCODING
+            )._get_dict_config()
+        )
         logging.config.dictConfig(dict_config)
 
     # Parse the hidden options:
@@ -1584,8 +1590,11 @@ def sigterm_handler(signum, frame):
 
     # Pass the logging_config to job_config to configure loggers of all worker
     # processes belonging to the job.
-    if logging_config:
-        job_config.set_py_logging_config(logging_config)
+    if logging_config or ray_constants.RAY_LOGGING_CONFIG_ENCODING:
+        job_config.set_py_logging_config(
+            logging_config
+            or LoggingConfig(encoding=ray_constants.RAY_LOGGING_CONFIG_ENCODING)
+        )
 
     redis_address, gcs_address = None, None
     bootstrap_address = services.canonicalize_bootstrap_address(address, _temp_dir)
diff --git a/python/ray/tests/test_logging_2.py b/python/ray/tests/test_logging_2.py
index 7adb26e976ee..92f9ea3acb44 100644
--- a/python/ray/tests/test_logging_2.py
+++ b/python/ray/tests/test_logging_2.py
@@ -262,6 +262,82 @@ def print_message(self):
         for s in should_exist:
             assert s in stderr
 
+    @pytest.mark.parametrize(
+        "ray_start_cluster_head_with_env_vars",
+        [
+            {
+                "env_vars": {
+                    "RAY_LOGGING_CONFIG_ENCODING": "TEXT",
+                },
+            }
+        ],
+        indirect=True,
+    )
+    def test_env_setup_logger_encoding(
+        self, ray_start_cluster_head_with_env_vars, shutdown_only
+    ):
+        script = """
+import ray
+import logging
+
+ray.init()
+
+@ray.remote
+class actor:
+    def __init__(self):
+        pass
+
+    def print_message(self):
+        logger = logging.getLogger(__name__)
+        logger.info("This is a Ray actor")
+
+actor_instance = actor.remote()
+ray.get(actor_instance.print_message.remote())
+"""
+        stderr = run_string_as_driver(script)
+        should_exist = [
+            "job_id",
+            "worker_id",
+            "node_id",
+            "actor_id",
+            "task_id",
+            "INFO",
+            "This is a Ray actor",
+        ]
+        for s in should_exist:
+            assert s in stderr
+
+    def test_logger_not_set(self, shutdown_only):
+        script = """
+import ray
+import logging
+
+ray.init()
+
+@ray.remote
+class actor:
+    def __init__(self):
+        pass
+
+    def print_message(self):
+        logger = logging.getLogger(__name__)
+        logger.info("This is a Ray actor")
+
+actor_instance = actor.remote()
+ray.get(actor_instance.print_message.remote())
+"""
+        stderr = run_string_as_driver(script)
+        should_not_exist = [
+            "job_id",
+            "worker_id",
+            "node_id",
+            "actor_id",
+            "task_id",
+            "This is a Ray actor",
+        ]
+        for s in should_not_exist:
+            assert s not in stderr
+
 
 if __name__ == "__main__":
     if os.environ.get("PARALLEL_CI"):