ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #26902 from BerriAI/litellm_pyroscope_tag_wrapper

Browse Source 
feat(proxy): add support for Grafana Cloud Pyroscope authentication
This commit is contained in:
harish-berri 2026-05-04 14:08:17 -07:00 committed by GitHub
commit f79d3fae2d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 215 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
litellm/proxy/proxy_server.py
View File

@ -476,6 +476,7 @@ from litellm.secret_managers.main import (
get_secret,
get_secret_bool,
get_secret_str,
normalize_nonempty_secret_str,
str_to_bool,
)
from litellm.types.integrations.slack_alerting import SlackAlertingArgs
@ -7487,6 +7488,7 @@ class ProxyStartupEvent:
Requires: pip install pyroscope-io (optional dependency).
When enabled, PYROSCOPE_SERVER_ADDRESS and PYROSCOPE_APP_NAME are required (no defaults).
Optional: PYROSCOPE_SAMPLE_RATE (parsed as integer) to set the sample rate.
Optional: PYROSCOPE_GRAFANA_USER and PYROSCOPE_GRAFANA_API_TOKEN for Grafana Cloud basic auth.
"""
if not get_secret_bool("LITELLM_ENABLE_PYROSCOPE", False):
verbose_proxy_logger.debug(
@ -7515,11 +7517,31 @@ class ProxyStartupEvent:
if env_name:
tags["environment"] = env_name
sample_rate_env = os.getenv("PYROSCOPE_SAMPLE_RATE")
grafana_pyroscope_user = normalize_nonempty_secret_str(
get_secret_str("PYROSCOPE_GRAFANA_USER", default_value=None)
)
grafana_api_token = normalize_nonempty_secret_str(
get_secret_str("PYROSCOPE_GRAFANA_API_TOKEN", default_value=None)
)
if grafana_api_token and not grafana_pyroscope_user:
raise ValueError(
"PYROSCOPE_GRAFANA_API_TOKEN is set but PYROSCOPE_GRAFANA_USER is not set. "
"Set PYROSCOPE_GRAFANA_USER to the Grafana Cloud Pyroscope user/tenant id."
)
if grafana_pyroscope_user and not grafana_api_token:
raise ValueError(
"PYROSCOPE_GRAFANA_USER is set but PYROSCOPE_GRAFANA_API_TOKEN is not set. "
"Set PYROSCOPE_GRAFANA_API_TOKEN to the Grafana Cloud API/access policy token."
)
configure_kwargs = {
"app_name": app_name,
"application_name": app_name,
"server_address": server_address,
"tags": tags if tags else None,
}
if grafana_api_token and grafana_pyroscope_user:
configure_kwargs["basic_auth_username"] = grafana_pyroscope_user
configure_kwargs["basic_auth_password"] = grafana_api_token
if sample_rate_env is not None:
try:
# pyroscope-io expects sample_rate as an integer

13
litellm/secret_managers/main.py
View File

@ -120,6 +120,19 @@ def get_secret_str(
return value
def normalize_nonempty_secret_str(val: Optional[str]) -> Optional[str]:
"""
Strip whitespace and treat None, '', and whitespace-only strings as unset.
Use when pairing secrets (mutual exclusion, optional auth) so whitespace-only
values do not count as present.
"""
if val is None:
return None
stripped = val.strip()
return stripped if stripped else None
def get_secret_bool(
secret_name: str,
default_value: Optional[bool] = None,

165
tests/test_litellm/proxy/test_pyroscope.py
View File

@ -2,11 +2,13 @@
import os
import sys
from typing import Optional
from unittest.mock import MagicMock, patch
import pytest
from litellm.proxy.proxy_server import ProxyStartupEvent
from litellm.secret_managers.main import get_secret_str as real_get_secret_str
def _mock_pyroscope_module():
@ -16,6 +18,22 @@ def _mock_pyroscope_module():
return m
def _patch_pyroscope_grafana_secrets(user: Optional[str], token: Optional[str]):
"""Patch proxy_server.get_secret_str for Grafana keys; defer other secrets to the real helper."""
def side_effect(secret_name: str, default_value=None):
if secret_name == "PYROSCOPE_GRAFANA_USER":
return user
if secret_name == "PYROSCOPE_GRAFANA_API_TOKEN":
return token
return real_get_secret_str(secret_name, default_value)
return patch(
"litellm.proxy.proxy_server.get_secret_str",
side_effect=side_effect,
)
def test_init_pyroscope_returns_cleanly_when_disabled():
"""When LITELLM_ENABLE_PYROSCOPE is false, _init_pyroscope returns without error."""
with (
@ -103,6 +121,8 @@ def test_init_pyroscope_raises_when_sample_rate_invalid():
"PYROSCOPE_APP_NAME": "myapp",
"PYROSCOPE_SERVER_ADDRESS": "http://localhost:4040",
"PYROSCOPE_SAMPLE_RATE": "not-a-number",
"PYROSCOPE_GRAFANA_API_TOKEN": "",
"PYROSCOPE_GRAFANA_USER": "",
},
clear=False,
),
@ -130,6 +150,8 @@ def test_init_pyroscope_accepts_integer_sample_rate():
"PYROSCOPE_APP_NAME": "myapp",
"PYROSCOPE_SERVER_ADDRESS": "http://localhost:4040",
"PYROSCOPE_SAMPLE_RATE": "100",
"PYROSCOPE_GRAFANA_API_TOKEN": "",
"PYROSCOPE_GRAFANA_USER": "",
},
clear=False,
),
@ -137,7 +159,7 @@ def test_init_pyroscope_accepts_integer_sample_rate():
ProxyStartupEvent._init_pyroscope()
mock_pyroscope.configure.assert_called_once()
call_kw = mock_pyroscope.configure.call_args[1]
assert call_kw["app_name"] == "myapp"
assert call_kw["application_name"] == "myapp"
assert call_kw["server_address"] == "http://localhost:4040"
assert call_kw["sample_rate"] == 100
@ -161,6 +183,8 @@ def test_init_pyroscope_accepts_float_sample_rate_parsed_as_int():
"PYROSCOPE_APP_NAME": "myapp",
"PYROSCOPE_SERVER_ADDRESS": "http://localhost:4040",
"PYROSCOPE_SAMPLE_RATE": "100.7",
"PYROSCOPE_GRAFANA_API_TOKEN": "",
"PYROSCOPE_GRAFANA_USER": "",
},
clear=False,
),
@ -168,3 +192,142 @@ def test_init_pyroscope_accepts_float_sample_rate_parsed_as_int():
ProxyStartupEvent._init_pyroscope()
call_kw = mock_pyroscope.configure.call_args[1]
assert call_kw["sample_rate"] == 100
def test_init_pyroscope_configures_grafana_cloud_basic_auth():
"""When Grafana Cloud credentials are set, passes them as Pyroscope basic auth."""
mock_pyroscope = _mock_pyroscope_module()
with (
patch(
"litellm.proxy.proxy_server.get_secret_bool",
return_value=True,
),
_patch_pyroscope_grafana_secrets("123456", "glc_test_token"),
patch.dict(
sys.modules,
{"pyroscope": mock_pyroscope},
),
patch.dict(
os.environ,
{
"LITELLM_ENABLE_PYROSCOPE": "true",
"PYROSCOPE_APP_NAME": "myapp",
"PYROSCOPE_SERVER_ADDRESS": "https://profiles-prod-001.grafana.net",
},
clear=False,
),
):
ProxyStartupEvent._init_pyroscope()
call_kw = mock_pyroscope.configure.call_args[1]
assert call_kw["basic_auth_username"] == "123456"
assert call_kw["basic_auth_password"] == "glc_test_token"
def test_init_pyroscope_raises_when_grafana_token_missing_user():
"""When Grafana token is set without a Pyroscope user, raises ValueError."""
mock_pyroscope = _mock_pyroscope_module()
with (
patch(
"litellm.proxy.proxy_server.get_secret_bool",
return_value=True,
),
_patch_pyroscope_grafana_secrets("", "glc_test_token"),
patch.dict(
sys.modules,
{"pyroscope": mock_pyroscope},
),
patch.dict(
os.environ,
{
"LITELLM_ENABLE_PYROSCOPE": "true",
"PYROSCOPE_APP_NAME": "myapp",
"PYROSCOPE_SERVER_ADDRESS": "https://profiles-prod-001.grafana.net",
},
clear=False,
),
):
with pytest.raises(ValueError, match="PYROSCOPE_GRAFANA_USER"):
ProxyStartupEvent._init_pyroscope()
def test_init_pyroscope_raises_when_grafana_user_missing_token():
"""When Grafana Pyroscope user is set without a token, raises ValueError."""
mock_pyroscope = _mock_pyroscope_module()
with (
patch(
"litellm.proxy.proxy_server.get_secret_bool",
return_value=True,
),
_patch_pyroscope_grafana_secrets("123456", ""),
patch.dict(
sys.modules,
{"pyroscope": mock_pyroscope},
),
patch.dict(
os.environ,
{
"LITELLM_ENABLE_PYROSCOPE": "true",
"PYROSCOPE_APP_NAME": "myapp",
"PYROSCOPE_SERVER_ADDRESS": "https://profiles-prod-001.grafana.net",
},
clear=False,
),
):
with pytest.raises(ValueError, match="PYROSCOPE_GRAFANA_API_TOKEN"):
ProxyStartupEvent._init_pyroscope()
def test_init_pyroscope_raises_when_grafana_user_whitespace_only_with_token():
"""Whitespace-only user id does not satisfy Grafana mutual exclusion."""
mock_pyroscope = _mock_pyroscope_module()
with (
patch(
"litellm.proxy.proxy_server.get_secret_bool",
return_value=True,
),
_patch_pyroscope_grafana_secrets(" \t", "glc_test_token"),
patch.dict(
sys.modules,
{"pyroscope": mock_pyroscope},
),
patch.dict(
os.environ,
{
"LITELLM_ENABLE_PYROSCOPE": "true",
"PYROSCOPE_APP_NAME": "myapp",
"PYROSCOPE_SERVER_ADDRESS": "https://profiles-prod-001.grafana.net",
},
clear=False,
),
):
with pytest.raises(ValueError, match="PYROSCOPE_GRAFANA_USER"):
ProxyStartupEvent._init_pyroscope()
def test_init_pyroscope_strips_grafana_credentials_for_basic_auth():
"""Leading/trailing whitespace on Grafana secrets is trimmed before configure."""
mock_pyroscope = _mock_pyroscope_module()
with (
patch(
"litellm.proxy.proxy_server.get_secret_bool",
return_value=True,
),
_patch_pyroscope_grafana_secrets(" 123456 ", " glc_test_token\n"),
patch.dict(
sys.modules,
{"pyroscope": mock_pyroscope},
),
patch.dict(
os.environ,
{
"LITELLM_ENABLE_PYROSCOPE": "true",
"PYROSCOPE_APP_NAME": "myapp",
"PYROSCOPE_SERVER_ADDRESS": "https://profiles-prod-001.grafana.net",
},
clear=False,
),
):
ProxyStartupEvent._init_pyroscope()
call_kw = mock_pyroscope.configure.call_args[1]
assert call_kw["basic_auth_username"] == "123456"
assert call_kw["basic_auth_password"] == "glc_test_token"

16
tests/test_litellm/secret_managers/test_secret_managers_main.py
View File

@ -4,7 +4,7 @@ from unittest.mock import Mock, patch
import pytest
from litellm.secret_managers.main import get_secret
from litellm.secret_managers.main import get_secret, normalize_nonempty_secret_str
# Set up logging for debugging
logging.basicConfig(level=logging.DEBUG)
@ -253,3 +253,17 @@ def test_unsupported_oidc_provider():
with pytest.raises(ValueError, match="Unsupported OIDC provider"):
get_secret(secret_name)
@pytest.mark.parametrize(
("raw", "expected"),
[
(None, None),
("", None),
(" \t\n", None),
("abc", "abc"),
(" xyz ", "xyz"),
],
)
def test_normalize_nonempty_secret_str(raw, expected):
assert normalize_nonempty_secret_str(raw) == expected