litellm/tests/sdk_function_trace/profiler.py
yujonglee 198906495f
refactor(python-bridge): split routes and add shared function tracing (#39031)
* refactor(python-bridge): split non-streaming bridge modules

* refactor(python-bridge): bring shared function tracing into route layer

* feat(dev): list Python route functions and call sites

* feat(dev): list Rust route functions and call sites

* docs(dev): record OCR parity gaps across Python and Rust

* feat(dev): list executed SDK calls with runtime tracing

* feat(dev): report Python vs Rust SDK pipeline steps in one CLI

* feat(dev): side-by-side pipeline step report in compare CLI

* fix(dev): drop invalid Final annotations in compare cell loop

* feat(dev): blue python-only and yellow rust-only steps in compare CLI

* feat(dev): vertical layout with section spacing in compare CLI

* fix(dev): validate SDK trace stages across sync and async routes

* refactor(rust): align SDK route call structure with Python

* refactor(python-bridge): share sync and async route call wrappers

* refactor(dev): split compare CLI into fixtures, runtime, and report modules

* fix(ci): run SDK trace tests and satisfy test lint
2026-09-02 16:26:35 -07:00

76 lines
2.7 KiB
Python

from __future__ import annotations
import sys
import threading
from collections.abc import Generator, Sequence
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from types import CodeType, FrameType, FunctionType
from typing import Final
@dataclass(frozen=True, slots=True)
class FunctionTraceEvent:
function: str
depth: int
ancestors: tuple[str, ...] | None = None
class PythonProfiler:
def __init__(self, functions: Sequence[FunctionType], source_root: Path | None = None) -> None:
self._source_root: Final = str(source_root.resolve()) + "/" if source_root is not None else None
self._names_by_code: Final = {function.__code__: function.__name__ for function in functions}
self._seen_frames: Final[set[FrameType]] = set()
self.events: Final[list[FunctionTraceEvent]] = []
def __call__(self, frame: FrameType, event: str, _arg: object) -> None:
if event != "call" or frame in self._seen_frames:
return
function_name: Final = self.function_name(frame.f_code)
if function_name is None:
return
ancestors: Final = tuple(
name for ancestor in _frame_ancestors(frame) if (name := self.function_name(ancestor.f_code)) is not None
)
self._seen_frames.add(frame)
self.events.append(
FunctionTraceEvent(
function=function_name,
depth=len(ancestors),
ancestors=ancestors if self._source_root is not None else None,
)
)
def function_name(self, code: CodeType) -> str | None:
if self._source_root is None:
return self._names_by_code.get(code)
if not code.co_filename.startswith(self._source_root):
return None
relative: Final = code.co_filename.removeprefix(self._source_root)
return f"{relative}:{code.co_firstlineno} {getattr(code, 'co_qualname', code.co_name)}"
def _frame_ancestors(frame: FrameType) -> Generator[FrameType]:
ancestor: Final = frame.f_back
if ancestor is not None:
yield ancestor
yield from _frame_ancestors(ancestor)
@contextmanager
def profile_python(
functions: Sequence[FunctionType] = (), *, source_root: Path | None = None, threads: bool = False
) -> Generator[PythonProfiler]:
profiler: Final = PythonProfiler(functions, source_root)
previous_thread: Final = threading.getprofile()
if threads:
threading.setprofile(profiler)
previous: Final = sys.getprofile()
sys.setprofile(profiler)
try:
yield profiler
finally:
sys.setprofile(previous)
if threads:
threading.setprofile(previous_thread)