Skip to content

Commit 617b274

Browse files
author
Roja Reddy Sareddy
committed
Merge branch 'fix/core-rich-logging-opt-in' into fix/core-stdout-logging-on-import
# Conflicts: # sagemaker-core/src/sagemaker/core/__init__.py
2 parents 2f00a11 + 4bfc1dd commit 617b274

3 files changed

Lines changed: 165 additions & 11 deletions

File tree

‎sagemaker-core/src/sagemaker/core/__init__.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
if not any(isinstance(_h, _logging.NullHandler) for _h in _sagemaker_root_logger.handlers):
1414
_sagemaker_root_logger.addHandler(_logging.NullHandler())
1515

16-
enable_textual_rich_console_and_traceback()
16+
enable_textual_rich_console_and_traceback() # opt-in; no-op unless SAGEMAKER_ENABLE_RICH_LOGGING is set
1717

1818
# Install the meta-path finder that gives actionable migration guidance for v2
1919
# modules removed in v3. sagemaker-core is the universal dependency of every v3

‎sagemaker-core/src/sagemaker/core/utils/utils.py‎

Lines changed: 50 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -137,12 +137,47 @@ def get_textual_rich_theme() -> Theme:
137137
)
138138

139139

140+
RICH_LOGGING_OPT_IN_ENV_VAR = "SAGEMAKER_ENABLE_RICH_LOGGING"
141+
142+
_TRUTHY_ENV_VALUES = frozenset({"1", "true", "yes", "on"})
143+
144+
145+
def is_rich_logging_enabled() -> bool:
146+
"""Whether the user opted in to sagemaker-core's rich console and tracebacks.
147+
148+
Reconfiguring the global rich console and calling ``rich.traceback.install()``
149+
override ``sys.excepthook`` and restyle the process-global console, so they are
150+
opt-in: merely importing the SDK must not change tracebacks or console styling.
151+
Enable by setting the ``SAGEMAKER_ENABLE_RICH_LOGGING`` environment variable to
152+
one of ``1``/``true``/``yes``/``on`` (case-insensitive).
153+
154+
Returns:
155+
bool: True if rich console/traceback output has been opted into.
156+
"""
157+
return os.environ.get(RICH_LOGGING_OPT_IN_ENV_VAR, "").strip().lower() in _TRUTHY_ENV_VALUES
158+
159+
140160
textual_rich_console_and_traceback_enabled = False
141161

142162

143-
def enable_textual_rich_console_and_traceback():
144-
"""Reconfigure the global textual rich console with the customized theme and enable textual rich error traceback"""
163+
def enable_textual_rich_console_and_traceback(force: bool = False):
164+
"""Reconfigure the global rich console and install rich error tracebacks.
165+
166+
This overrides ``sys.excepthook`` (via ``rich.traceback.install``) and restyles
167+
the process-global rich console. Because those are process-wide side effects, it
168+
is opt-in and a no-op unless the user opts in via the
169+
``SAGEMAKER_ENABLE_RICH_LOGGING`` environment variable
170+
(see :func:`is_rich_logging_enabled`) or the caller passes ``force=True``. This
171+
keeps ``import sagemaker`` free of global traceback/console side effects by
172+
default.
173+
174+
Args:
175+
force (bool): Enable regardless of the environment variable, for callers
176+
that explicitly want rich output. Defaults to False.
177+
"""
145178
global textual_rich_console_and_traceback_enabled
179+
if not (force or is_rich_logging_enabled()):
180+
return
146181
if not textual_rich_console_and_traceback_enabled:
147182
theme = get_textual_rich_theme()
148183
reconfigure(theme=theme)
@@ -159,23 +194,28 @@ def get_rich_handler():
159194

160195

