mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
fae3d224eb
# Conflicts: # basedpyright-code-budget.json # tests/sdk_function_trace/profiler.py # tests/sdk_function_trace/test_profiler.py
141 lines
3.9 KiB
Python
141 lines
3.9 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import sys
|
|
import threading
|
|
from collections.abc import Callable
|
|
from functools import wraps
|
|
from pathlib import Path
|
|
from types import FunctionType
|
|
from typing import Final, ParamSpec, TypeVar, cast
|
|
|
|
import pytest
|
|
|
|
from .profiler import (
|
|
FunctionTraceEvent,
|
|
PythonProfiler,
|
|
_module_qualnames,
|
|
profile_python,
|
|
profile_python_function_usage,
|
|
)
|
|
|
|
_P = ParamSpec("_P")
|
|
_T = TypeVar("_T")
|
|
|
|
|
|
def _passthrough(function: Callable[_P, _T]) -> Callable[_P, _T]:
|
|
@wraps(function)
|
|
def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _T:
|
|
return function(*args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
|
|
class Decorated:
|
|
@_passthrough
|
|
def call(self) -> None:
|
|
return None
|
|
|
|
|
|
def _events_named(profiler: PythonProfiler, name: str) -> tuple[FunctionTraceEvent, ...]:
|
|
return tuple(event for event in profiler.events if event.function.endswith(name))
|
|
|
|
|
|
def test_profiler_keeps_repeated_calls() -> None:
|
|
def called() -> None:
|
|
return None
|
|
|
|
with profile_python(Path(__file__).parent) as profiler:
|
|
called()
|
|
called()
|
|
|
|
assert len(_events_named(profiler, "called")) == 2
|
|
|
|
|
|
def test_profiler_qualifies_decorated_methods_by_class() -> None:
|
|
with profile_python(Path(__file__).parent) as profiler:
|
|
Decorated().call()
|
|
|
|
assert any(event.function.endswith(" Decorated.call") for event in profiler.events)
|
|
assert _module_qualnames(__name__)[cast(FunctionType, Decorated.call.__wrapped__).__code__] == "Decorated.call"
|
|
|
|
|
|
def test_profiler_records_real_frame_ancestry() -> None:
|
|
def called() -> None:
|
|
return None
|
|
|
|
def outer() -> None:
|
|
called()
|
|
|
|
with profile_python(Path(__file__).parent) as profiler:
|
|
outer()
|
|
|
|
outer_event, called_event = (event for event in profiler.events if event.function.endswith(("outer", "called")))
|
|
assert called_event.parent_id == outer_event.id
|
|
|
|
|
|
def test_profiler_restores_previous_profiler_after_failure() -> None:
|
|
previous: Final = sys.getprofile()
|
|
|
|
with pytest.raises(RuntimeError, match="stop"):
|
|
with profile_python(Path(__file__).parent):
|
|
raise RuntimeError("stop")
|
|
|
|
assert sys.getprofile() is previous
|
|
|
|
|
|
def test_profiler_does_not_count_coroutine_resumption_as_another_call() -> None:
|
|
async def suspended() -> None:
|
|
await asyncio.sleep(0)
|
|
await asyncio.sleep(0)
|
|
|
|
with profile_python(Path(__file__).parent) as profiler:
|
|
asyncio.run(suspended())
|
|
|
|
assert len(_events_named(profiler, "suspended")) == 1
|
|
|
|
|
|
def test_profiler_preserves_parent_across_coroutine_suspension() -> None:
|
|
def called() -> None:
|
|
return None
|
|
|
|
async def suspended() -> None:
|
|
await asyncio.sleep(0)
|
|
called()
|
|
|
|
with profile_python(Path(__file__).parent) as profiler:
|
|
asyncio.run(suspended())
|
|
|
|
suspended_event: Final = _events_named(profiler, "suspended")[0]
|
|
called_event: Final = _events_named(profiler, "called")[0]
|
|
assert called_event.parent_id == suspended_event.id
|
|
|
|
|
|
def test_profiler_captures_worker_threads_when_enabled() -> None:
|
|
def called() -> None:
|
|
return None
|
|
|
|
with profile_python(Path(__file__).parent, threads=True) as profiler:
|
|
thread: Final = threading.Thread(target=called)
|
|
thread.start()
|
|
thread.join()
|
|
|
|
called_event: Final = _events_named(profiler, "called")[0]
|
|
assert called_event.parent_id is None
|
|
|
|
|
|
def test_function_usage_profiler_records_only_selected_functions() -> None:
|
|
def selected() -> None:
|
|
return None
|
|
|
|
def ignored() -> None:
|
|
return None
|
|
|
|
source_root: Final = Path(__file__).parent
|
|
function: Final = f"{Path(__file__).name}:{selected.__code__.co_firstlineno} {selected.__qualname__}"
|
|
|
|
with profile_python_function_usage(source_root, frozenset((function,))) as profiler:
|
|
selected()
|
|
ignored()
|
|
|
|
assert profiler.called == {function}
|