From ad5b87eb9149de28b39cd1b1980fa093f408d673 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 14 Sep 2026 11:19:28 -0700 Subject: [PATCH] fix: capture reused worker threads in Python traces --- .../shared/tracing/profiler.py | 68 +++++++---- .../shared/tracing/test_profiler.py | 113 ++++++++++++++++-- 2 files changed, 150 insertions(+), 31 deletions(-) diff --git a/tests/rust-python-harness/shared/tracing/profiler.py b/tests/rust-python-harness/shared/tracing/profiler.py index abfb6a2425d..e72a36a5289 100644 --- a/tests/rust-python-harness/shared/tracing/profiler.py +++ b/tests/rust-python-harness/shared/tracing/profiler.py @@ -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) diff --git a/tests/rust-python-harness/shared/tracing/test_profiler.py b/tests/rust-python-harness/shared/tracing/test_profiler.py index 616e9c23e75..b9a07f249a3 100644 --- a/tests/rust-python-harness/shared/tracing/test_profiler.py +++ b/tests/rust-python-harness/shared/tracing/test_profiler.py @@ -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