161196
def get_textual_rich_logger(name: str, log_level: str = "INFO") -> logging.Logger:
162-
"""Get a logger with textual rich handler.
197+
"""Get a logger, attaching a rich handler only when rich logging is opted in.
198+
199+
Rich logging (a ``RichHandler`` on the root logger via ``logging.basicConfig``,
200+
plus the themed console/traceback) is opt-in, so that importing the SDK does not
201+
reconfigure the root logger or change process-wide log formatting/level. When the
202+
user has not opted in (see :func:`is_rich_logging_enabled`), this returns the named
203+
logger without configuring handlers or levels, leaving logging to the application.
163204
164205
Args:
165206
name (str): The name of the logger
166-
log_level (str): The log level to set.
207+
log_level (str): The log level to set when rich logging is enabled.
167208
Accepted values are: "DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL".
168209
Defaults to the value of "INFO".
169210
170211
Return:
171-
logging.Logger: A textial rich logger.
212+
logging.Logger: The requested logger.
172213
"""
173214
enable_textual_rich_console_and_traceback()
174-
handler = get_rich_handler()
175-
logging.basicConfig(level=getattr(logging, log_level), handlers=[handler])
176-
logger = logging.getLogger(name)
177-
178-
return logger
215+
if is_rich_logging_enabled():
216+
handler = get_rich_handler()
217+
logging.basicConfig(level=getattr(logging, log_level), handlers=[handler])
218+
return logging.getLogger(name)
179219

180220

181221
logger = get_textual_rich_logger(__name__)
Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
"""Unit tests for opt-in rich console/traceback behavior in sagemaker.core.utils.utils.
2+
3+
These lock in that importing the SDK does not override sys.excepthook or restyle the
4+
process-global rich console unless the user explicitly opts in via
5+
SAGEMAKER_ENABLE_RICH_LOGGING (or a force=True call).
6+
"""
7+
8+
import os
9+
from unittest.mock import patch
10+
11+
import pytest
12+
13+
from sagemaker.core.utils import utils
14+
from sagemaker.core.utils.utils import (
15+
RICH_LOGGING_OPT_IN_ENV_VAR,
16+
enable_textual_rich_console_and_traceback,
17+
is_rich_logging_enabled,
18+
)
19+
20+
21+
@pytest.fixture(autouse=True)
22+
def _reset_latch():
23+
"""Reset the one-shot 'already enabled' latch around every test."""
24+
saved = utils.textual_rich_console_and_traceback_enabled
25+
utils.textual_rich_console_and_traceback_enabled = False
26+
try:
27+
yield
28+
finally:
29+
utils.textual_rich_console_and_traceback_enabled = saved
30+
31+
32+
@pytest.mark.parametrize("value", ["1", "true", "TRUE", "Yes", "on", " on "])
33+
def test_is_rich_logging_enabled_truthy(value):
34+
with patch.dict(os.environ, {RICH_LOGGING_OPT_IN_ENV_VAR: value}):
35+
assert is_rich_logging_enabled() is True
36+
37+
38+
@pytest.mark.parametrize("value", ["", "0", "false", "no", "off", "nope"])
39+
def test_is_rich_logging_enabled_falsey(value):
40+
with patch.dict(os.environ, {RICH_LOGGING_OPT_IN_ENV_VAR: value}):
41+
assert is_rich_logging_enabled() is False
42+
43+
44+
def test_disabled_by_default_is_noop():
45+
# Env var absent -> importing/using the SDK must not touch the global console
46+
# or install rich tracebacks (no sys.excepthook override).
47+
with patch.dict(os.environ, {}, clear=False):
48+
os.environ.pop(RICH_LOGGING_OPT_IN_ENV_VAR, None)
49+
with (
50+
patch.object(utils, "reconfigure") as mock_reconfigure,
51+
patch.object(utils, "install") as mock_install,
52+
):
53+
enable_textual_rich_console_and_traceback()
54+
mock_reconfigure.assert_not_called()
55+
mock_install.assert_not_called()
56+
assert utils.textual_rich_console_and_traceback_enabled is False
57+
58+
59+
def test_enabled_when_opted_in_via_env():
60+
with patch.dict(os.environ, {RICH_LOGGING_OPT_IN_ENV_VAR: "true"}):
61+
with (
62+
patch.object(utils, "reconfigure") as mock_reconfigure,
63+
patch.object(utils, "install") as mock_install,
64+
):
65+
enable_textual_rich_console_and_traceback()
66+
mock_reconfigure.assert_called_once()
67+
mock_install.assert_called_once()
68+
assert utils.textual_rich_console_and_traceback_enabled is True
69+
70+
71+
def test_force_enables_regardless_of_env():
72+
with patch.dict(os.environ, {}, clear=False):
73+
os.environ.pop(RICH_LOGGING_OPT_IN_ENV_VAR, None)
74+
with (
75+
patch.object(utils, "reconfigure") as mock_reconfigure,
76+
patch.object(utils, "install") as mock_install,
77+
):
78+
enable_textual_rich_console_and_traceback(force=True)
79+
mock_reconfigure.assert_called_once()
80+
mock_install.assert_called_once()
81+
82+
83+
def test_enable_is_idempotent_when_opted_in():
84+
with patch.dict(os.environ, {RICH_LOGGING_OPT_IN_ENV_VAR: "1"}):
85+
with (
86+
patch.object(utils, "reconfigure") as mock_reconfigure,
87+
patch.object(utils, "install") as mock_install,
88+
):
89+
enable_textual_rich_console_and_traceback()
90+
enable_textual_rich_console_and_traceback()
91+
# The one-shot latch prevents re-installing on the second call.
92+
mock_reconfigure.assert_called_once()
93+
mock_install.assert_called_once()
94+
95+
96+
def test_get_logger_does_not_call_basicconfig_when_opted_out():
97+
# Getting a module logger must not reconfigure the root logger by default.
98+
with patch.dict(os.environ, {}, clear=False):
99+
os.environ.pop(RICH_LOGGING_OPT_IN_ENV_VAR, None)
100+
with patch.object(utils.logging, "basicConfig") as mock_basic_config:
101+
returned = utils.get_textual_rich_logger("sagemaker.core.test.optout")
102+
mock_basic_config.assert_not_called()
103+
assert returned is utils.logging.getLogger("sagemaker.core.test.optout")
104+
105+
106+
def test_get_logger_calls_basicconfig_when_opted_in():
107+
with patch.dict(os.environ, {RICH_LOGGING_OPT_IN_ENV_VAR: "1"}):
108+
with (
109+
patch.object(utils, "reconfigure"),
110+
patch.object(utils, "install"),
111+
patch.object(utils.logging, "basicConfig") as mock_basic_config,
112+
):
113+
utils.get_textual_rich_logger("sagemaker.core.test.optin")
114+
mock_basic_config.assert_called_once()

0 commit comments

Comments
 (0)