diff --git a/gateway/routes/allowlist.py b/gateway/routes/allowlist.py index 92b73867e67..65a0a259111 100644 --- a/gateway/routes/allowlist.py +++ b/gateway/routes/allowlist.py @@ -109,6 +109,7 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = ( # Health & ops "/health", "/metrics", + "/debug/asyncio-tasks", "/watsonx", ) diff --git a/litellm/proxy/common_utils/debug_utils.py b/litellm/proxy/common_utils/debug_utils.py index 1d4024bb84e..a4a34d8ad1f 100644 --- a/litellm/proxy/common_utils/debug_utils.py +++ b/litellm/proxy/common_utils/debug_utils.py @@ -1,26 +1,143 @@ # Start tracing memory allocations import asyncio import gc +import itertools import json import os import sys import tracemalloc from collections import Counter -from collections.abc import Mapping, Sequence -from typing import Any, Final, NamedTuple, Protocol, TypedDict +from collections.abc import Iterator, Mapping, Sequence +from types import AsyncGeneratorType, CoroutineType, FrameType, GeneratorType +from typing import Any, Final, NamedTuple, Protocol, TypeAlias from fastapi import APIRouter, Depends, HTTPException, Query -from typing_extensions import ReadOnly +from typing_extensions import ReadOnly, TypedDict from litellm import get_secret_str from litellm._logging import verbose_proxy_logger from litellm.constants import PYTHON_GC_THRESHOLD -from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy._types import UserAPIKeyAuth, user_api_key_has_admin_view from litellm.proxy.auth.user_api_key_auth import user_api_key_auth router: Final = APIRouter() +class _Frame(TypedDict): + file: ReadOnly[str] + line: ReadOnly[int] + function: ReadOnly[str] + + +class _TaskStackGroup(TypedDict): + count: ReadOnly[int] + coroutine: ReadOnly[str] + task_names: ReadOnly[tuple[str, ...]] + stack: ReadOnly[tuple[_Frame, ...]] + + +class _TaskStackDump(TypedDict): + worker_pid: ReadOnly[int] + total_active_tasks: ReadOnly[int] + groups: ReadOnly[tuple[_TaskStackGroup, ...]] + + +_TaskStackKey: TypeAlias = tuple[tuple[str, int, str], ...] +_TaskStackRecord: TypeAlias = tuple[_TaskStackKey, tuple[_Frame, ...], asyncio.Task[object]] + + +def _frame_from_stack_frame(frame: FrameType) -> _Frame: + result: Final[_Frame] = { + "file": frame.f_code.co_filename, + "line": frame.f_lineno, + "function": frame.f_code.co_name, + } + return result + + +def _awaitable_frame(awaitable: object) -> FrameType | None: + if isinstance(awaitable, CoroutineType): + return awaitable.cr_frame + if isinstance(awaitable, GeneratorType): + return awaitable.gi_frame + if isinstance(awaitable, AsyncGeneratorType): + return awaitable.ag_frame + return None + + +def _awaited(awaitable: object) -> object | None: + if isinstance(awaitable, CoroutineType): + return awaitable.cr_await + if isinstance(awaitable, GeneratorType): + return awaitable.gi_yieldfrom + if isinstance(awaitable, AsyncGeneratorType): + return awaitable.ag_await + return None + + +def _next_awaitable(awaitable: object, _: object) -> object | None: + return _awaited(awaitable) + + +def _awaitable_frame_tuple(awaitable: object) -> tuple[FrameType, ...]: + frame: Final = _awaitable_frame(awaitable) + return (frame,) if frame is not None else () + + +def _awaited_chain(awaitable: object) -> Iterator[FrameType]: + awaitables: Final = itertools.takewhile( + lambda current: _awaitable_frame(current) is not None, + itertools.accumulate(itertools.repeat(None), _next_awaitable, initial=awaitable), + ) + return itertools.chain.from_iterable(_awaitable_frame_tuple(current) for current in awaitables) + + +def _task_stack(task: asyncio.Task[object], max_frames: int) -> tuple[_Frame, ...]: + awaited_frames: Final = tuple(itertools.islice(_awaited_chain(task.get_coro()), max_frames)) + frames: Final = awaited_frames or tuple(task.get_stack(limit=max_frames)) + return tuple(_frame_from_stack_frame(frame) for frame in frames) + + +def _task_stack_key(stack: tuple[_Frame, ...]) -> _TaskStackKey: + return tuple((frame["file"], frame["line"], frame["function"]) for frame in stack) + + +def _task_coroutine_name(task: asyncio.Task[object]) -> str: + coroutine: Final = task.get_coro() + coroutine_name: Final = getattr(coroutine, "__qualname__", None) + return coroutine_name if isinstance(coroutine_name, str) else repr(coroutine) + + +def _task_stack_group(records: tuple[_TaskStackRecord, ...]) -> _TaskStackGroup: + sample_record: Final = records[0] + sample_tasks: Final = tuple(record[2] for record in records) + result: Final[_TaskStackGroup] = { + "count": len(records), + "coroutine": _task_coroutine_name(sample_tasks[0]), + "task_names": tuple(task.get_name() for task in sample_tasks[:5]), + "stack": sample_record[1], + } + return result + + +def _group_task_stacks(tasks: tuple[asyncio.Task[object], ...], max_frames: int) -> tuple[_TaskStackGroup, ...]: + records: Final = tuple( + sorted( + ( + (stack_key, stack, task) + for task in tasks + for stack in (_task_stack(task, max_frames),) + for stack_key in (_task_stack_key(stack),) + ), + key=lambda record: record[0], + ) + ) + groups: Final = tuple( + _task_stack_group(tuple(group)) for _, group in itertools.groupby(records, key=lambda record: record[0]) + ) + return tuple(sorted(groups, key=lambda group: group["count"], reverse=True)) + + # Configure garbage collection thresholds from environment variables def configure_gc_thresholds(): """Configure Python garbage collection thresholds from environment variables.""" @@ -87,6 +204,27 @@ async def get_active_tasks_stats(): } +@router.get("/debug/asyncio-tasks/stacks", include_in_schema=False) +async def get_active_task_stacks( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + max_frames: int = Query(default=40, ge=1, le=200), +) -> _TaskStackDump: + if not user_api_key_has_admin_view(user_api_key_dict): + raise HTTPException(status_code=403, detail="Only proxy admins can read asyncio task stacks") + + max_tasks_to_check: Final = 5000 + active_tasks: Final = tuple(task for task in asyncio.all_tasks() if not task.done()) + current_task: Final = asyncio.current_task() + tasks: Final = tuple(task for task in active_tasks if task is not current_task)[:max_tasks_to_check] + groups: Final = _group_task_stacks(tasks, max_frames) + result: Final[_TaskStackDump] = { + "worker_pid": os.getpid(), + "total_active_tasks": len(active_tasks), + "groups": groups, + } + return result + + if os.environ.get("LITELLM_PROFILE", "false").lower() == "true": try: import objgraph diff --git a/tests/test_litellm/proxy/common_utils/test_debug_utils.py b/tests/test_litellm/proxy/common_utils/test_debug_utils.py new file mode 100644 index 00000000000..75b2950a9d9 --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_debug_utils.py @@ -0,0 +1,113 @@ +import asyncio + +import httpx +import pytest +from fastapi import FastAPI + +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.debug_utils import router as debug_router + + +async def _park_for_test() -> None: + await asyncio.sleep(30) + + +async def _park_inner(event: asyncio.Event) -> None: + await event.wait() + + +async def _park_outer(event: asyncio.Event) -> None: + await _park_inner(event) + + +def _test_app(user_role: LitellmUserRoles) -> FastAPI: + app = FastAPI() + app.include_router(debug_router) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=user_role) + return app + + +def _client(app: FastAPI) -> httpx.AsyncClient: + return httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://testserver") + + +@pytest.mark.asyncio +async def test_task_stacks_require_proxy_admin() -> None: + app = _test_app(LitellmUserRoles.INTERNAL_USER) + async with _client(app) as client: + response = await client.get("/debug/asyncio-tasks/stacks") + + assert response.status_code == 403 + assert response.json()["detail"] == "Only proxy admins can read asyncio task stacks" + + +@pytest.mark.asyncio +async def test_task_stacks_allow_proxy_admin_view_only() -> None: + app = _test_app(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY) + async with _client(app) as client: + response = await client.get("/debug/asyncio-tasks/stacks") + + assert response.status_code == 200 + assert "groups" in response.json() + + +@pytest.mark.asyncio +async def test_task_stacks_include_parked_task() -> None: + app = _test_app(LitellmUserRoles.PROXY_ADMIN) + parked_task = asyncio.create_task(_park_for_test(), name="parked-test-task") + try: + async with _client(app) as client: + response = await client.get("/debug/asyncio-tasks/stacks") + finally: + parked_task.cancel() + await asyncio.gather(parked_task, return_exceptions=True) + + assert response.status_code == 200 + body = response.json() + parked_group = next(group for group in body["groups"] if "_park_for_test" in group["coroutine"]) + assert any(frame["function"] == "_park_for_test" for frame in parked_group["stack"]) + assert any(frame["file"].endswith("test_debug_utils.py") for frame in parked_group["stack"]) + assert body["total_active_tasks"] >= 1 + + +@pytest.mark.asyncio +async def test_task_stacks_include_nested_await_frames() -> None: + app = _test_app(LitellmUserRoles.PROXY_ADMIN) + event = asyncio.Event() + parked_task = asyncio.create_task(_park_outer(event), name="nested-parked-test-task") + await asyncio.sleep(0) + try: + async with _client(app) as client: + response = await client.get("/debug/asyncio-tasks/stacks") + finally: + parked_task.cancel() + await asyncio.gather(parked_task, return_exceptions=True) + + assert response.status_code == 200 + nested_group = next( + group for group in response.json()["groups"] if group["task_names"] == ["nested-parked-test-task"] + ) + frames = nested_group["stack"] + functions = [frame["function"] for frame in frames] + assert functions.index("_park_outer") < functions.index("_park_inner") + assert any(frame["file"].endswith("asyncio/locks.py") or frame["function"] == "wait" for frame in frames) + + +@pytest.mark.asyncio +async def test_task_stacks_respect_max_frames() -> None: + app = _test_app(LitellmUserRoles.PROXY_ADMIN) + parked_task = asyncio.create_task(_park_for_test(), name="parked-test-task") + try: + async with _client(app) as client: + response = await client.get("/debug/asyncio-tasks/stacks?max_frames=1") + finally: + parked_task.cancel() + await asyncio.gather(parked_task, return_exceptions=True) + + assert response.status_code == 200 + groups = response.json()["groups"] + parked_group = next(group for group in groups if "_park_for_test" in group["coroutine"]) + assert len(parked_group["stack"]) == 1 + assert parked_group["stack"][0]["function"] == "_park_for_test" + assert all(len(group["stack"]) <= 1 for group in groups) diff --git a/tests/test_litellm/proxy/test_component_allowlists.py b/tests/test_litellm/proxy/test_component_allowlists.py index 0fdb43d60da..8618f2f4438 100644 --- a/tests/test_litellm/proxy/test_component_allowlists.py +++ b/tests/test_litellm/proxy/test_component_allowlists.py @@ -26,6 +26,10 @@ RDS IAM token when ``IAM_TOKEN_DB_AUTH`` is set). import os import sys +import httpx +import pytest +from fastapi import FastAPI + # Importing ``litellm.proxy.proxy_server`` runs its module-level setup, which # reads ``DATABASE_URL`` (Prisma) and ``LITELLM_MASTER_KEY``. Tier-zero CI # runners don't set these. We pin throwaway values before the import so the @@ -59,6 +63,8 @@ from gateway.routes.allowlist import ( GATEWAY_MOUNT_PATHS, GATEWAY_PATH_PREFIXES, ) +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.proxy_server import app for _key, _previous in _PRE_EXISTING_ENV.items(): @@ -196,6 +202,34 @@ def test_gateway_drops_ui_and_swagger_mounts(): f"Mount {path} must not be served by the gateway" +@pytest.mark.asyncio +async def test_gateway_serves_asyncio_task_debug_routes() -> None: + previous_override = app.dependency_overrides.get(user_api_key_auth) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) + trimmed = FastAPI(routes=[route for route in app.router.routes if _is_gateway_route(route)]) + management_route = next(route for route in app.router.routes if getattr(route, "path", None) == "/key/info") + try: + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=trimmed), + base_url="http://testserver", + ) as client: + stacks_response = await client.get("/debug/asyncio-tasks/stacks") + count_response = await client.get("/debug/asyncio-tasks") + management_response = await client.get("/key/info") + finally: + if previous_override is None: + app.dependency_overrides.pop(user_api_key_auth, None) + else: + app.dependency_overrides[user_api_key_auth] = previous_override + + assert stacks_response.status_code == 200 + assert stacks_response.json()["worker_pid"] + assert "groups" in stacks_response.json() + assert count_response.status_code == 200 + assert not _is_gateway_route(management_route) + assert management_response.status_code == 404 + + def test_every_app_mount_is_assigned_to_a_component(): """Every Mount on the proxy app must be consciously assigned to a component. diff --git a/tests/test_litellm/proxy/test_sensitive_route_auth.py b/tests/test_litellm/proxy/test_sensitive_route_auth.py index 19998e52779..e420041da45 100644 --- a/tests/test_litellm/proxy/test_sensitive_route_auth.py +++ b/tests/test_litellm/proxy/test_sensitive_route_auth.py @@ -21,6 +21,7 @@ def _get_route_dependency_calls(router, path: str, method: str): def test_sensitive_debug_routes_require_auth_dependency(): for path, method in ( ("/debug/asyncio-tasks", "GET"), + ("/debug/asyncio-tasks/stacks", "GET"), ("/otel-spans", "GET"), ): assert user_api_key_auth in _get_route_dependency_calls( diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index dd455cba44a..0df232f69af 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -4067,6 +4067,23 @@ export interface paths { patch?: never; trace?: never; }; + "/debug/asyncio-tasks/stacks": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** Get Active Task Stacks */ + get: operations["get_active_task_stacks_debug_asyncio_tasks_stacks_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/debug/memory/details": { parameters: { query?: never; @@ -39578,6 +39595,35 @@ export interface components { /** Status */ status?: ("pending" | "running" | "paused" | "completed" | "failed") | null; }; + /** _Frame */ + _Frame: { + /** File */ + file: string; + /** Function */ + function: string; + /** Line */ + line: number; + }; + /** _TaskStackDump */ + _TaskStackDump: { + /** Groups */ + groups: components["schemas"]["_TaskStackGroup"][]; + /** Total Active Tasks */ + total_active_tasks: number; + /** Worker Pid */ + worker_pid: number; + }; + /** _TaskStackGroup */ + _TaskStackGroup: { + /** Coroutine */ + coroutine: string; + /** Count */ + count: number; + /** Stack */ + stack: components["schemas"]["_Frame"][]; + /** Task Names */ + task_names: string[]; + }; /** ModelInfo */ litellm__proxy___types__ModelInfo: { /** Base Model */ @@ -45851,6 +45897,37 @@ export interface operations { }; }; }; + get_active_task_stacks_debug_asyncio_tasks_stacks_get: { + parameters: { + query?: { + max_frames?: number; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["_TaskStackDump"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; get_memory_details_debug_memory_details_get: { parameters: { query?: {