scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-15 23:31:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 165

fix: capture reused worker threads in Python traces

Browse source
This commit is contained in:
Yujong Lee 2026-09-14 11:19:28 -07:00
parent cab1e113f7
commit ad5b87eb91
2 changed files with 150 additions and 31 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

68
tests/rust-python-harness/shared/tracing/profiler.py
View file

@ -2,7 +2,8 @@ from __future__ import annotations
import sys
import threading
from collections.abc import Generator, Iterator, Mapping
import warnings
from collections.abc import Callable, Generator, Iterator, Mapping
from contextlib import contextmanager
from dataclasses import dataclass
from functools import lru_cache
@ -32,6 +33,7 @@ class PythonProfiler:
self._source_root: Final = str(source_root.resolve()) + "/"
self._seen_frames: Final[set[FrameType]] = set()
self._event_ids: Final[dict[FrameType, int]] = {}
self._lock: Final = threading.Lock()
self.events: Final[list[FunctionTraceEvent]] = []
def __call__(self, frame: FrameType, event: str, _arg: object) -> None:
@ -40,14 +42,15 @@ class PythonProfiler:
function_name: Final = self.function_name(frame)
if function_name is None:
return
event_id: Final = len(self.events)
parent_id: Final = next(
(self._event_ids[ancestor] for ancestor in _frame_ancestors(frame) if ancestor in self._event_ids),
None,
)
self._seen_frames.add(frame)
self._event_ids[frame] = event_id
self.events.append(FunctionTraceEvent(id=event_id, parent_id=parent_id, function=function_name))
with self._lock:
event_id: Final = len(self.events)
parent_id: Final = next(
(self._event_ids[ancestor] for ancestor in _frame_ancestors(frame) if ancestor in self._event_ids),
None,
)
self._seen_frames.add(frame)
self._event_ids[frame] = event_id
self.events.append(FunctionTraceEvent(id=event_id, parent_id=parent_id, function=function_name))
def function_name(self, frame: FrameType) -> str | None:
code: Final = frame.f_code
@ -137,21 +140,51 @@ def _frame_ancestors(frame: FrameType) -> Generator[FrameType]:
@contextmanager
def profile_python(source_root: Path, *, threads: bool = False) -> Generator[PythonProfiler]:
profiler: Final = PythonProfiler(source_root)
def _installed_profiler(profiler: Callable[[FrameType, str, object], None], *, threads: bool) -> Generator[None]:
if threads and sys.version_info >= (3, 12):
tool_id: Final = next((slot for slot in (2, 3, 4, 0, 1, 5) if sys.monitoring.get_tool(slot) is None), None)
if tool_id is None:
raise RuntimeError("no sys.monitoring tool ID is available for Python trace collection")
def started(_code: CodeType, _offset: int) -> None:
profiler(sys._getframe(1), "call", None)
sys.monitoring.use_tool_id(tool_id, "litellm-python-trace")
try:
sys.monitoring.register_callback(tool_id, sys.monitoring.events.PY_START, started)
sys.monitoring.set_events(tool_id, sys.monitoring.events.PY_START)
yield
finally:
sys.monitoring.set_events(tool_id, 0)
sys.monitoring.register_callback(tool_id, sys.monitoring.events.PY_START, None)
sys.monitoring.free_tool_id(tool_id)
return
if threads:
warnings.warn(
"Python <3.12 cannot trace existing worker threads; use Python 3.12+ for complete threaded traces",
RuntimeWarning,
stacklevel=3,
)
previous_thread: Final = threading.getprofile()
if threads:
threading.setprofile(profiler)
previous: Final = sys.getprofile()
sys.setprofile(profiler)
try:
yield profiler
yield
finally:
sys.setprofile(previous)
if threads:
threading.setprofile(previous_thread)
@contextmanager
def profile_python(source_root: Path, *, threads: bool = False) -> Generator[PythonProfiler]:
profiler: Final = PythonProfiler(source_root)
with _installed_profiler(profiler, threads=threads):
yield profiler
@contextmanager
def profile_python_function_usage(
source_root: Path,
@ -160,14 +193,5 @@ def profile_python_function_usage(
threads: bool = False,
) -> Generator[PythonFunctionUsageProfiler]:
profiler: Final = PythonFunctionUsageProfiler(source_root, functions)
previous_thread: Final = threading.getprofile()
if threads:
threading.setprofile(profiler)
previous: Final = sys.getprofile()
sys.setprofile(profiler)
try:
with _installed_profiler(profiler, threads=threads):
yield profiler
finally:
sys.setprofile(previous)
if threads:
threading.setprofile(previous_thread)

113
tests/rust-python-harness/shared/tracing/test_profiler.py
View file

@ -4,9 +4,10 @@ import asyncio
import sys
import threading
from collections.abc import Callable
from concurrent.futures import ThreadPoolExecutor
from functools import wraps
from pathlib import Path
from types import FunctionType
from types import FrameType, FunctionType
from typing import Final, ParamSpec, TypeVar, cast
import pytest
@ -41,11 +42,12 @@ def _events_named(profiler: PythonProfiler, name: str) -> tuple[FunctionTraceEve
return tuple(event for event in profiler.events if event.function.endswith(name))
def test_profiler_keeps_repeated_calls() -> None:
@pytest.mark.parametrize("threads", (False, True))
def test_profiler_keeps_repeated_calls(threads: bool) -> None:
def called() -> None:
return None
with profile_python(Path(__file__).parent) as profiler:
with profile_python(Path(__file__).parent, threads=threads) as profiler:
called()
called()
@ -60,14 +62,15 @@ def test_profiler_qualifies_decorated_methods_by_class() -> None:
assert _module_qualnames(__name__)[cast(FunctionType, Decorated.call.__wrapped__).__code__] == "Decorated.call"
def test_profiler_records_real_frame_ancestry() -> None:
@pytest.mark.parametrize("threads", (False, True))
def test_profiler_records_real_frame_ancestry(threads: bool) -> None:
def called() -> None:
return None
def outer() -> None:
called()
with profile_python(Path(__file__).parent) as profiler:
with profile_python(Path(__file__).parent, threads=threads) as profiler:
outer()
outer_event, called_event = (event for event in profiler.events if event.function.endswith(("outer", "called")))
@ -84,18 +87,20 @@ def test_profiler_restores_previous_profiler_after_failure() -> None:
assert sys.getprofile() is previous
def test_profiler_does_not_count_coroutine_resumption_as_another_call() -> None:
@pytest.mark.parametrize("threads", (False, True))
def test_profiler_does_not_count_coroutine_resumption_as_another_call(threads: bool) -> None:
async def suspended() -> None:
await asyncio.sleep(0)
await asyncio.sleep(0)
with profile_python(Path(__file__).parent) as profiler:
with profile_python(Path(__file__).parent, threads=threads) as profiler:
asyncio.run(suspended())
assert len(_events_named(profiler, "suspended")) == 1
def test_profiler_preserves_parent_across_coroutine_suspension() -> None:
@pytest.mark.parametrize("threads", (False, True))
def test_profiler_preserves_parent_across_coroutine_suspension(threads: bool) -> None:
def called() -> None:
return None
@ -103,7 +108,7 @@ def test_profiler_preserves_parent_across_coroutine_suspension() -> None:
await asyncio.sleep(0)
called()
with profile_python(Path(__file__).parent) as profiler:
with profile_python(Path(__file__).parent, threads=threads) as profiler:
asyncio.run(suspended())
suspended_event: Final = _events_named(profiler, "suspended")[0]
@ -124,6 +129,96 @@ def test_profiler_captures_worker_threads_when_enabled() -> None:
assert called_event.parent_id is None
@pytest.mark.skipif(sys.version_info < (3, 12), reason="existing worker capture requires sys.monitoring")
@pytest.mark.parametrize("prewarm", (False, True))
def test_profiler_captures_reused_workers_without_leaking_between_sessions(prewarm: bool) -> None:
def called() -> None:
return None
with ThreadPoolExecutor(max_workers=1) as executor:
if prewarm:
executor.submit(called).result(timeout=5)
with profile_python(Path(__file__).parent, threads=True) as first:
executor.submit(called).result(timeout=5)
executor.submit(called).result(timeout=5)
with profile_python(Path(__file__).parent, threads=True) as second:
executor.submit(called).result(timeout=5)
executor.submit(called).result(timeout=5)
assert len(_events_named(first, "called")) == 1
assert len(_events_named(second, "called")) == 1
def test_profiler_restores_main_and_worker_hooks_after_failure() -> None:
previous: Final = sys.getprofile()
previous_thread: Final = threading.getprofile()
with ThreadPoolExecutor(max_workers=1) as executor:
worker_previous: Final = executor.submit(sys.getprofile).result(timeout=5)
with pytest.raises(RuntimeError, match="stop"):
with profile_python(Path(__file__).parent, threads=True):
raise RuntimeError("stop")
assert executor.submit(sys.getprofile).result(timeout=5) is worker_previous
assert sys.getprofile() is previous
assert threading.getprofile() is previous_thread
@pytest.mark.skipif(sys.version_info < (3, 12), reason="existing worker capture requires sys.monitoring")
def test_function_usage_profiler_captures_reused_workers() -> None:
def selected() -> None:
return None
function: Final = f"{Path(__file__).name}:{selected.__code__.co_firstlineno} {selected.__qualname__}"
with ThreadPoolExecutor(max_workers=1) as executor:
executor.submit(selected).result(timeout=5)
with profile_python_function_usage(Path(__file__).parent, frozenset((function,)), threads=True) as profiler:
executor.submit(selected).result(timeout=5)
assert profiler.called == {function}
@pytest.mark.skipif(sys.version_info < (3, 12), reason="independent thread hooks require sys.monitoring")
def test_threaded_profiler_preserves_custom_worker_hook_and_releases_monitoring_slot() -> None:
def worker_hook(_frame: FrameType, _event: str, _arg: object) -> None:
return None
def fail_with_profile(executor: ThreadPoolExecutor) -> None:
with profile_python(Path(__file__).parent, threads=True):
assert executor.submit(sys.getprofile).result(timeout=5) is worker_hook
raise RuntimeError("stop")
tools_before: Final = tuple(sys.monitoring.get_tool(slot) for slot in range(6))
with ThreadPoolExecutor(max_workers=1, initializer=lambda: sys.setprofile(worker_hook)) as executor:
assert executor.submit(sys.getprofile).result(timeout=5) is worker_hook
with pytest.raises(RuntimeError, match="stop"):
fail_with_profile(executor)
assert executor.submit(sys.getprofile).result(timeout=5) is worker_hook
assert tuple(sys.monitoring.get_tool(slot) for slot in range(6)) == tools_before
@pytest.mark.skipif(sys.version_info < (3, 12), reason="existing worker capture requires sys.monitoring")
def test_threaded_profiler_keeps_concurrent_event_ids_and_parent_links() -> None:
def child() -> None:
return None
def parent() -> None:
child()
with ThreadPoolExecutor(max_workers=4) as executor:
with profile_python(Path(__file__).parent, threads=True) as profiler:
futures: Final = tuple(executor.submit(parent) for _ in range(200))
for future in futures:
future.result(timeout=5)
parent_ids: Final = frozenset(event.id for event in _events_named(profiler, "parent"))
children: Final = _events_named(profiler, "child")
assert len(parent_ids) == len(children) == 200
assert frozenset(event.parent_id for event in children) == parent_ids
assert tuple(event.id for event in profiler.events) == tuple(range(len(profiler.events)))
def test_function_usage_profiler_records_only_selected_functions() -> None:
def selected() -> None:
return None