From 99c91fe41f0c9ee66d85d73e736c4b45f221fada Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 3 Jun 2025 12:03:08 -0700 Subject: [PATCH] [Feat]: Performance add DD profiler to monitor python profile of LiteLLM CPU% (#11375) * feat: add DD profile * fix: test_should_use_dd_profiler * docs dd profiler * docs DD profiler --- docs/my-website/docs/proxy/logging.md | 9 +++ litellm/litellm_core_utils/dd_tracing.py | 5 ++ litellm/proxy/proxy_server.py | 60 ++++++++++++------- .../litellm_core_utils/test_dd_tracing.py | 27 ++++++++- 4 files changed, 77 insertions(+), 24 deletions(-) diff --git a/docs/my-website/docs/proxy/logging.md b/docs/my-website/docs/proxy/logging.md index b4e889119a..3bf5ed1230 100644 --- a/docs/my-website/docs/proxy/logging.md +++ b/docs/my-website/docs/proxy/logging.md @@ -1484,12 +1484,21 @@ Expected output on Datadog Use `ddtrace-run` to enable [Datadog Tracing](https://ddtrace.readthedocs.io/en/stable/installation_quickstart.html) on litellm proxy +**DD Tracer** Pass `USE_DDTRACE=true` to the docker run command. When `USE_DDTRACE=true`, the proxy will run `ddtrace-run litellm` as the `ENTRYPOINT` instead of just `litellm` +**DD Profiler** + +Pass `USE_DDPROFILER=true` to the docker run command. When `USE_DDPROFILER=true`, the proxy will activate the [Datadog Profiler](https://docs.datadoghq.com/profiler/enabling/python/). This is useful for debugging CPU% and memory usage. + +We don't recommend using `USE_DDPROFILER` in production. It is only recommended for debugging CPU% and memory usage. + + ```bash docker run \ -v $(pwd)/litellm_config.yaml:/app/config.yaml \ -e USE_DDTRACE=true \ + -e USE_DDPROFILER=true \ -p 4000:4000 \ ghcr.io/berriai/litellm:main-latest \ --config /app/config.yaml --detailed_debug diff --git a/litellm/litellm_core_utils/dd_tracing.py b/litellm/litellm_core_utils/dd_tracing.py index 1f866a998a..ce784ecf6a 100644 --- a/litellm/litellm_core_utils/dd_tracing.py +++ b/litellm/litellm_core_utils/dd_tracing.py @@ -57,6 +57,11 @@ def _should_use_dd_tracer(): return get_secret_bool("USE_DDTRACE", False) is True +def _should_use_dd_profiler(): + """Returns True if `USE_DDPROFILER` is set to True in .env""" + return get_secret_bool("USE_DDPROFILER", False) is True + + # Initialize tracer should_use_dd_tracer = _should_use_dd_tracer() tracer: Union[NullTracer, DD_TRACER] = NullTracer() diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 2dc086ef97..3e01bf2b1a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -874,9 +874,9 @@ model_max_budget_limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter( dual_cache=user_api_key_cache ) litellm.logging_callback_manager.add_litellm_callback(model_max_budget_limiter) -redis_usage_cache: Optional[ - RedisCache -] = None # redis cache used for tracking spend, tpm/rpm limits +redis_usage_cache: Optional[RedisCache] = ( + None # redis cache used for tracking spend, tpm/rpm limits +) user_custom_auth = None user_custom_key_generate = None user_custom_sso = None @@ -1203,9 +1203,9 @@ async def update_cache( # noqa: PLR0915 _id = "team_id:{}".format(team_id) try: # Fetch the existing cost for the given user - existing_spend_obj: Optional[ - LiteLLM_TeamTable - ] = await user_api_key_cache.async_get_cache(key=_id) + existing_spend_obj: Optional[LiteLLM_TeamTable] = ( + await user_api_key_cache.async_get_cache(key=_id) + ) if existing_spend_obj is None: # do nothing if team not in api key cache return @@ -2780,10 +2780,10 @@ class ProxyConfig: ) try: - guardrails_in_db: List[ - Guardrail - ] = await GuardrailRegistry.get_all_guardrails_from_db( - prisma_client=prisma_client + guardrails_in_db: List[Guardrail] = ( + await GuardrailRegistry.get_all_guardrails_from_db( + prisma_client=prisma_client + ) ) verbose_proxy_logger.debug( "guardrails from the DB %s", str(guardrails_in_db) @@ -3003,9 +3003,9 @@ async def initialize( # noqa: PLR0915 user_api_base = api_base dynamic_config[user_model]["api_base"] = api_base if api_version: - os.environ[ - "AZURE_API_VERSION" - ] = api_version # set this for azure - litellm can read this from the env + os.environ["AZURE_API_VERSION"] = ( + api_version # set this for azure - litellm can read this from the env + ) if max_tokens: # model-specific param dynamic_config[user_model]["max_tokens"] = max_tokens if temperature: # model-specific param @@ -3460,13 +3460,23 @@ class ProxyStartupEvent: DD tracer is used to trace Python applications. Doc: https://docs.datadoghq.com/tracing/trace_collection/automatic_instrumentation/dd_libraries/python/ """ - from litellm.litellm_core_utils.dd_tracing import _should_use_dd_tracer + from litellm.litellm_core_utils.dd_tracing import ( + _should_use_dd_profiler, + _should_use_dd_tracer, + ) if _should_use_dd_tracer(): import ddtrace ddtrace.patch_all(logging=True, openai=False) + if _should_use_dd_profiler(): + from ddtrace.profiling import Profiler + + prof = Profiler() + prof.start() + verbose_proxy_logger.debug("Datadog Profiler started......") + #### API ENDPOINTS #### @router.get( @@ -3651,9 +3661,11 @@ async def chat_completion( # noqa: PLR0915 return StreamingResponse( selected_data_generator, media_type="text/event-stream", - status_code=e.status_code - if hasattr(e, "status_code") - else status.HTTP_400_BAD_REQUEST, + status_code=( + e.status_code + if hasattr(e, "status_code") + else status.HTTP_400_BAD_REQUEST + ), ) _usage = litellm.Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0) _chat_response.usage = _usage # type: ignore @@ -3763,9 +3775,11 @@ async def completion( # noqa: PLR0915 selected_data_generator, media_type="text/event-stream", headers={}, - status_code=e.status_code - if hasattr(e, "status_code") - else status.HTTP_400_BAD_REQUEST, + status_code=( + e.status_code + if hasattr(e, "status_code") + else status.HTTP_400_BAD_REQUEST + ), ) else: _response = litellm.TextCompletionResponse() @@ -7840,9 +7854,9 @@ async def get_config_list( hasattr(sub_field_info, "description") and sub_field_info.description is not None ): - nested_fields[ - idx - ].field_description = sub_field_info.description + nested_fields[idx].field_description = ( + sub_field_info.description + ) idx += 1 _stored_in_db = None diff --git a/tests/test_litellm/litellm_core_utils/test_dd_tracing.py b/tests/test_litellm/litellm_core_utils/test_dd_tracing.py index 9efb1fe7fa..455ad033af 100644 --- a/tests/test_litellm/litellm_core_utils/test_dd_tracing.py +++ b/tests/test_litellm/litellm_core_utils/test_dd_tracing.py @@ -9,7 +9,10 @@ sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system path -from litellm.litellm_core_utils.dd_tracing import _should_use_dd_tracer +from litellm.litellm_core_utils.dd_tracing import ( + _should_use_dd_profiler, + _should_use_dd_tracer, +) from litellm.litellm_core_utils.dd_tracing import tracer as dd_tracer @@ -89,3 +92,25 @@ def test_should_use_dd_tracer(): mock_get_secret.return_value = False assert _should_use_dd_tracer() is False mock_get_secret.assert_called_once_with("USE_DDTRACE", False) + + +def test_should_use_dd_profiler(): + """ + Test that the should_use_dd_profiler function works as expected + """ + with patch( + "litellm.litellm_core_utils.dd_tracing.get_secret_bool" + ) as mock_get_secret: + # Test when USE_DDPROFILER is True + + mock_get_secret.return_value = True + assert _should_use_dd_profiler() is True + mock_get_secret.assert_called_once_with("USE_DDPROFILER", False) + + # Reset the mock for the next test + mock_get_secret.reset_mock() + + # Test when USE_DDPROFILER is False + mock_get_secret.return_value = False + assert _should_use_dd_profiler() is False + mock_get_secret.assert_called_once_with("USE_DDPROFILER", False)