From 1299434357b52ff12941fc5bc06d5f2862a05f00 Mon Sep 17 00:00:00 2001 From: Roja Reddy Sareddy Date: Sun, 27 Sep 2026 17:38:49 -0700 Subject: [PATCH 1/3] fix(core): make rich console and traceback opt-in Importing the SDK ran enable_textual_rich_console_and_traceback() at module import (sagemaker/core/__init__.py, and again via the module-level logger in utils.py), which called rich.traceback.install() -- overriding sys.excepthook -- and reconfigured the process-global rich console. Users who did not want rich tracebacks/console styling had no way to turn it off, and it changed traceback behavior for the whole process just by importing sagemaker. Gate those process-wide side effects behind an opt-in: they now run only when SAGEMAKER_ENABLE_RICH_LOGGING is set (1/true/yes/on) or a caller passes force=True. Default import no longer touches sys.excepthook or the global console. Add is_rich_logging_enabled() and unit tests covering the default no-op, env opt-in, force, and idempotency. Closes #5179 Closes #5561 --- sagemaker-core/src/sagemaker/core/__init__.py | 2 +- .../src/sagemaker/core/utils/utils.py | 41 ++++++++- .../tests/unit/utils/test_rich_logging.py | 88 +++++++++++++++++++ 3 files changed, 128 insertions(+), 3 deletions(-) create mode 100644 sagemaker-core/tests/unit/utils/test_rich_logging.py diff --git a/sagemaker-core/src/sagemaker/core/__init__.py b/sagemaker-core/src/sagemaker/core/__init__.py index fe829739b3..d5bfb4586a 100644 --- a/sagemaker-core/src/sagemaker/core/__init__.py +++ b/sagemaker-core/src/sagemaker/core/__init__.py @@ -3,7 +3,7 @@ from sagemaker.core.utils.utils import enable_textual_rich_console_and_traceback from sagemaker.core.deprecations import register_removed_module_finder -enable_textual_rich_console_and_traceback() +enable_textual_rich_console_and_traceback() # opt-in; no-op unless SAGEMAKER_ENABLE_RICH_LOGGING is set # Install the meta-path finder that gives actionable migration guidance for v2 # modules removed in v3. sagemaker-core is the universal dependency of every v3 diff --git a/sagemaker-core/src/sagemaker/core/utils/utils.py b/sagemaker-core/src/sagemaker/core/utils/utils.py index dfc815190a..2eb9d44c37 100644 --- a/sagemaker-core/src/sagemaker/core/utils/utils.py +++ b/sagemaker-core/src/sagemaker/core/utils/utils.py @@ -137,12 +137,49 @@ def get_textual_rich_theme() -> Theme: ) +RICH_LOGGING_OPT_IN_ENV_VAR = "SAGEMAKER_ENABLE_RICH_LOGGING" + +_TRUTHY_ENV_VALUES = frozenset({"1", "true", "yes", "on"}) + + +def is_rich_logging_enabled() -> bool: + """Whether the user opted in to sagemaker-core's rich console and tracebacks. + + Reconfiguring the global rich console and calling ``rich.traceback.install()`` + override ``sys.excepthook`` and restyle the process-global console, so they are + opt-in: merely importing the SDK must not change tracebacks or console styling. + Enable by setting the ``SAGEMAKER_ENABLE_RICH_LOGGING`` environment variable to + one of ``1``/``true``/``yes``/``on`` (case-insensitive). + + Returns: + bool: True if rich console/traceback output has been opted into. + """ + return ( + os.environ.get(RICH_LOGGING_OPT_IN_ENV_VAR, "").strip().lower() in _TRUTHY_ENV_VALUES + ) + + textual_rich_console_and_traceback_enabled = False -def enable_textual_rich_console_and_traceback(): - """Reconfigure the global textual rich console with the customized theme and enable textual rich error traceback""" +def enable_textual_rich_console_and_traceback(force: bool = False): + """Reconfigure the global rich console and install rich error tracebacks. + + This overrides ``sys.excepthook`` (via ``rich.traceback.install``) and restyles + the process-global rich console. Because those are process-wide side effects, it + is opt-in and a no-op unless the user opts in via the + ``SAGEMAKER_ENABLE_RICH_LOGGING`` environment variable + (see :func:`is_rich_logging_enabled`) or the caller passes ``force=True``. This + keeps ``import sagemaker`` free of global traceback/console side effects by + default. + + Args: + force (bool): Enable regardless of the environment variable, for callers + that explicitly want rich output. Defaults to False. + """ global textual_rich_console_and_traceback_enabled + if not (force or is_rich_logging_enabled()): + return if not textual_rich_console_and_traceback_enabled: theme = get_textual_rich_theme() reconfigure(theme=theme) diff --git a/sagemaker-core/tests/unit/utils/test_rich_logging.py b/sagemaker-core/tests/unit/utils/test_rich_logging.py new file mode 100644 index 0000000000..b5016bbdec --- /dev/null +++ b/sagemaker-core/tests/unit/utils/test_rich_logging.py @@ -0,0 +1,88 @@ +"""Unit tests for opt-in rich console/traceback behavior in sagemaker.core.utils.utils. + +These lock in that importing the SDK does not override sys.excepthook or restyle the +process-global rich console unless the user explicitly opts in via +SAGEMAKER_ENABLE_RICH_LOGGING (or a force=True call). +""" +import os +from unittest.mock import patch + +import pytest + +from sagemaker.core.utils import utils +from sagemaker.core.utils.utils import ( + RICH_LOGGING_OPT_IN_ENV_VAR, + enable_textual_rich_console_and_traceback, + is_rich_logging_enabled, +) + + +@pytest.fixture(autouse=True) +def _reset_latch(): + """Reset the one-shot 'already enabled' latch around every test.""" + saved = utils.textual_rich_console_and_traceback_enabled + utils.textual_rich_console_and_traceback_enabled = False + try: + yield + finally: + utils.textual_rich_console_and_traceback_enabled = saved + + +@pytest.mark.parametrize("value", ["1", "true", "TRUE", "Yes", "on", " on "]) +def test_is_rich_logging_enabled_truthy(value): + with patch.dict(os.environ, {RICH_LOGGING_OPT_IN_ENV_VAR: value}): + assert is_rich_logging_enabled() is True + + +@pytest.mark.parametrize("value", ["", "0", "false", "no", "off", "nope"]) +def test_is_rich_logging_enabled_falsey(value): + with patch.dict(os.environ, {RICH_LOGGING_OPT_IN_ENV_VAR: value}): + assert is_rich_logging_enabled() is False + + +def test_disabled_by_default_is_noop(): + # Env var absent -> importing/using the SDK must not touch the global console + # or install rich tracebacks (no sys.excepthook override). + with patch.dict(os.environ, {}, clear=False): + os.environ.pop(RICH_LOGGING_OPT_IN_ENV_VAR, None) + with patch.object(utils, "reconfigure") as mock_reconfigure, patch.object( + utils, "install" + ) as mock_install: + enable_textual_rich_console_and_traceback() + mock_reconfigure.assert_not_called() + mock_install.assert_not_called() + assert utils.textual_rich_console_and_traceback_enabled is False + + +def test_enabled_when_opted_in_via_env(): + with patch.dict(os.environ, {RICH_LOGGING_OPT_IN_ENV_VAR: "true"}): + with patch.object(utils, "reconfigure") as mock_reconfigure, patch.object( + utils, "install" + ) as mock_install: + enable_textual_rich_console_and_traceback() + mock_reconfigure.assert_called_once() + mock_install.assert_called_once() + assert utils.textual_rich_console_and_traceback_enabled is True + + +def test_force_enables_regardless_of_env(): + with patch.dict(os.environ, {}, clear=False): + os.environ.pop(RICH_LOGGING_OPT_IN_ENV_VAR, None) + with patch.object(utils, "reconfigure") as mock_reconfigure, patch.object( + utils, "install" + ) as mock_install: + enable_textual_rich_console_and_traceback(force=True) + mock_reconfigure.assert_called_once() + mock_install.assert_called_once() + + +def test_enable_is_idempotent_when_opted_in(): + with patch.dict(os.environ, {RICH_LOGGING_OPT_IN_ENV_VAR: "1"}): + with patch.object(utils, "reconfigure") as mock_reconfigure, patch.object( + utils, "install" + ) as mock_install: + enable_textual_rich_console_and_traceback() + enable_textual_rich_console_and_traceback() + # The one-shot latch prevents re-installing on the second call. + mock_reconfigure.assert_called_once() + mock_install.assert_called_once() From eca6c7013b7430f33232193ccb18b001c928e074 Mon Sep 17 00:00:00 2001 From: Roja Reddy Sareddy Date: Mon, 28 Sep 2026 11:27:51 -0700 Subject: [PATCH 2/3] style(core): apply black to rich-logging changes Reformat test_rich_logging.py (blank line after docstring, parenthesized context managers) and collapse a one-line return in utils.py to satisfy black --check in the codestyle-doc-tests CI job. --- .../src/sagemaker/core/utils/utils.py | 4 +-- .../tests/unit/utils/test_rich_logging.py | 29 +++++++++++-------- 2 files changed, 18 insertions(+), 15 deletions(-) diff --git a/sagemaker-core/src/sagemaker/core/utils/utils.py b/sagemaker-core/src/sagemaker/core/utils/utils.py index 2eb9d44c37..f0d270edac 100644 --- a/sagemaker-core/src/sagemaker/core/utils/utils.py +++ b/sagemaker-core/src/sagemaker/core/utils/utils.py @@ -154,9 +154,7 @@ def is_rich_logging_enabled() -> bool: Returns: bool: True if rich console/traceback output has been opted into. """ - return ( - os.environ.get(RICH_LOGGING_OPT_IN_ENV_VAR, "").strip().lower() in _TRUTHY_ENV_VALUES - ) + return os.environ.get(RICH_LOGGING_OPT_IN_ENV_VAR, "").strip().lower() in _TRUTHY_ENV_VALUES textual_rich_console_and_traceback_enabled = False diff --git a/sagemaker-core/tests/unit/utils/test_rich_logging.py b/sagemaker-core/tests/unit/utils/test_rich_logging.py index b5016bbdec..ba35d4d051 100644 --- a/sagemaker-core/tests/unit/utils/test_rich_logging.py +++ b/sagemaker-core/tests/unit/utils/test_rich_logging.py @@ -4,6 +4,7 @@ process-global rich console unless the user explicitly opts in via SAGEMAKER_ENABLE_RICH_LOGGING (or a force=True call). """ + import os from unittest.mock import patch @@ -45,9 +46,10 @@ def test_disabled_by_default_is_noop(): # or install rich tracebacks (no sys.excepthook override). with patch.dict(os.environ, {}, clear=False): os.environ.pop(RICH_LOGGING_OPT_IN_ENV_VAR, None) - with patch.object(utils, "reconfigure") as mock_reconfigure, patch.object( - utils, "install" - ) as mock_install: + with ( + patch.object(utils, "reconfigure") as mock_reconfigure, + patch.object(utils, "install") as mock_install, + ): enable_textual_rich_console_and_traceback() mock_reconfigure.assert_not_called() mock_install.assert_not_called() @@ -56,9 +58,10 @@ def test_disabled_by_default_is_noop(): def test_enabled_when_opted_in_via_env(): with patch.dict(os.environ, {RICH_LOGGING_OPT_IN_ENV_VAR: "true"}): - with patch.object(utils, "reconfigure") as mock_reconfigure, patch.object( - utils, "install" - ) as mock_install: + with ( + patch.object(utils, "reconfigure") as mock_reconfigure, + patch.object(utils, "install") as mock_install, + ): enable_textual_rich_console_and_traceback() mock_reconfigure.assert_called_once() mock_install.assert_called_once() @@ -68,9 +71,10 @@ def test_enabled_when_opted_in_via_env(): def test_force_enables_regardless_of_env(): with patch.dict(os.environ, {}, clear=False): os.environ.pop(RICH_LOGGING_OPT_IN_ENV_VAR, None) - with patch.object(utils, "reconfigure") as mock_reconfigure, patch.object( - utils, "install" - ) as mock_install: + with ( + patch.object(utils, "reconfigure") as mock_reconfigure, + patch.object(utils, "install") as mock_install, + ): enable_textual_rich_console_and_traceback(force=True) mock_reconfigure.assert_called_once() mock_install.assert_called_once() @@ -78,9 +82,10 @@ def test_force_enables_regardless_of_env(): def test_enable_is_idempotent_when_opted_in(): with patch.dict(os.environ, {RICH_LOGGING_OPT_IN_ENV_VAR: "1"}): - with patch.object(utils, "reconfigure") as mock_reconfigure, patch.object( - utils, "install" - ) as mock_install: + with ( + patch.object(utils, "reconfigure") as mock_reconfigure, + patch.object(utils, "install") as mock_install, + ): enable_textual_rich_console_and_traceback() enable_textual_rich_console_and_traceback() # The one-shot latch prevents re-installing on the second call. From 4bfc1dd4ef95f1d49a27cea59631385cc6df3bd8 Mon Sep 17 00:00:00 2001 From: Roja Reddy Sareddy Date: Mon, 28 Sep 2026 13:40:46 -0700 Subject: [PATCH 3/3] fix(core): gate import-time root logger config behind opt-in Addresses review feedback: get_textual_rich_logger() still called logging.basicConfig(level=INFO, handlers=[RichHandler(...)]) at import, which runs on a plain "import sagemaker.core" (via the module-level loggers in utils.py, resources.py, and intelligent_defaults_helper.py). That installed a RichHandler on the root logger and forced the root level to INFO for the whole process, defeating the goal of not touching global logging state on import. Gate the basicConfig/RichHandler behind is_rich_logging_enabled() (the same opt-in as the console/traceback). When not opted in, get_textual_rich_logger() just returns logging.getLogger(name) without configuring handlers or levels. Add tests asserting basicConfig is not called on the opt-out path and is called when opted in. Verified in a fresh process: default "import sagemaker.core" leaves the root logger level and handlers unchanged and sys.excepthook as the stdlib default. Relates-to #5179 Relates-to #5561 --- .../src/sagemaker/core/utils/utils.py | 21 ++++++++++++------- .../tests/unit/utils/test_rich_logging.py | 21 +++++++++++++++++++ 2 files changed, 34 insertions(+), 8 deletions(-) diff --git a/sagemaker-core/src/sagemaker/core/utils/utils.py b/sagemaker-core/src/sagemaker/core/utils/utils.py index f0d270edac..d807523935 100644 --- a/sagemaker-core/src/sagemaker/core/utils/utils.py +++ b/sagemaker-core/src/sagemaker/core/utils/utils.py @@ -194,23 +194,28 @@ def get_rich_handler(): def get_textual_rich_logger(name: str, log_level: str = "INFO") -> logging.Logger: - """Get a logger with textual rich handler. + """Get a logger, attaching a rich handler only when rich logging is opted in. + + Rich logging (a ``RichHandler`` on the root logger via ``logging.basicConfig``, + plus the themed console/traceback) is opt-in, so that importing the SDK does not + reconfigure the root logger or change process-wide log formatting/level. When the + user has not opted in (see :func:`is_rich_logging_enabled`), this returns the named + logger without configuring handlers or levels, leaving logging to the application. Args: name (str): The name of the logger - log_level (str): The log level to set. + log_level (str): The log level to set when rich logging is enabled. Accepted values are: "DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL". Defaults to the value of "INFO". Return: - logging.Logger: A textial rich logger. + logging.Logger: The requested logger. """ enable_textual_rich_console_and_traceback() - handler = get_rich_handler() - logging.basicConfig(level=getattr(logging, log_level), handlers=[handler]) - logger = logging.getLogger(name) - - return logger + if is_rich_logging_enabled(): + handler = get_rich_handler() + logging.basicConfig(level=getattr(logging, log_level), handlers=[handler]) + return logging.getLogger(name) logger = get_textual_rich_logger(__name__) diff --git a/sagemaker-core/tests/unit/utils/test_rich_logging.py b/sagemaker-core/tests/unit/utils/test_rich_logging.py index ba35d4d051..e875ea9ec1 100644 --- a/sagemaker-core/tests/unit/utils/test_rich_logging.py +++ b/sagemaker-core/tests/unit/utils/test_rich_logging.py @@ -91,3 +91,24 @@ def test_enable_is_idempotent_when_opted_in(): # The one-shot latch prevents re-installing on the second call. mock_reconfigure.assert_called_once() mock_install.assert_called_once() + + +def test_get_logger_does_not_call_basicconfig_when_opted_out(): + # Getting a module logger must not reconfigure the root logger by default. + with patch.dict(os.environ, {}, clear=False): + os.environ.pop(RICH_LOGGING_OPT_IN_ENV_VAR, None) + with patch.object(utils.logging, "basicConfig") as mock_basic_config: + returned = utils.get_textual_rich_logger("sagemaker.core.test.optout") + mock_basic_config.assert_not_called() + assert returned is utils.logging.getLogger("sagemaker.core.test.optout") + + +def test_get_logger_calls_basicconfig_when_opted_in(): + with patch.dict(os.environ, {RICH_LOGGING_OPT_IN_ENV_VAR: "1"}): + with ( + patch.object(utils, "reconfigure"), + patch.object(utils, "install"), + patch.object(utils.logging, "basicConfig") as mock_basic_config, + ): + utils.get_textual_rich_logger("sagemaker.core.test.optin") + mock_basic_config.assert_called_once()