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..d807523935 100644 --- a/sagemaker-core/src/sagemaker/core/utils/utils.py +++ b/sagemaker-core/src/sagemaker/core/utils/utils.py @@ -137,12 +137,47 @@ 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) @@ -159,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 new file mode 100644 index 0000000000..e875ea9ec1 --- /dev/null +++ b/sagemaker-core/tests/unit/utils/test_rich_logging.py @@ -0,0 +1,114 @@ +"""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() + + +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()