allow seeing warnings with traces optionally

This commit is contained in:
Ashwin Bharambe 2024-11-13 12:27:19 -08:00
parent 96e7ef646f
commit 7f6ac2fbd7

View file

@ -13,6 +13,7 @@ import re
import signal import signal
import sys import sys
import traceback import traceback
import warnings
from contextlib import asynccontextmanager from contextlib import asynccontextmanager
from ssl import SSLError from ssl import SSLError
@ -45,6 +46,16 @@ from llama_stack.distribution.stack import construct_stack
from .endpoints import get_all_api_endpoints from .endpoints import get_all_api_endpoints
def warn_with_traceback(message, category, filename, lineno, file=None, line=None):
log = file if hasattr(file, "write") else sys.stderr
traceback.print_stack(file=log)
log.write(warnings.formatwarning(message, category, filename, lineno, line))
if os.environ.get("LLAMA_STACK_TRACE_WARNINGS"):
warnings.showwarning = warn_with_traceback
def create_sse_event(data: Any) -> str: def create_sse_event(data: Any) -> str:
if isinstance(data, BaseModel): if isinstance(data, BaseModel):
data = data.model_dump_json() data = data.model_dump_json()