refactor(proxy): group asyncio task stacks in one pass

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/proxy/common_utils/debug_utils.py

@@ -1,6 +1,7 @@
 # Start tracing memory allocations
 import asyncio
 import gc
+import itertools
 import json
 import os
 import sys
@@ -68,12 +69,11 @@ def _task_coroutine_name(task: asyncio.Task[object]) -> str:
     return coroutine_name if isinstance(coroutine_name, str) else repr(coroutine)
 
 
-def _task_stack_group(stack_key: _TaskStackKey, records: tuple[_TaskStackRecord, ...]) -> _TaskStackGroup:
-    matching_records: Final = tuple(record for record in records if record[0] == stack_key)
-    sample_record: Final = matching_records[0]
-    sample_tasks: Final = tuple(record[2] for record in matching_records)
+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(matching_records),
+        "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],
@@ -83,13 +83,19 @@ def _task_stack_group(stack_key: _TaskStackKey, records: tuple[_TaskStackRecord,
 
 def _group_task_stacks(tasks: tuple[asyncio.Task[object], ...], max_frames: int) -> tuple[_TaskStackGroup, ...]:
     records: Final = tuple(
-        (stack_key, stack, task)
-        for task in tasks
-        for stack in (_task_stack(task, max_frames),)
-        for stack_key in (_task_stack_key(stack),)
+        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])
     )
-    stack_keys: Final = tuple(dict.fromkeys(record[0] for record in records))
-    groups: Final = tuple(_task_stack_group(stack_key, records) for stack_key in stack_keys)
     return tuple(sorted(groups, key=lambda group: group["count"], reverse=True))
 
 

tests/test_litellm/proxy/test_component_allowlists.py

@@ -108,8 +108,12 @@ def test_gateway_plus_backend_covers_full_app():
         for r in app.router.routes
         if not isinstance(r, Mount) and getattr(r, "path", None) is not None
     }
-    gateway_paths = _component_paths(app.router.routes, GATEWAY_EXACT_PATHS, GATEWAY_PATH_PREFIXES)
-    backend_paths = _component_paths(app.router.routes, BACKEND_EXACT_PATHS, BACKEND_PATH_PREFIXES)
+    gateway_paths = _component_paths(
+        app.router.routes, GATEWAY_EXACT_PATHS, GATEWAY_PATH_PREFIXES
+    )
+    backend_paths = _component_paths(
+        app.router.routes, BACKEND_EXACT_PATHS, BACKEND_PATH_PREFIXES
+    )
 
     uncovered = all_paths - (gateway_paths | backend_paths)
 
@@ -122,15 +126,16 @@ def test_gateway_plus_backend_covers_full_app():
 
 def test_backend_mount_paths_defined():
     """BACKEND_MOUNT_PATHS constant must exist and be a frozenset."""
-    assert isinstance(BACKEND_MOUNT_PATHS, frozenset), (
+    assert isinstance(BACKEND_MOUNT_PATHS, frozenset), \
         f"BACKEND_MOUNT_PATHS must be a frozenset, got {type(BACKEND_MOUNT_PATHS)}"
-    )
-    assert len(BACKEND_MOUNT_PATHS) > 0, "BACKEND_MOUNT_PATHS must contain at least one Mount path"
+    assert len(BACKEND_MOUNT_PATHS) > 0, \
+        "BACKEND_MOUNT_PATHS must contain at least one Mount path"
 
 
 def test_swagger_mount_in_backend_allowlist():
     """The /swagger Mount must be in BACKEND_MOUNT_PATHS."""
-    assert "/swagger" in BACKEND_MOUNT_PATHS, "/swagger Mount path must be in BACKEND_MOUNT_PATHS"
+    assert "/swagger" in BACKEND_MOUNT_PATHS, \
+        "/swagger Mount path must be in BACKEND_MOUNT_PATHS"
 
 
 def test_backend_keeps_swagger_mount():
@@ -140,31 +145,32 @@ def test_backend_keeps_swagger_mount():
         for r in app.router.routes
         if isinstance(r, Mount) and getattr(r, "path", None) in BACKEND_MOUNT_PATHS
     }
