From bde3a98a552121eb54d94862da540a12069347aa Mon Sep 17 00:00:00 2001 From: harish-berri Date: Thu, 30 Apr 2026 18:13:03 +0000 Subject: [PATCH] feat(proxy): add support for Grafana Cloud Pyroscope authentication - Introduced optional environment variables `PYROSCOPE_GRAFANA_USER` and `PYROSCOPE_GRAFANA_API_TOKEN` for Grafana Cloud integration. - Updated documentation to reflect new configuration options for Pyroscope profiling. - Enhanced error handling to ensure both credentials are provided when using Grafana Cloud. - Added tests to validate Grafana Cloud authentication scenarios. --- docs/my-website/docs/proxy/config_settings.md | 2 + .../docs/proxy/pyroscope_profiling.md | 14 +++ litellm/proxy/proxy_server.py | 22 ++++- tests/test_litellm/proxy/test_pyroscope.py | 94 ++++++++++++++++++- 4 files changed, 130 insertions(+), 2 deletions(-) diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index a886a754f5e..f7a2f0a7379 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -855,6 +855,8 @@ router_settings: | PYROSCOPE_APP_NAME | Application name reported to Pyroscope. Required when LITELLM_ENABLE_PYROSCOPE is true. No default. | PYROSCOPE_SERVER_ADDRESS | Pyroscope server URL to send profiles to. Required when LITELLM_ENABLE_PYROSCOPE is true. No default. | PYROSCOPE_SAMPLE_RATE | Optional. Sample rate for Pyroscope profiling (integer). No default; when unset, the pyroscope-io library default is used. +| PYROSCOPE_GRAFANA_USER | Optional. Grafana Cloud Pyroscope user/tenant ID for basic auth. Required when PYROSCOPE_GRAFANA_API_TOKEN is set. +| PYROSCOPE_GRAFANA_API_TOKEN | Optional. Grafana Cloud API/access policy token for Pyroscope basic auth. Required when PYROSCOPE_GRAFANA_USER is set. | LITELLM_MASTER_KEY | Master key for proxy authentication | LITELLM_MAX_BUDGET_PER_SESSION_TTL | TTL in seconds for session budget counters used by the max-budget-per-session limiter. Default is 3600 (1 hour) | LITELLM_MAX_ITERATIONS_TTL | TTL in seconds for session iteration counters used by the max-iterations limiter. Default is 3600 (1 hour) diff --git a/docs/my-website/docs/proxy/pyroscope_profiling.md b/docs/my-website/docs/proxy/pyroscope_profiling.md index 19d12ba24ea..ab6367903bb 100644 --- a/docs/my-website/docs/proxy/pyroscope_profiling.md +++ b/docs/my-website/docs/proxy/pyroscope_profiling.md @@ -24,6 +24,8 @@ LiteLLM proxy can send continuous CPU profiles to [Grafana Pyroscope](https://gr | `PYROSCOPE_APP_NAME` | Yes (when enabled) | Application name shown in the Pyroscope UI. | | `PYROSCOPE_SERVER_ADDRESS` | Yes (when enabled) | Pyroscope server URL (e.g. `http://localhost:4040`). | | `PYROSCOPE_SAMPLE_RATE` | No | Sample rate (integer). If unset, the pyroscope-io library default is used. | + | `PYROSCOPE_GRAFANA_USER` | No | Grafana Cloud Pyroscope user/tenant ID. Required when `PYROSCOPE_GRAFANA_API_TOKEN` is set. | + | `PYROSCOPE_GRAFANA_API_TOKEN` | No | Grafana Cloud API/access policy token. Used as the Pyroscope basic auth password. | 3. **Start the proxy**; profiling will begin automatically when the proxy starts. @@ -34,6 +36,18 @@ LiteLLM proxy can send continuous CPU profiles to [Grafana Pyroscope](https://gr litellm --config config.yaml ``` + For Grafana Cloud Pyroscope, use the Profiles endpoint as `PYROSCOPE_SERVER_ADDRESS` + and set the Grafana Cloud credentials: + + ```bash + export LITELLM_ENABLE_PYROSCOPE=true + export PYROSCOPE_APP_NAME=litellm-proxy + export PYROSCOPE_SERVER_ADDRESS=https://profiles-prod-.grafana.net + export PYROSCOPE_GRAFANA_USER= + export PYROSCOPE_GRAFANA_API_TOKEN= + litellm --config config.yaml + ``` + 4. **View profiles** in the Pyroscope (or Grafana) UI and select your `PYROSCOPE_APP_NAME`. ## Notes diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 00c2cf9e3d6..073f2ad2d03 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7001,6 +7001,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( @@ -7029,11 +7030,30 @@ class ProxyStartupEvent: if env_name: tags["environment"] = env_name sample_rate_env = os.getenv("PYROSCOPE_SAMPLE_RATE") + grafana_pyroscope_user = get_secret_str( + "PYROSCOPE_GRAFANA_USER", default_value=None + ) + grafana_api_token = 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 diff --git a/tests/test_litellm/proxy/test_pyroscope.py b/tests/test_litellm/proxy/test_pyroscope.py index e7b16125dc1..fc2499101ca 100644 --- a/tests/test_litellm/proxy/test_pyroscope.py +++ b/tests/test_litellm/proxy/test_pyroscope.py @@ -103,6 +103,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 +132,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 +141,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 +165,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 +174,89 @@ 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.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", + "PYROSCOPE_GRAFANA_USER": "123456", + "PYROSCOPE_GRAFANA_API_TOKEN": "glc_test_token", + }, + 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.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", + "PYROSCOPE_GRAFANA_USER": "", + "PYROSCOPE_GRAFANA_API_TOKEN": "glc_test_token", + }, + 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.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", + "PYROSCOPE_GRAFANA_USER": "123456", + "PYROSCOPE_GRAFANA_API_TOKEN": "", + }, + clear=False, + ), + ): + with pytest.raises(ValueError, match="PYROSCOPE_GRAFANA_API_TOKEN"): + ProxyStartupEvent._init_pyroscope()