From bb6f43d12e9a14fcafb2653b5b1defa44ea30cdf Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 20 Feb 2025 13:41:56 -0800 Subject: [PATCH] (Bug fix) - Cache Health not working when configured with prometheus service logger (#8687) * fix serialize on safe json dumps * test_non_standard_dict_keys_complex * ui fix HealthCheckCacheParams * fix HealthCheckCacheParams * fix code qa * test_cache_ping_failure * test_cache_ping_health_check_includes_only_cache_attributes * test_cache_ping_health_check_includes_only_cache_attributes --- litellm/litellm_core_utils/safe_json_dumps.py | 3 +- litellm/proxy/caching_routes.py | 44 +++++++++++--- litellm/types/caching.py | 18 +++++- .../test_safe_json_dumps.py | 60 +++++++++++++++++++ tests/litellm/proxy/test_caching_routes.py | 36 ++++++++++- .../src/components/cache_health.tsx | 23 ++++--- 6 files changed, 164 insertions(+), 20 deletions(-) diff --git a/litellm/litellm_core_utils/safe_json_dumps.py b/litellm/litellm_core_utils/safe_json_dumps.py index dda4a54047..990c0ed561 100644 --- a/litellm/litellm_core_utils/safe_json_dumps.py +++ b/litellm/litellm_core_utils/safe_json_dumps.py @@ -23,7 +23,8 @@ def safe_dumps(data: Any, max_depth: int = 10) -> str: if isinstance(obj, dict): result = {} for k, v in obj.items(): - result[k] = _serialize(v, seen, depth + 1) + if isinstance(k, (str)): + result[k] = _serialize(v, seen, depth + 1) seen.remove(id(obj)) return result elif isinstance(obj, list): diff --git a/litellm/proxy/caching_routes.py b/litellm/proxy/caching_routes.py index f4b65974f1..3a22ce2fa8 100644 --- a/litellm/proxy/caching_routes.py +++ b/litellm/proxy/caching_routes.py @@ -9,7 +9,7 @@ from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker from litellm.proxy._types import ProxyErrorTypes, ProxyException from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.types.caching import CachePingResponse +from litellm.types.caching import CachePingResponse, HealthCheckCacheParams masker = SensitiveDataMasker() @@ -19,6 +19,36 @@ router = APIRouter( ) +def _extract_cache_params() -> Dict[str, Any]: + """ + Safely extracts and cleans cache parameters. + + The health check UI needs to display specific cache parameters, to show users how they set up their cache. + + eg. + { + "host": "localhost", + "port": 6379, + "redis_kwargs": {"db": 0}, + "namespace": "test", + } + + Returns: + Dict containing cleaned and masked cache parameters + """ + if litellm.cache is None: + return {} + try: + cache_params = vars(litellm.cache.cache) + cleaned_params = ( + HealthCheckCacheParams(**cache_params).model_dump() if cache_params else {} + ) + return masker.mask_dict(cleaned_params) + except (AttributeError, TypeError) as e: + verbose_proxy_logger.debug(f"Error extracting cache params: {str(e)}") + return {} + + @router.get( "/ping", response_model=CachePingResponse, @@ -29,7 +59,7 @@ async def cache_ping(): Endpoint for checking if cache can be pinged """ litellm_cache_params: Dict[str, Any] = {} - specific_cache_params: Dict[str, Any] = {} + cleaned_cache_params: Dict[str, Any] = {} try: if litellm.cache is None: raise HTTPException( @@ -38,17 +68,13 @@ async def cache_ping(): litellm_cache_params = masker.mask_dict(vars(litellm.cache)) # remove field that might reference itself litellm_cache_params.pop("cache", None) - specific_cache_params = ( - masker.mask_dict(vars(litellm.cache.cache)) if litellm.cache else {} - ) + cleaned_cache_params = _extract_cache_params() if litellm.cache.type == "redis": - # ping the redis cache ping_response = await litellm.cache.ping() verbose_proxy_logger.debug( "/cache/ping: ping_response: " + str(ping_response) ) - # making a set cache call # add cache does not return anything await litellm.cache.async_add_cache( result="test_key", @@ -63,7 +89,7 @@ async def cache_ping(): ping_response=True, set_cache_response="success", litellm_cache_params=safe_dumps(litellm_cache_params), - redis_cache_params=safe_dumps(specific_cache_params), + health_check_cache_params=cleaned_cache_params, ) else: return CachePingResponse( @@ -78,7 +104,7 @@ async def cache_ping(): error_message = { "message": f"Service Unhealthy ({str(e)})", "litellm_cache_params": safe_dumps(litellm_cache_params), - "redis_cache_params": safe_dumps(specific_cache_params), + "health_check_cache_params": safe_dumps(cleaned_cache_params), "traceback": traceback.format_exc(), } raise ProxyException( diff --git a/litellm/types/caching.py b/litellm/types/caching.py index ba439d9249..743c2c0937 100644 --- a/litellm/types/caching.py +++ b/litellm/types/caching.py @@ -1,5 +1,5 @@ from enum import Enum -from typing import Literal, Optional, TypedDict +from typing import Any, Dict, Literal, Optional, TypedDict, Union from pydantic import BaseModel @@ -61,4 +61,18 @@ class CachePingResponse(BaseModel): ping_response: Optional[bool] = None set_cache_response: Optional[str] = None litellm_cache_params: Optional[str] = None - redis_cache_params: Optional[str] = None + + # intentionally a dict, since we run masker.mask_dict() on HealthCheckCacheParams + health_check_cache_params: Optional[dict] = None + + +class HealthCheckCacheParams(BaseModel): + """ + Cache Params returned on /cache/ping call + """ + + host: Optional[str] = None + port: Optional[Union[str, int]] = None + redis_kwargs: Optional[Dict[str, Any]] = None + namespace: Optional[str] = None + redis_version: Optional[str] = None diff --git a/tests/litellm/litellm_core_utils/test_safe_json_dumps.py b/tests/litellm/litellm_core_utils/test_safe_json_dumps.py index 858150ad32..56937e6266 100644 --- a/tests/litellm/litellm_core_utils/test_safe_json_dumps.py +++ b/tests/litellm/litellm_core_utils/test_safe_json_dumps.py @@ -78,3 +78,63 @@ def test_unserializable_object(): obj = TestClass() result = json.loads(safe_dumps(obj)) assert result == "Unserializable Object" + + +def test_non_standard_dict_keys(): + try: + # Test handling of dictionaries with non-standard keys + class GCCollector: + def __str__(self): + return "GCCollector" + + data = {GCCollector(): "value", "test": "test"} + json_dump = safe_dumps(data) + print(json_dump) + result = json.loads(json_dump) + assert result["test"] == "test" + except Exception as e: + print(e) + import traceback + + traceback.print_exc() + raise e + + +def test_non_standard_dict_keys_complex(): + try: + # Test handling of dictionaries with non-standard keys + class GCCollector: + def __str__(self): + return "GCCollector" + + data = [ + {"test": "test"}, + GCCollector(), + { + "bad_key": "bad_value", + }, + { + GCCollector(): "value", + }, + { + "bad_key": GCCollector(), + }, + (GCCollector(), GCCollector()), + ] + json_dump = safe_dumps(data) + print(json_dump) + result = json.loads(json_dump) + print("result=", json.dumps(result, indent=4)) + assert result[0]["test"] == "test" + assert result[1] == "GCCollector" + assert result[2]["bad_key"] == "bad_value" + assert result[3] == {} + assert result[4]["bad_key"] == "GCCollector" + assert result[5][0] == "GCCollector" + assert result[5][1] == "GCCollector" + except Exception as e: + print(e) + import traceback + + traceback.print_exc() + raise e diff --git a/tests/litellm/proxy/test_caching_routes.py b/tests/litellm/proxy/test_caching_routes.py index 7a6dab2e68..291cb0b87c 100644 --- a/tests/litellm/proxy/test_caching_routes.py +++ b/tests/litellm/proxy/test_caching_routes.py @@ -122,7 +122,7 @@ def test_cache_ping_failure(mock_redis_failure): error_details = json.loads(error["message"]) assert "message" in error_details assert "litellm_cache_params" in error_details - assert "redis_cache_params" in error_details + assert "health_check_cache_params" in error_details assert "traceback" in error_details # Verify specific error message @@ -150,3 +150,37 @@ def test_cache_ping_no_cache_initialized(): # Restore original cache litellm.cache = original_cache + + +def test_cache_ping_health_check_includes_only_cache_attributes(mock_redis_success): + """ + Ensure that the /cache/ping endpoint only pulls HealthCheckCacheParams from litellm.cache.cache, + and not from other attributes on litellm.cache. + """ + # Add an unrelated field directly to the cache mock; it should NOT appear in health_check_cache_params + mock_redis_success.some_unrelated_field = "should-not-appear-in-health-check" + + # Add a field on the underlying `cache` object that SHOULD appear + mock_redis_success.cache.redis_kwargs = {"host": "localhost", "port": 6379} + + response = client.get("/cache/ping", headers={"Authorization": "Bearer sk-1234"}) + assert ( + response.status_code == 200 + ), f"Unexpected status code: {response.status_code}" + + data = response.json() + print("/cache/ping response data=", json.dumps(data, indent=4)) + health_check_cache_params = data.get("health_check_cache_params", {}) + # The unrelated field we attached at the top-level of litellm.cache should *not* be present + assert ( + "some_unrelated_field" not in health_check_cache_params + ), "Found an unexpected field from the mock_redis_success object in health_check_cache_params" + + # The field we attached to 'mock_redis_success.cache' should be present and correctly reported + assert ( + "redis_kwargs" in health_check_cache_params + ), "Expected field on `litellm.cache.cache` was not found in health_check_cache_params" + assert health_check_cache_params["redis_kwargs"] == { + "host": "localhost", + "port": 6379, + } diff --git a/ui/litellm-dashboard/src/components/cache_health.tsx b/ui/litellm-dashboard/src/components/cache_health.tsx index ad9816982a..975bf0b49e 100644 --- a/ui/litellm-dashboard/src/components/cache_health.tsx +++ b/ui/litellm-dashboard/src/components/cache_health.tsx @@ -68,6 +68,7 @@ interface RedisDetails { redis_port?: string; redis_version?: string; startup_nodes?: string; + namespace?: string; } // Add new interface for Error Details @@ -75,7 +76,7 @@ interface ErrorDetails { message: string; traceback: string; litellm_params?: any; - redis_cache_params?: any; + health_check_cache_params?: any; } // Update HealthCheckDetails component to handle errors @@ -96,23 +97,23 @@ const HealthCheckDetails: React.FC<{ response: any }> = ({ response }) => { message: errorMessage?.message || 'Unknown error', traceback: errorMessage?.traceback || 'No traceback available', litellm_params: errorMessage?.litellm_cache_params || {}, - redis_cache_params: errorMessage?.redis_cache_params || {} + health_check_cache_params: errorMessage?.health_check_cache_params || {} }; parsedLitellmParams = deepParse(errorDetails.litellm_params) || {}; - parsedRedisParams = deepParse(errorDetails.redis_cache_params) || {}; + parsedRedisParams = deepParse(errorDetails.health_check_cache_params) || {}; } catch (e) { console.warn("Error parsing error details:", e); errorDetails = { message: String(response.error.message || 'Unknown error'), traceback: 'Error parsing details', litellm_params: {}, - redis_cache_params: {} + health_check_cache_params: {} }; } } else { parsedLitellmParams = deepParse(response?.litellm_cache_params) || {}; - parsedRedisParams = deepParse(response?.redis_cache_params) || {}; + parsedRedisParams = deepParse(response?.health_check_cache_params) || {}; } } catch (e) { console.warn("Error in response parsing:", e); @@ -126,11 +127,13 @@ const HealthCheckDetails: React.FC<{ response: any }> = ({ response }) => { redis_host: parsedRedisParams?.redis_client?.connection_pool?.connection_kwargs?.host || parsedRedisParams?.redis_async_client?.connection_pool?.connection_kwargs?.host || parsedRedisParams?.connection_kwargs?.host || + parsedRedisParams?.host || "N/A", redis_port: parsedRedisParams?.redis_client?.connection_pool?.connection_kwargs?.port || parsedRedisParams?.redis_async_client?.connection_pool?.connection_kwargs?.port || parsedRedisParams?.connection_kwargs?.port || + parsedRedisParams?.port || "N/A", redis_version: parsedRedisParams?.redis_version || "N/A", @@ -148,7 +151,9 @@ const HealthCheckDetails: React.FC<{ response: any }> = ({ response }) => { } catch (e) { return "N/A"; } - })() + })(), + + namespace: parsedRedisParams?.namespace || "N/A", }; return ( @@ -229,6 +234,10 @@ const HealthCheckDetails: React.FC<{ response: any }> = ({ response }) => { label="Startup Nodes" value={redisDetails.startup_nodes || "N/A"} /> + )} @@ -244,7 +253,7 @@ const HealthCheckDetails: React.FC<{ response: any }> = ({ response }) => { const data = { ...response, litellm_cache_params: parsedLitellmParams, - redis_cache_params: parsedRedisParams + health_check_cache_params: parsedRedisParams }; // First parse any string JSON values const prettyData = JSON.parse(JSON.stringify(data, (key, value) => {