-    assert "/swagger" in backend_mounts, (
+    assert "/swagger" in backend_mounts, \
         "/swagger Mount is expected on the proxy app and should be in BACKEND_MOUNT_PATHS"
-    )
 
 
 def test_backend_drops_non_allowlisted_mounts():
     """Verify that Mounts NOT in BACKEND_MOUNT_PATHS would be dropped from backend."""
     all_mounts = {
-        getattr(r, "path") for r in app.router.routes if isinstance(r, Mount) and getattr(r, "path", None) is not None
+        getattr(r, "path")
+        for r in app.router.routes
+        if isinstance(r, Mount) and getattr(r, "path", None) is not None
     }
     non_backend_mounts = all_mounts - BACKEND_MOUNT_PATHS
 
-    assert len(non_backend_mounts) > 0, (
+    assert len(non_backend_mounts) > 0, \
         "Expected at least one non-backend Mount (e.g., /ui, /_next) to verify filtering logic"
-    )
     for mount_path in non_backend_mounts:
-        assert mount_path not in BACKEND_MOUNT_PATHS, f"Mount {mount_path} should not be in BACKEND_MOUNT_PATHS"
+        assert mount_path not in BACKEND_MOUNT_PATHS, \
+            f"Mount {mount_path} should not be in BACKEND_MOUNT_PATHS"
 
 
 def test_gateway_mount_paths_defined():
     """GATEWAY_MOUNT_PATHS constant must exist and expose /metrics."""
-    assert isinstance(GATEWAY_MOUNT_PATHS, frozenset), (
+    assert isinstance(GATEWAY_MOUNT_PATHS, frozenset), \
         f"GATEWAY_MOUNT_PATHS must be a frozenset, got {type(GATEWAY_MOUNT_PATHS)}"
-    )
-    assert "/metrics" in GATEWAY_MOUNT_PATHS, "/metrics Mount path must be in GATEWAY_MOUNT_PATHS"
+    assert "/metrics" in GATEWAY_MOUNT_PATHS, \
+        "/metrics Mount path must be in GATEWAY_MOUNT_PATHS"
 
 
 def test_gateway_trim_keeps_metrics_mount():
@@ -179,15 +185,15 @@ def test_gateway_trim_keeps_metrics_mount():
     metrics_mount = Mount("/metrics", app=make_asgi_app())
     routes = [*app.router.routes, metrics_mount]
     trimmed = [r for r in routes if _is_gateway_route(r)]
-    assert metrics_mount in trimmed, "/metrics Mount must survive the gateway route trim"
+    assert metrics_mount in trimmed, \
+        "/metrics Mount must survive the gateway route trim"
 
 
 def test_gateway_drops_ui_and_swagger_mounts():
     """UI static and swagger Mounts must still be trimmed from the gateway."""
     for path in ("/ui", "/_next", "/litellm-asset-prefix/_next", "/swagger"):
-        assert not _is_gateway_route(Mount(path, app=make_asgi_app())), (
+        assert not _is_gateway_route(Mount(path, app=make_asgi_app())), \
             f"Mount {path} must not be served by the gateway"
-        )
 
 
 def test_gateway_keeps_asyncio_task_stacks_route():

tests/test_litellm/proxy/test_sensitive_route_auth.py

@@ -9,7 +9,11 @@ from litellm.proxy.spend_tracking.spend_management_endpoints import (
 
 def _get_route_dependency_calls(router, path: str, method: str):
     for route in router.routes:
-        if isinstance(route, APIRoute) and route.path == path and method in route.methods:
+        if (
+            isinstance(route, APIRoute)
+            and route.path == path
+            and method in route.methods
+        ):
             return [dependency.call for dependency in route.dependant.dependencies]
     raise AssertionError(f"Route {method} {path} not found")
 
@@ -20,8 +24,12 @@ def test_sensitive_debug_routes_require_auth_dependency():
         ("/debug/asyncio-tasks/stacks", "GET"),
         ("/otel-spans", "GET"),
     ):
-        assert user_api_key_auth in _get_route_dependency_calls(debug_router, path, method)
+        assert user_api_key_auth in _get_route_dependency_calls(
+            debug_router, path, method
+        )
 
 
 def test_provider_budgets_requires_auth_dependency():
-    assert user_api_key_auth in _get_route_dependency_calls(spend_router, "/provider/budgets", "GET")
+    assert user_api_key_auth in _get_route_dependency_calls(
+        spend_router, "/provider/budgets", "GET"
+    )