Importing sagemaker seems to enable rich tracebacks

[paw-lu]

Describe the bug

Whenever sagemaker is imported, it seems to modify and automatically enable rich tracebacks.

To reproduce

Make a new environment

$ python -m venv .venv
$ source /.venv/bin/activate
$ python -m pip install sagemaker rich

Make a file that returns a traceback

"""rich_sage.py"""
import sagemaker


def main() -> None:
    raise ValueError


if __name__ == "__main__":
    main()

Run the file

$ python rich_sage.py
sagemaker.config INFO - Not applying SDK defaults from location: /Library/Application Support/sagemaker/config.yaml
sagemaker.config INFO - Not applying SDK defaults from location: /Users/pawlu/Library/Application Support/sagemaker/config.yaml
╭─────────────────────────────── Traceback (most recent call last) ────────────────────────────────╮
│ /Users/pawlu/Documents/scratch/rich_sage/sage_me.py:9 in <module>                                │
│                                                                                                  │
│    6                                                                                             │
│    7                                                                                             │
│    8 if __name__ == "__main__":                                                                  │
│ ❱  9 │   main()                                                                                  │
│   10                                                                                             │
│                                                                                                  │
│ /Users/pawlu/Documents/scratch/rich_sage/sage_me.py:5 in main                                    │
│                                                                                                  │
│    2                                                                                             │
│    3                                                                                             │
│    4 def main() -> None:                                                                         │
│ ❱  5 │   raise ValueError                                                                        │
│    6                                                                                             │
│    7                                                                                             │
│    8 if __name__ == "__main__":                                                                  │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
ValueError

If the import statement is removed,

- import sagemaker

then the traceback renders as expected

$ python rich_sage.py
Traceback (most recent call last):
  File "/Users/pawlu/Documents/scratch/rich_sage/sage_me.py", line 8, in <module>
    main()
  File "/Users/pawlu/Documents/scratch/rich_sage/sage_me.py", line 4, in main
    raise ValueError
ValueError

Expected behavior

Ideally, tracebacks shouldn't be modified without users input. At worst there should be a way to opt out of this. But we shouldn't be modifying tracebacks automatically at import. It's a tricky behavior that's hard to debug.

Screenshots or logs

System information A description of your system. Please provide:

• SageMaker Python SDK version: 2.244.0
• Framework name (eg. PyTorch) or algorithm (eg. KMeans): Rich
• Framework version: 14.0.0
• Python version: 3.12.9
• CPU or GPU: M3 (macbook)
• Custom Docker image (Y/N): N

---

Thank you very much!

[cosmic-heart]

Is there any workaround to turn off rich? can you please make rich traceback logs optional?

[paw-lu]

hey @pintaoz-aws —thanks for labeling this issue

to be honest I view this as a bug though, not a feature request. ideally a third party dependency shouldn't be modifying the output traceback on import, and in cases this can make logging completely unreadable

[mollyheamazon]

Hi all, there are two workaround choices I can provide:

1. Local fix (per-script): Add the following when you are importing sagemaker in your rich_sage.py:

import sagemaker
import sys
sys.excepthook = sys.__excepthook__

2. Global fix (per-environment): In your virtual‑env’s site‑packages/ directory (or anywhere on your PYTHONPATH), create a file called:

 …/site‑packages/sitecustomize.py

put the following in sitecustomize.py:

import sys

try:
    import rich.traceback
    rich.traceback.install = lambda *args, **kwargs: None
except ImportError:
    pass

sys.excepthook = sys.__excepthook__

Let me know if you have further questions.

[tadamcz]

Thank you so much @paw-lu! I spend days trying to debug this in a larger application 🤯

@mollyheamazon What about fixing the bug? It's egregious behaviour that sagemaker does this at import-time!

[tadamcz]

Oh wait, does https://github.com/aws/sagemaker-core/pull/325 close this?