tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-05 08:23:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Feat]: Performance add DD profiler to monitor python profile of LiteLLM CPU% (#11375)

Browse Source 
* feat: add DD profile

* fix: test_should_use_dd_profiler

* docs dd profiler

* docs DD profiler
This commit is contained in:
Ishaan Jaff
2025-06-03 12:03:08 -07:00
committed by GitHub
parent 41a2a62511
commit 99c91fe41f
4 changed files with 77 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/my-website/docs/proxy/logging.md
+9
View File
@@ -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
litellm/litellm_core_utils/dd_tracing.py
+5
View File
@@ -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()
litellm/proxy/proxy_server.py
+37 -23
View File
@@ -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
tests/test_litellm/litellm_core_utils/test_dd_tracing.py
+26 -1
View File
@@ -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)