mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
fix: stable deterministic IDs + correct cleanup for pass-through registry leak
Fixes #24833 — _registered_pass_through_routes grows unbounded causing CPU to reach 100% when pass-through endpoints are stored in the DB. Two root-cause bugs: 1. **UUID churn** — DB-sourced endpoints have no uid=1000(ubuntu) gid=1000(ubuntu) groups=1000(ubuntu),27(sudo) field, so generated a fresh on every 30-second reload cycle. Each new UUID produced a different route key, so old entries never matched during cleanup and accumulated forever. With N endpoints and a 30-second cycle, the dict grew by N entries every 30 s and the O(n²) scan in consumed all available CPU. Fix: derive a stable, deterministic ID from the endpoint's path + methods using SHA-256 (first 16 hex chars, prefixed with 'auto-'). An unchanged endpoint now always maps to the same registry key across reloads. 2. **Wrong argument to remove_endpoint_routes** — the cleanup loop in passed the full route key (e.g. 'uuid:exact:/path:GET,POST') to , which expects only the endpoint_id prefix. The inner linear scan compared value['endpoint_id'] == full_route_key, which never matched any stored entry. Cleanup was a complete no-op regardless of whether the ID changed. Fix: extract the endpoint_id by splitting the route key on ':' with maxsplit=1 before calling remove_endpoint_routes(). Also added at module level (previously would have been a repeated inline import). Tests: 5 new unit tests in tests/pass_through_unit_tests/test_passthrough_registry_leak.py covering: - stable ID for same path across reloads - different paths produce different IDs - remove_endpoint_routes works when called with extracted endpoint_id - registry size does not grow across 5 consecutive reload cycles - explicit IDs are preserved (no regression)
This commit is contained in:
parent
7046a58885
commit
5866c862be
2 changed files with 276 additions and 2 deletions
|
|
@ -1,6 +1,7 @@
|
|||
import ast
|
||||
import asyncio
|
||||
import copy
|
||||
import hashlib
|
||||
import json
|
||||
import traceback
|
||||
from base64 import b64encode
|
||||
|
|
@ -2152,7 +2153,18 @@ async def _register_pass_through_endpoint(
|
|||
endpoint_data = endpoint
|
||||
|
||||
if endpoint_data.get("id") is None:
|
||||
endpoint_data["id"] = str(uuid.uuid4())
|
||||
# Derive a stable, deterministic ID from path + methods so that
|
||||
# repeated reload cycles (e.g. every 30 s from DB) always produce the
|
||||
# same ID. Random UUIDs caused the registry to grow unboundedly:
|
||||
# old entries could never be matched during cleanup because the UUID
|
||||
# changed on every reload. Using a content-based ID means an
|
||||
# unchanged endpoint always maps to the same registry key.
|
||||
_path_for_id = endpoint_data.get("path") or ""
|
||||
_methods_for_id = sorted(endpoint_data.get("methods") or [])
|
||||
_stable_key = f"path:{_path_for_id}|methods:{','.join(_methods_for_id)}"
|
||||
endpoint_data["id"] = "auto-" + hashlib.sha256(
|
||||
_stable_key.encode()
|
||||
).hexdigest()[:16]
|
||||
endpoint_id = cast(str, endpoint_data["id"])
|
||||
|
||||
target = endpoint_data.get("target")
|
||||
|
|
@ -2291,7 +2303,14 @@ async def initialize_pass_through_endpoints(
|
|||
# remove the ones that are not visited from the list
|
||||
for endpoint_key in registered_pass_through_endpoints:
|
||||
if endpoint_key not in visited_endpoints:
|
||||
InitPassThroughEndpointHelpers.remove_endpoint_routes(endpoint_key)
|
||||
# Route keys are formatted as "{endpoint_id}:exact:{path}:{methods}"
|
||||
# or "{endpoint_id}:subpath:{path}:{methods}".
|
||||
# remove_endpoint_routes() expects only the endpoint_id prefix, so
|
||||
# we must split it out here. Previously the full key was passed
|
||||
# verbatim, which never matched any stored endpoint_id and silently
|
||||
# left stale routes in the registry forever.
|
||||
stale_endpoint_id = endpoint_key.split(":", 1)[0]
|
||||
InitPassThroughEndpointHelpers.remove_endpoint_routes(stale_endpoint_id)
|
||||
|
||||
|
||||
def _get_pass_through_endpoints_from_config() -> List[PassThroughGenericEndpoint]:
|
||||
|
|
|
|||
255
tests/pass_through_unit_tests/test_passthrough_registry_leak.py
Normal file
255
tests/pass_through_unit_tests/test_passthrough_registry_leak.py
Normal file
|
|
@ -0,0 +1,255 @@
|
|||
"""
|
||||
Regression tests for pass-through endpoint registry unbounded growth.
|
||||
|
||||
Issue #24833: When endpoints come from the DB they have no `id` field.
|
||||
Before the fix, a new random UUID was generated on every 30-second reload
|
||||
cycle, so:
|
||||
1. Every reload produced a different route key in _registered_pass_through_routes
|
||||
2. The old entries never matched during cleanup (UUID changed) → registry grew unboundedly
|
||||
3. Additionally, the cleanup loop passed the full route key to remove_endpoint_routes()
|
||||
which expected only the endpoint_id prefix → cleanup was a no-op
|
||||
|
||||
Both issues are covered here.
|
||||
"""
|
||||
import asyncio
|
||||
import hashlib
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
|
||||
InitPassThroughEndpointHelpers,
|
||||
_register_pass_through_endpoint,
|
||||
_registered_pass_through_routes,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helper: deterministic ID that the fixed code should produce
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _expected_auto_id(path: str, methods=None) -> str:
|
||||
"""Mirror the deterministic-ID logic introduced in the fix."""
|
||||
_methods = sorted(methods or [])
|
||||
stable_key = f"path:{path}|methods:{','.join(_methods)}"
|
||||
return "auto-" + hashlib.sha256(stable_key.encode()).hexdigest()[:16]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 1 – stable ID for DB endpoints (no `id` field)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_db_endpoint_without_id_gets_stable_deterministic_id():
|
||||
"""
|
||||
A DB-sourced endpoint (no `id` field) must always receive the same ID
|
||||
across multiple calls with identical path+methods. Random UUIDs would
|
||||
break cleanup.
|
||||
"""
|
||||
|
||||
async def _run():
|
||||
app_mock = MagicMock()
|
||||
visited: set = set()
|
||||
|
||||
endpoint_a = {
|
||||
"path": "/my-service",
|
||||
"target": "http://backend:8080",
|
||||
}
|
||||
endpoint_b = dict(endpoint_a) # fresh copy simulating second reload
|
||||
|
||||
await _register_pass_through_endpoint(
|
||||
endpoint=endpoint_a,
|
||||
app=app_mock,
|
||||
premium_user=False,
|
||||
visited_endpoints=visited,
|
||||
)
|
||||
id_first = endpoint_a["id"]
|
||||
|
||||
# Reset visited so the second call is treated as a fresh reload
|
||||
visited.clear()
|
||||
await _register_pass_through_endpoint(
|
||||
endpoint=endpoint_b,
|
||||
app=app_mock,
|
||||
premium_user=False,
|
||||
visited_endpoints=visited,
|
||||
)
|
||||
id_second = endpoint_b["id"]
|
||||
|
||||
assert id_first == id_second, (
|
||||
f"Expected stable ID across reloads but got {id_first!r} then {id_second!r}. "
|
||||
"Random UUIDs cause registry leak."
|
||||
)
|
||||
assert id_first == _expected_auto_id("/my-service"), (
|
||||
f"ID format changed — expected {_expected_auto_id('/my-service')!r}, got {id_first!r}"
|
||||
)
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 2 – different paths produce different IDs
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_different_paths_produce_different_ids():
|
||||
async def _run():
|
||||
app_mock = MagicMock()
|
||||
|
||||
ep1 = {"path": "/svc-alpha", "target": "http://a"}
|
||||
ep2 = {"path": "/svc-beta", "target": "http://b"}
|
||||
visited: set = set()
|
||||
|
||||
await _register_pass_through_endpoint(
|
||||
endpoint=ep1, app=app_mock, premium_user=False, visited_endpoints=visited
|
||||
)
|
||||
await _register_pass_through_endpoint(
|
||||
endpoint=ep2, app=app_mock, premium_user=False, visited_endpoints=visited
|
||||
)
|
||||
|
||||
assert ep1["id"] != ep2["id"], "Different paths must produce different IDs"
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 3 – remove_endpoint_routes works when called with full route key prefix
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_remove_endpoint_routes_called_with_endpoint_id_prefix():
|
||||
"""
|
||||
The cleanup loop extracts the endpoint_id by splitting on ':'.
|
||||
Verify that remove_endpoint_routes() correctly deletes all routes for
|
||||
that endpoint_id, regardless of how many route keys exist for it.
|
||||
"""
|
||||
endpoint_id = "test-remove-prefix-ep"
|
||||
path = "/test-remove-prefix"
|
||||
methods_str = "DELETE,GET,PATCH,POST,PUT"
|
||||
exact_key = f"{endpoint_id}:exact:{path}:{methods_str}"
|
||||
sub_key = f"{endpoint_id}:subpath:{path}:{methods_str}"
|
||||
|
||||
# Manually insert fake entries
|
||||
_registered_pass_through_routes[exact_key] = {
|
||||
"endpoint_id": endpoint_id,
|
||||
"path": path,
|
||||
"type": "exact",
|
||||
"passthrough_params": {},
|
||||
}
|
||||
_registered_pass_through_routes[sub_key] = {
|
||||
"endpoint_id": endpoint_id,
|
||||
"path": path,
|
||||
"type": "subpath",
|
||||
"passthrough_params": {},
|
||||
}
|
||||
|
||||
try:
|
||||
assert exact_key in _registered_pass_through_routes
|
||||
assert sub_key in _registered_pass_through_routes
|
||||
|
||||
# Simulate what the fixed cleanup loop does:
|
||||
# route_key → split to get endpoint_id → remove_endpoint_routes(endpoint_id)
|
||||
full_route_key = exact_key
|
||||
extracted_id = full_route_key.split(":", 1)[0]
|
||||
assert extracted_id == endpoint_id
|
||||
|
||||
InitPassThroughEndpointHelpers.remove_endpoint_routes(extracted_id)
|
||||
|
||||
# Both keys should be gone
|
||||
assert exact_key not in _registered_pass_through_routes, (
|
||||
"Exact route was not removed after cleanup"
|
||||
)
|
||||
assert sub_key not in _registered_pass_through_routes, (
|
||||
"Subpath route was not removed after cleanup"
|
||||
)
|
||||
finally:
|
||||
# Safety cleanup
|
||||
_registered_pass_through_routes.pop(exact_key, None)
|
||||
_registered_pass_through_routes.pop(sub_key, None)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 4 – registry does NOT grow across multiple simulated reload cycles
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_registry_does_not_grow_across_reload_cycles():
|
||||
"""
|
||||
Simulate 5 consecutive reload cycles of the same DB endpoint.
|
||||
Registry size should remain constant (1 exact route), not grow by 1
|
||||
each cycle as it did before the fix.
|
||||
"""
|
||||
|
||||
async def _run():
|
||||
app_mock = MagicMock()
|
||||
|
||||
path = "/stable-reload-test"
|
||||
target = "http://backend:9000"
|
||||
|
||||
# Count entries related to our test path before we start
|
||||
def _count_our_routes():
|
||||
return sum(
|
||||
1
|
||||
for k in _registered_pass_through_routes
|
||||
if path in k
|
||||
)
|
||||
|
||||
initial_count = _count_our_routes()
|
||||
|
||||
for _ in range(5):
|
||||
endpoint = {"path": path, "target": target} # fresh dict each cycle
|
||||
visited: set = set()
|
||||
await _register_pass_through_endpoint(
|
||||
endpoint=endpoint,
|
||||
app=app_mock,
|
||||
premium_user=False,
|
||||
visited_endpoints=visited,
|
||||
)
|
||||
|
||||
final_count = _count_our_routes()
|
||||
|
||||
# Should have added exactly 1 route, not 5
|
||||
assert final_count == initial_count + 1, (
|
||||
f"Registry grew from {initial_count} to {final_count} after 5 reload cycles. "
|
||||
f"Expected exactly {initial_count + 1}. "
|
||||
"This indicates the UUID-churn regression is still present."
|
||||
)
|
||||
|
||||
# Cleanup
|
||||
keys_to_del = [k for k in _registered_pass_through_routes if path in k]
|
||||
for k in keys_to_del:
|
||||
del _registered_pass_through_routes[k]
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 5 – explicitly-assigned IDs are preserved (no regression)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_explicit_id_is_preserved():
|
||||
"""
|
||||
Endpoints that already carry an `id` field must keep that ID.
|
||||
The deterministic-ID path must only trigger when `id` is absent.
|
||||
"""
|
||||
|
||||
async def _run():
|
||||
app_mock = MagicMock()
|
||||
custom_id = "my-custom-endpoint-id-12345"
|
||||
endpoint = {
|
||||
"path": "/explicit-id-test",
|
||||
"target": "http://svc",
|
||||
"id": custom_id,
|
||||
}
|
||||
visited: set = set()
|
||||
|
||||
await _register_pass_through_endpoint(
|
||||
endpoint=endpoint,
|
||||
app=app_mock,
|
||||
premium_user=False,
|
||||
visited_endpoints=visited,
|
||||
)
|
||||
|
||||
assert endpoint["id"] == custom_id, (
|
||||
f"Explicit ID was overwritten! Got {endpoint['id']!r} instead of {custom_id!r}"
|
||||
)
|
||||
|
||||
# Cleanup
|
||||
keys_to_del = [k for k in _registered_pass_through_routes if custom_id in k]
|
||||
for k in keys_to_del:
|
||||
del _registered_pass_through_routes[k]
|
||||
|
||||
asyncio.run(_run())
|
||||
Loading…
Add table
Reference in a new issue