mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
Litellm ishaan april4 2 (#25150)
* feat(router): integrate allowed_fails_policy into health check failures (#24988) * feat(router): integrate allowed_fails_policy into health check failures Health check failures now increment the same per-deployment failure counters used by allowed_fails_policy, so users can control how many health check failures of each error type are required before a deployment enters cooldown. - ahealth_check() preserves the original exception in its return dict - run_with_timeout() returns a litellm.Timeout on health check timeout - _perform_health_check() propagates exceptions to unhealthy endpoints - _write_health_state_to_router_cache() calls _set_cooldown_deployments for each unhealthy endpoint that has an exception - When allowed_fails_policy is set, the binary health check filter is bypassed so cooldown is the sole routing exclusion mechanism - Safety net: if all deployments are in cooldown with enable_health_check_routing=True, the cooldown filter is bypassed Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * feat(router): add health_check_ignore_transient_errors flag When enabled, health check failures with 429 (rate limit) or 408 (timeout) status codes are skipped from the cooldown pipeline. These are transient load issues, not broken deployments. Auth errors (401), 404, and 5xx errors still increment counters and trigger cooldown as before. Config (general_settings): health_check_ignore_transient_errors: true Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix(router): also exclude 429/408 from health state cache when ignore_transient_errors set The previous fix only skipped cooldown counter increments. The health state cache was still marking 429/408 endpoints as is_healthy=False, causing the binary health check filter to exclude them from routing. Now, when health_check_ignore_transient_errors=True, 429/408 endpoints are also excluded from the unhealthy list passed to build_deployment_health_states(), so the binary filter treats them as unaffected (not unhealthy). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * docs(router): add health check driven routing guide New standalone page covering the full health check routing feature: allowed_fails_policy integration, health_check_ignore_transient_errors, architecture SVG, step-by-step setup, and gotchas (TTL, AllowedFails semantics). Replaces the inline section in health.md with a link to the new page. Added to the Routing & Load Balancing sidebar. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix(health-check-routing): fix three CI failures - Add "exception" to ILLEGAL_DISPLAY_PARAMS in health_check.py so the exception object is stripped before the health endpoint serializes results to JSON (fixes TypeError: 'URL' object is not iterable) - Add allowed_fails_policy = None to FakeRouter stubs in test_router_health_check_routing.py (fixes AttributeError) - Add health_check_ignore_transient_errors to config_settings.md router settings reference table (fixes documentation test) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * Fix litellm/tests/proxy_unit_tests/test_proxy_server.py * fix(router): address greptile review comments - Narrow cooldown safety-net bypass: only fires when allowed_fails_policy is set (cooldown is health-check driven). Without a policy, cooldowns are from real request failures and must not be bypassed. - Restore cooldown deployments DEBUG log that was accidentally removed. - Fix test_health TypeError: move exception extraction to a separate exceptions_by_model_id dict returned alongside endpoints, so exception objects never appear in the endpoint dicts that get JSON-serialized by the /health response. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix(health-check-routing): properly isolate exceptions from health response Return exceptions_by_model_id as a separate third value from _perform_health_check / perform_health_check so exception objects (which contain non-JSON-serializable httpx URL types) never appear in the endpoint dicts that get serialized by the /health response. Callers updated: _health_endpoints.py, shared_health_check_manager.py, proxy_server.py background loop. All use the exceptions dict only for cooldown integration, not for display. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix(shared-health-check): fix remaining 2-value return sites and update type annotation * fix(health-check-routing): fix P0 cooldown integration never firing The cooldown loop was reading endpoint.get("exception") which is always None because exceptions are now returned via exceptions_by_model_id, not stored in endpoint dicts. Fixed to use _exceptions.get(model_id). Also fixes the transient-error filter to use _exceptions instead of endpoint.get("exception"), and fixes all remaining 2-value return sites in shared_health_check_manager.py. Tests updated to pass exceptions via exceptions_by_model_id parameter instead of endpoint dicts. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix(health-check-routing): fix P1 transient-error filter broken on cache hits When SharedHealthCheckManager returns cached results, exceptions_by_model_id is always {} so the transient-error filter defaulted to status 500 for all endpoints, incorrectly marking 429/408 endpoints as unhealthy. Fix: store integer exception_status on each unhealthy endpoint dict in _perform_health_check. _get_endpoint_exception_status() uses the live exception object when available (direct path) and falls back to the stored integer (cache-hit path). The integer is JSON-serializable and survives the shared cache round-trip. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix(health-check-routing): gate cooldown loop behind allowed_fails_policy Without the policy, cooldown is not the routing exclusion mechanism. Firing _set_cooldown_deployments for all enable_health_check_routing users was a backwards-incompatible change — 401s would immediately cooldown deployments that the binary filter would have recovered on the next cycle. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * revert: undo allowed_fails_policy gate on cooldown loop Cooldown integration via health checks is intentional for all enable_health_check_routing users, not just those with allowed_fails_policy. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix(docs+tests): fix health_check_ignore_transient_errors doc section and test coverage - Move health_check_ignore_transient_errors from router_settings to general_settings in config_settings.md (code reads it from general_settings) - Remove duplicate enable_health_check_routing / health_check_staleness_threshold entries that were incorrectly listed under router_settings - Replace TestHealthCheckEndpointExceptionPropagation tests with ones that exercise the real _perform_health_check code path via mocked ahealth_check, verifying exceptions appear in exceptions_by_model_id and NOT in endpoint dicts Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix(tests+docs): fix tuple unpacking and docs test failures - Update test mocks that return (healthy, unhealthy) to return (healthy, unhealthy, {}) to match the new 3-value signature - Update test unpackings of perform_shared_health_check to use healthy, unhealthy, _ = ... - Add health_check_ignore_transient_errors to router_settings section in config_settings.md (it is a Router constructor param, so the doc test requires it there; it also lives in general_settings for proxy use) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * Fix CodeQL errors * fix(tests): fix 2-value unpackings of _perform_health_check in test_health_check.py * fix(tests): fix mock _perform_health_check returning 2-tuple instead of 3 * fix team routing --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: add distributed lock for key rotation job (#23364) * fix: add distributed lock for key rotation job * fix: address Greptile review feedback on key rotation lock (#23834) * fix: address Greptile review feedback on key rotation lock * fix req changes greptile * feat(proxy): Optional on_error for guardrail pipeline (API / technical failures) (#24831) * guardrails fallback * docs * docs: add LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS to environment variables reference * fix(mypy): accept Union[Dict, Any] in _get_deployment_order and use typed list to fix min() type error * fix(mypy): use Optional[str] for api_base in PydanticAI provider to match superclass signature --------- Co-authored-by: Sameer Kankute <sameer@berri.ai> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Co-authored-by: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Co-authored-by: Shivam Rawat <shivam@berri.ai> Co-authored-by: yuneng-jiang <yuneng@berri.ai>
This commit is contained in:
ishaan-berri
GitHub
parent
8d9e5ff3d4
commit
51876292a0
31 changed files with 2873 additions and 254 deletions
|
|
@ -289,6 +289,7 @@ router_settings:
|
|||
| database_connection_pool_timeout | integer | Database connection pool timeout in seconds |
|
||||
| disable_error_logs | boolean | If true, suppresses error tracking and storage in the database |
|
||||
| enable_health_check_routing | boolean | If true, enables health check-driven request routing to avoid unhealthy deployments |
|
||||
| health_check_ignore_transient_errors | boolean | If true, 429 (rate limit) and 408 (timeout) health check failures are ignored and do not affect routing or cooldown |
|
||||
| enable_mcp_registry | boolean | If true, enables access to the centralized MCP server registry |
|
||||
| enforce_rbac | boolean | If true, enables role-based access control (RBAC) for all proxy operations |
|
||||
| forward_llm_provider_auth_headers | boolean | If true, forwards provider-specific auth headers to LLM API calls |
|
||||
|
|
@ -397,6 +398,7 @@ router_settings:
|
|||
| guardrail_list | List[GuardrailTypedDict] | List of guardrail configurations for guardrail load balancing. Enables load balancing across multiple guardrail deployments with the same guardrail_name. [Further Docs](./guardrails/guardrail_load_balancing.md) |
|
||||
| enable_health_check_routing | boolean | If true, enables health check-driven deployment filtering to avoid routing requests to unhealthy deployments |
|
||||
| health_check_staleness_threshold | integer | Maximum age in seconds for cached health check results before marking deployments as stale |
|
||||
| health_check_ignore_transient_errors | boolean | If true, 429 (rate limit) and 408 (timeout) health check failures are ignored and do not affect routing or cooldown |
|
||||
|
||||
|
||||
### environment variables - Reference
|
||||
|
|
@ -821,6 +823,7 @@ router_settings:
|
|||
| LITELLM_KEY_ROTATION_ENABLED | Enable auto-key rotation for LiteLLM (boolean). Default is false.
|
||||
| LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS | Interval in seconds for how often to run job that auto-rotates keys. Default is 86400 (24 hours).
|
||||
| LITELLM_KEY_ROTATION_GRACE_PERIOD | Duration to keep old key valid after rotation (e.g. "24h", "2d"). Default is empty (immediate revoke). Used for scheduled rotations and as fallback when not specified in regenerate request.
|
||||
| LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS | TTL in seconds for the distributed lock used by the key rotation job. Default is 600 (10 minutes).
|
||||
| LITELLM_LICENSE | License key for LiteLLM usage
|
||||
| LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS | Set to `True` to use the local bundled Anthropic beta headers config only, disabling remote fetching. Default is `False`
|
||||
| LITELLM_LOCAL_BLOG_POSTS | When set to `True`, uses the local bundled blog posts only, disabling remote fetching from GitHub. Default is `False`
|
||||
|
|
|
|||
|
|
@ -311,7 +311,7 @@ Response:
|
|||
|
||||
## Policy Flow Builder
|
||||
|
||||
For conditional execution (e.g., run a second guardrail only if the first fails), use the [Policy Flow Builder](./policy_flow_builder) to define pipelines with per-step pass/fail actions.
|
||||
For conditional execution (e.g., run a second guardrail only if the first fails), use the [Policy Flow Builder](./policy_flow_builder) to define pipelines with per-step **pass**, **fail**, and optional **error** actions (`on_pass`, `on_fail`, `on_error`).
|
||||
|
||||
## Config Reference
|
||||
|
||||
|
|
@ -337,7 +337,7 @@ policies:
|
|||
| `guardrails.add` | `list[string]` | Guardrails to enable. |
|
||||
| `guardrails.remove` | `list[string]` | Guardrails to disable (useful with inheritance). |
|
||||
| `condition.model` | `string` or `list[string]` | Optional. Only apply when model matches. Supports regex. |
|
||||
| `pipeline` | `object` | Optional. Ordered guardrail execution with per-step actions. See [Policy Flow Builder](./policy_flow_builder). |
|
||||
| `pipeline` | `object` | Optional. Ordered guardrail execution with per-step actions (`on_pass`, `on_fail`, optional `on_error`). See [Policy Flow Builder](./policy_flow_builder). |
|
||||
|
||||
### `policy_attachments`
|
||||
|
||||
|
|
|
|||
|
|
@ -1,8 +1,8 @@
|
|||
# Policy Flow Builder
|
||||
|
||||
The Policy Flow Builder lets you design guardrail pipelines with **conditional execution**. Instead of running guardrails independently, you chain them into ordered steps and control what happens when each guardrail passes or fails.
|
||||
The Policy Flow Builder lets you design guardrail pipelines with **conditional execution**. Instead of running guardrails independently, you chain them into ordered steps and control what happens when each guardrail **passes**, **fails a policy check** (content intervention), or hits a **technical error** (e.g. timeout, unreachable provider, missing guardrail).
|
||||
|
||||
Two powerful patterns it enables: **guardrail fallbacks** (try a different guardrail when one fails) and **retrying the same guardrail** (run the same guardrail again if it fails, e.g. to handle transient errors).
|
||||
Two powerful patterns it enables: **guardrail fallbacks** (try a different guardrail when one fails) and **retrying the same guardrail** (run the same guardrail again if it fails, e.g. to handle transient errors). With **`on_error`**, you can treat **technical** failures differently from **policy** failures—for example, fall back to another provider when the primary API errors, while still blocking on flagged content.
|
||||
|
||||
## When to use the Flow Builder
|
||||
|
||||
|
|
@ -19,6 +19,7 @@ Use the Flow Builder when you need:
|
|||
- **Custom responses** — return a specific message when a guardrail fails instead of a generic block
|
||||
- **Data chaining** — pass modified data (e.g., PII-masked content) from one step to the next
|
||||
- **Fine-grained control** — different actions on pass vs. fail per step
|
||||
- **Technical-error routing** — set `on_error` separately from `on_fail` so outages or timeouts can **allow**, **block**, **go to the next step**, or return a **custom response** without conflating them with content violations
|
||||
|
||||
## Concepts
|
||||
|
||||
|
|
@ -29,24 +30,37 @@ A pipeline has:
|
|||
- **Mode**: `pre_call` (before the LLM) or `post_call` (after the LLM)
|
||||
- **Steps**: Ordered list of guardrail steps
|
||||
|
||||
### Outcomes: pass, fail, and error
|
||||
|
||||
Each step run produces one of three outcomes:
|
||||
|
||||
| Outcome | Meaning | Typical cause |
|
||||
|--------|---------|----------------|
|
||||
| **pass** | Guardrail completed without blocking | Content allowed, or data was modified and returned |
|
||||
| **fail** | Policy intervention | Guardrail raised an intervention (e.g. flagged content, blocked request) |
|
||||
| **error** | Technical failure | Timeouts, network errors, guardrail not registered, or other non-intervention exceptions |
|
||||
|
||||
`on_pass` and `on_fail` apply to **pass** and **fail** respectively. **`on_error`** applies only to **error**. If `on_error` is omitted, the pipeline uses **`on_fail`** for error outcomes (backward compatible).
|
||||
|
||||
### Step actions
|
||||
|
||||
Each step defines what happens when the guardrail **passes** and when it **fails**:
|
||||
For each step you choose an action for **pass**, **fail**, and optionally **error**. Allowed values are: `next`, `allow`, `block`, `modify_response`.
|
||||
|
||||
| Action | Description |
|
||||
|--------|-------------|
|
||||
| **Next Step** | Continue to the next guardrail in the pipeline |
|
||||
| **Allow** | Stop the pipeline and allow the request to proceed |
|
||||
| **Block** | Stop the pipeline and block the request |
|
||||
| **Custom Response** | Return a custom message instead of the default block |
|
||||
| **Next Step** (`next`) | Continue to the next guardrail in the pipeline |
|
||||
| **Allow** (`allow`) | Stop the pipeline and allow the request to proceed |
|
||||
| **Block** (`block`) | Stop the pipeline and block the request |
|
||||
| **Custom Response** (`modify_response`) | Return a custom message instead of the default block |
|
||||
|
||||
### Step options
|
||||
|
||||
| Field | Type | Description |
|
||||
|-------|------|--------------|
|
||||
| `guardrail` | `string` | Name of the guardrail to run |
|
||||
| `on_pass` | `string` | Action when guardrail passes: `next`, `allow`, `block`, `modify_response` |
|
||||
| `on_fail` | `string` | Action when guardrail fails: `next`, `allow`, `block`, `modify_response` |
|
||||
| `on_pass` | `string` | Action when outcome is **pass**: `next`, `allow`, `block`, `modify_response` |
|
||||
| `on_fail` | `string` | Action when outcome is **fail** (policy intervention): `next`, `allow`, `block`, `modify_response` |
|
||||
| `on_error` | `string` (optional) | Action when outcome is **error** (technical). If omitted, **error** uses `on_fail`. |
|
||||
| `pass_data` | `boolean` | Forward modified request data (e.g., PII-masked) to the next step |
|
||||
| `modify_response_message` | `string` | Custom message when using `modify_response` action |
|
||||
|
||||
|
|
@ -57,7 +71,7 @@ Each step defines what happens when the guardrail **passes** and when it **fails
|
|||
3. Select **Flow Builder** (instead of the simple form)
|
||||
4. Design your flow:
|
||||
- **Trigger** — Incoming LLM request (runs when the policy matches)
|
||||
- **Steps** — Add guardrails, set ON PASS and ON FAIL actions per step
|
||||
- **Steps** — Add guardrails, set **ON PASS**, **ON FAIL**, and **ON ERROR** actions per step (ON ERROR is optional; when unset, errors follow ON FAIL)
|
||||
- **End** — Request proceeds to the LLM
|
||||
5. Use the **+** between steps to insert new steps
|
||||
6. Use the **Test** panel to run sample messages through the pipeline before saving
|
||||
|
|
@ -151,6 +165,37 @@ policies:
|
|||
|
||||
First attempt passes → allow. First attempt fails → retry the same guardrail; second pass → allow, second fail → block.
|
||||
|
||||
## Technical errors vs policy failures (`on_error`)
|
||||
|
||||
Use **`on_error`** when you want different behavior for **API/infra problems** than for **content policy** violations.
|
||||
|
||||
- **`on_fail`** — Runs when the guardrail **intervenes** (e.g. toxic content, PII detected).
|
||||
- **`on_error`** — Runs when the step ends in **error** (timeout, connection failure, guardrail not loaded, etc.). If you omit `on_error`, **error** outcomes use **`on_fail`**.
|
||||
|
||||
Example: block on bad content, but if the primary scanner is down, fall back to a second guardrail instead of blocking every request:
|
||||
|
||||
```yaml
|
||||
policies:
|
||||
error-fallback-policy:
|
||||
guardrails:
|
||||
add:
|
||||
- primary_scanner
|
||||
- backup_scanner
|
||||
pipeline:
|
||||
mode: pre_call
|
||||
steps:
|
||||
- guardrail: primary_scanner
|
||||
on_pass: allow
|
||||
on_fail: block
|
||||
on_error: next
|
||||
- guardrail: backup_scanner
|
||||
on_pass: allow
|
||||
on_fail: block
|
||||
on_error: allow
|
||||
```
|
||||
|
||||
If `primary_scanner` errors → run `backup_scanner`. If `backup_scanner` errors → allow the request (set `on_error` to `block` if you prefer fail-closed).
|
||||
|
||||
## Example: Custom response on fail
|
||||
|
||||
Return a branded message instead of a generic block:
|
||||
|
|
|
|||
|
|
@ -316,86 +316,9 @@ general_settings:
|
|||
|
||||
## Health Check Driven Routing
|
||||
|
||||
By default, background health checks are observability-only — they populate the `/health` endpoint but don't affect routing. Unhealthy deployments still receive traffic until request failures trigger cooldown.
|
||||
Route traffic away from unhealthy deployments proactively — before user requests hit them. Supports per-error-type failure thresholds, transient error suppression, and automatic safety nets.
|
||||
|
||||
With `enable_health_check_routing: true`, the router **excludes deployments that failed their last background health check** before selecting a candidate. This gives you proactive failover instead of reactive cooldown.
|
||||
|
||||
### How it works
|
||||
|
||||
1. Background health checks run on their configured interval
|
||||
2. After each cycle, every deployment is marked healthy or unhealthy
|
||||
3. On each incoming request, the router filters out unhealthy deployments **before** cooldown filtering and load balancing
|
||||
4. If all deployments are unhealthy, the filter is bypassed (safety net — never causes a total outage)
|
||||
5. If health state is stale (older than `health_check_staleness_threshold`), it is ignored
|
||||
|
||||
### Quick start
|
||||
|
||||
```yaml
|
||||
model_list:
|
||||
- model_name: gpt-4
|
||||
litellm_params:
|
||||
model: openai/gpt-4
|
||||
api_key: os.environ/OPENAI_API_KEY
|
||||
- model_name: gpt-4
|
||||
litellm_params:
|
||||
model: openai/gpt-4
|
||||
api_key: os.environ/OPENAI_API_KEY_SECONDARY
|
||||
|
||||
general_settings:
|
||||
background_health_checks: true
|
||||
health_check_interval: 60
|
||||
enable_health_check_routing: true
|
||||
```
|
||||
|
||||
### Configuration
|
||||
|
||||
| Setting | Where | Default | Description |
|
||||
|---------|-------|---------|-------------|
|
||||
| `enable_health_check_routing` | `general_settings` | `false` | Enable/disable health-check-driven routing |
|
||||
| `health_check_staleness_threshold` | `general_settings` | `health_check_interval * 2` | Seconds before health state is considered stale and ignored |
|
||||
| `background_health_checks` | `general_settings` | `false` | Must be `true` for health check routing to work |
|
||||
| `health_check_interval` | `general_settings` | `300` | Seconds between health check cycles |
|
||||
|
||||
### Interaction with cooldown
|
||||
|
||||
Health check filtering and cooldown are **additive**. A deployment can be excluded by either mechanism:
|
||||
|
||||
- **Health check filter** — proactive, runs on the configured interval, excludes deployments that failed the last check
|
||||
- **Cooldown** — reactive, triggered by request failures, excludes deployments for a short TTL
|
||||
|
||||
This means request failures still provide fast detection between health check intervals.
|
||||
|
||||
### Staleness
|
||||
|
||||
If a health check result is older than `health_check_staleness_threshold`, it is ignored and the deployment is treated as eligible. This prevents stale data from permanently excluding a deployment if the health check loop stops or slows down.
|
||||
|
||||
The default staleness threshold is `health_check_interval * 2`. For a 60s interval, health state expires after 120s.
|
||||
|
||||
### Example: custom staleness
|
||||
|
||||
```yaml
|
||||
general_settings:
|
||||
background_health_checks: true
|
||||
health_check_interval: 30
|
||||
enable_health_check_routing: true
|
||||
health_check_staleness_threshold: 90 # ignore health state older than 90s
|
||||
```
|
||||
|
||||
### Debugging
|
||||
|
||||
Run the proxy with `--detailed_debug` and look for:
|
||||
|
||||
```
|
||||
health_check_routing_state_updated healthy=3 unhealthy=1
|
||||
```
|
||||
|
||||
This is logged after each health check cycle when routing state is written.
|
||||
|
||||
If the safety net triggers (all deployments unhealthy), you'll see:
|
||||
|
||||
```
|
||||
All deployments marked unhealthy by health checks, bypassing health filter
|
||||
```
|
||||
See the full guide: [Health Check Driven Routing](./health_check_routing.md)
|
||||
|
||||
## Health Check Timeout
|
||||
|
||||
|
|
|
|||
340
docs/my-website/docs/proxy/health_check_routing.md
Normal file
340
docs/my-website/docs/proxy/health_check_routing.md
Normal file
|
|
@ -0,0 +1,340 @@
|
|||
# Health Check Driven Routing
|
||||
|
||||
Route traffic away from unhealthy deployments before users hit errors. Background health checks run on a configurable interval, and any deployment that fails gets removed from the routing pool proactively, not after a user request already failed.
|
||||
|
||||
|
||||
## Architecture
|
||||
|
||||
<svg viewBox="0 0 860 600" xmlns="http://www.w3.org/2000/svg" style={{maxWidth: '100%', fontFamily: 'system-ui, sans-serif'}}>
|
||||
{/* Background */}
|
||||
<rect width="860" height="600" fill="#f8fafc" rx="12"/>
|
||||
|
||||
{/* LEFT PANEL: Background health check loop */}
|
||||
<rect x="20" y="20" width="240" height="560" fill="#eff6ff" rx="10" stroke="#bfdbfe" strokeWidth="1.5"/>
|
||||
<text x="140" y="48" textAnchor="middle" fill="#1d4ed8" fontSize="13" fontWeight="600">Background Loop</text>
|
||||
<text x="140" y="64" textAnchor="middle" fill="#3b82f6" fontSize="11">every health_check_interval seconds</text>
|
||||
|
||||
{/* Deployment A */}
|
||||
<rect x="40" y="82" width="200" height="50" fill="white" rx="8" stroke="#93c5fd" strokeWidth="1.5"/>
|
||||
<text x="140" y="102" textAnchor="middle" fill="#1e40af" fontSize="12" fontWeight="500">Deployment A</text>
|
||||
<text x="140" y="120" textAnchor="middle" fill="#64748b" fontSize="11">ahealth_check() → 200 ✓</text>
|
||||
|
||||
{/* Deployment B */}
|
||||
<rect x="40" y="148" width="200" height="50" fill="white" rx="8" stroke="#fca5a5" strokeWidth="1.5"/>
|
||||
<text x="140" y="168" textAnchor="middle" fill="#991b1b" fontSize="12" fontWeight="500">Deployment B</text>
|
||||
<text x="140" y="186" textAnchor="middle" fill="#64748b" fontSize="11">ahealth_check() → 401 ✗</text>
|
||||
|
||||
{/* Deployment C */}
|
||||
<rect x="40" y="214" width="200" height="50" fill="white" rx="8" stroke="#fde68a" strokeWidth="1.5"/>
|
||||
<text x="140" y="234" textAnchor="middle" fill="#92400e" fontSize="12" fontWeight="500">Deployment C</text>
|
||||
<text x="140" y="252" textAnchor="middle" fill="#64748b" fontSize="11">ahealth_check() → 429 ⚡</text>
|
||||
|
||||
{/* ignore_transient box */}
|
||||
<rect x="40" y="282" width="200" height="68" fill="#fefce8" rx="8" stroke="#fde047" strokeWidth="1.5"/>
|
||||
<text x="140" y="302" textAnchor="middle" fill="#713f12" fontSize="11" fontWeight="600">ignore_transient_errors: true</text>
|
||||
<text x="140" y="320" textAnchor="middle" fill="#92400e" fontSize="11">429 / 408 → ignored</text>
|
||||
<text x="140" y="338" textAnchor="middle" fill="#92400e" fontSize="11">not written to cache</text>
|
||||
|
||||
{/* allowed_fails_policy box */}
|
||||
<rect x="40" y="368" width="200" height="84" fill="#f0fdf4" rx="8" stroke="#86efac" strokeWidth="1.5"/>
|
||||
<text x="140" y="388" textAnchor="middle" fill="#166534" fontSize="11" fontWeight="600">allowed_fails_policy</text>
|
||||
<text x="140" y="406" textAnchor="middle" fill="#15803d" fontSize="11">401 → increment counter</text>
|
||||
<text x="140" y="424" textAnchor="middle" fill="#15803d" fontSize="11">counter > threshold</text>
|
||||
<text x="140" y="442" textAnchor="middle" fill="#15803d" fontSize="11">→ cooldown triggered</text>
|
||||
|
||||
{/* CENTER PANEL: Shared State */}
|
||||
<rect x="300" y="20" width="220" height="560" fill="#f5f3ff" rx="10" stroke="#c4b5fd" strokeWidth="1.5"/>
|
||||
<text x="410" y="48" textAnchor="middle" fill="#6d28d9" fontSize="13" fontWeight="600">Shared State</text>
|
||||
|
||||
{/* Health State Cache */}
|
||||
<rect x="320" y="62" width="180" height="116" fill="white" rx="8" stroke="#a78bfa" strokeWidth="1.5"/>
|
||||
<text x="410" y="84" textAnchor="middle" fill="#5b21b6" fontSize="12" fontWeight="600">DeploymentHealthCache</text>
|
||||
<text x="410" y="104" textAnchor="middle" fill="#64748b" fontSize="11">A → healthy ✓</text>
|
||||
<text x="410" y="122" textAnchor="middle" fill="#64748b" fontSize="11">B → unhealthy ✗</text>
|
||||
<text x="410" y="140" textAnchor="middle" fill="#64748b" fontSize="11">C → not written (ignored)</text>
|
||||
<text x="410" y="164" textAnchor="middle" fill="#94a3b8" fontSize="10">TTL: staleness_threshold × 1.5</text>
|
||||
|
||||
{/* Cooldown Cache */}
|
||||
<rect x="320" y="196" width="180" height="104" fill="white" rx="8" stroke="#a78bfa" strokeWidth="1.5"/>
|
||||
<text x="410" y="218" textAnchor="middle" fill="#5b21b6" fontSize="12" fontWeight="600">Cooldown Cache</text>
|
||||
<text x="410" y="238" textAnchor="middle" fill="#64748b" fontSize="11">B → cooling down</text>
|
||||
<text x="410" y="256" textAnchor="middle" fill="#64748b" fontSize="11">(after policy threshold)</text>
|
||||
<text x="410" y="278" textAnchor="middle" fill="#94a3b8" fontSize="10">TTL: cooldown_time</text>
|
||||
|
||||
{/* failed_calls counter */}
|
||||
<rect x="320" y="318" width="180" height="90" fill="white" rx="8" stroke="#a78bfa" strokeWidth="1.5"/>
|
||||
<text x="410" y="340" textAnchor="middle" fill="#5b21b6" fontSize="12" fontWeight="600">failed_calls counter</text>
|
||||
<text x="410" y="360" textAnchor="middle" fill="#64748b" fontSize="11">B: 2 / AuthAllowedFails: 1</text>
|
||||
<text x="410" y="378" textAnchor="middle" fill="#64748b" fontSize="11">→ threshold exceeded</text>
|
||||
<text x="410" y="398" textAnchor="middle" fill="#94a3b8" fontSize="10">TTL: cooldown_time (must > interval)</text>
|
||||
|
||||
{/* RIGHT PANEL: Request path */}
|
||||
<rect x="560" y="20" width="280" height="560" fill="#fff7ed" rx="10" stroke="#fed7aa" strokeWidth="1.5"/>
|
||||
<text x="700" y="48" textAnchor="middle" fill="#c2410c" fontSize="13" fontWeight="600">Request Path</text>
|
||||
|
||||
{/* Incoming request */}
|
||||
<rect x="580" y="62" width="240" height="38" fill="#fff" rx="7" stroke="#fb923c" strokeWidth="1.5"/>
|
||||
<text x="700" y="85" textAnchor="middle" fill="#9a3412" fontSize="12" fontWeight="500">Incoming request</text>
|
||||
|
||||
{/* All deployments */}
|
||||
<rect x="580" y="120" width="240" height="38" fill="#fff" rx="7" stroke="#fb923c" strokeWidth="1.5"/>
|
||||
<text x="700" y="143" textAnchor="middle" fill="#9a3412" fontSize="12">All deployments [A, B, C]</text>
|
||||
|
||||
<line x1="700" y1="100" x2="700" y2="120" stroke="#fb923c" strokeWidth="1.5" markerEnd="url(#arrow-orange)"/>
|
||||
|
||||
{/* Health check filter */}
|
||||
<rect x="580" y="178" width="240" height="62" fill="#fff" rx="7" stroke="#fb923c" strokeWidth="1.5"/>
|
||||
<text x="700" y="200" textAnchor="middle" fill="#9a3412" fontSize="12" fontWeight="600">① Health Check Filter</text>
|
||||
<text x="700" y="218" textAnchor="middle" fill="#64748b" fontSize="11">if policy set → bypass</text>
|
||||
<text x="700" y="234" textAnchor="middle" fill="#64748b" fontSize="11">else → remove unhealthy</text>
|
||||
|
||||
<line x1="700" y1="158" x2="700" y2="178" stroke="#fb923c" strokeWidth="1.5" markerEnd="url(#arrow-orange)"/>
|
||||
|
||||
{/* Cooldown filter */}
|
||||
<rect x="580" y="262" width="240" height="50" fill="#fff" rx="7" stroke="#fb923c" strokeWidth="1.5"/>
|
||||
<text x="700" y="284" textAnchor="middle" fill="#9a3412" fontSize="12" fontWeight="600">② Cooldown Filter</text>
|
||||
<text x="700" y="302" textAnchor="middle" fill="#64748b" fontSize="11">remove deployments in cooldown</text>
|
||||
|
||||
<line x1="700" y1="240" x2="700" y2="262" stroke="#fb923c" strokeWidth="1.5" markerEnd="url(#arrow-orange)"/>
|
||||
|
||||
{/* Safety net */}
|
||||
<rect x="580" y="334" width="240" height="52" fill="#fef9c3" rx="7" stroke="#fbbf24" strokeWidth="1.5"/>
|
||||
<text x="700" y="356" textAnchor="middle" fill="#713f12" fontSize="12" fontWeight="600">Safety Net</text>
|
||||
<text x="700" y="376" textAnchor="middle" fill="#713f12" fontSize="11">if all removed → return all</text>
|
||||
|
||||
<line x1="700" y1="312" x2="700" y2="334" stroke="#fb923c" strokeWidth="1.5" markerEnd="url(#arrow-orange)"/>
|
||||
|
||||
{/* Load balancer */}
|
||||
<rect x="580" y="408" width="240" height="38" fill="#fff" rx="7" stroke="#fb923c" strokeWidth="1.5"/>
|
||||
<text x="700" y="431" textAnchor="middle" fill="#9a3412" fontSize="12" fontWeight="600">③ Load Balancer</text>
|
||||
|
||||
<line x1="700" y1="386" x2="700" y2="408" stroke="#fb923c" strokeWidth="1.5" markerEnd="url(#arrow-orange)"/>
|
||||
|
||||
{/* Selected deployment */}
|
||||
<rect x="580" y="468" width="240" height="38" fill="#dcfce7" rx="7" stroke="#4ade80" strokeWidth="1.5"/>
|
||||
<text x="700" y="491" textAnchor="middle" fill="#14532d" fontSize="12" fontWeight="600">Selected: Deployment A ✓</text>
|
||||
|
||||
<line x1="700" y1="446" x2="700" y2="468" stroke="#4ade80" strokeWidth="1.5" markerEnd="url(#arrow-green)"/>
|
||||
|
||||
{/* ARROWS: left → center */}
|
||||
<line x1="240" y1="107" x2="320" y2="110" stroke="#3b82f6" strokeWidth="1.5" strokeDasharray="4,3" markerEnd="url(#arrow-blue)"/>
|
||||
<line x1="240" y1="173" x2="320" y2="240" stroke="#ef4444" strokeWidth="1.5" strokeDasharray="4,3" markerEnd="url(#arrow-red)"/>
|
||||
<line x1="240" y1="173" x2="320" y2="348" stroke="#ef4444" strokeWidth="1.5" strokeDasharray="4,3" markerEnd="url(#arrow-red)"/>
|
||||
<line x1="240" y1="316" x2="320" y2="130" stroke="#eab308" strokeWidth="1.5" strokeDasharray="4,3" markerEnd="url(#arrow-yellow)"/>
|
||||
|
||||
{/* ARROWS: center → right */}
|
||||
<line x1="500" y1="120" x2="580" y2="190" stroke="#8b5cf6" strokeWidth="1.5" strokeDasharray="4,3" markerEnd="url(#arrow-purple)"/>
|
||||
<line x1="500" y1="248" x2="580" y2="274" stroke="#8b5cf6" strokeWidth="1.5" strokeDasharray="4,3" markerEnd="url(#arrow-purple)"/>
|
||||
|
||||
{/* Arrow markers */}
|
||||
<defs>
|
||||
<marker id="arrow-orange" markerWidth="8" markerHeight="8" refX="6" refY="3" orient="auto">
|
||||
<path d="M0,0 L0,6 L8,3 z" fill="#fb923c"/>
|
||||
</marker>
|
||||
<marker id="arrow-blue" markerWidth="8" markerHeight="8" refX="6" refY="3" orient="auto">
|
||||
<path d="M0,0 L0,6 L8,3 z" fill="#3b82f6"/>
|
||||
</marker>
|
||||
<marker id="arrow-red" markerWidth="8" markerHeight="8" refX="6" refY="3" orient="auto">
|
||||
<path d="M0,0 L0,6 L8,3 z" fill="#ef4444"/>
|
||||
</marker>
|
||||
<marker id="arrow-yellow" markerWidth="8" markerHeight="8" refX="6" refY="3" orient="auto">
|
||||
<path d="M0,0 L0,6 L8,3 z" fill="#eab308"/>
|
||||
</marker>
|
||||
<marker id="arrow-purple" markerWidth="8" markerHeight="8" refX="6" refY="3" orient="auto">
|
||||
<path d="M0,0 L0,6 L8,3 z" fill="#8b5cf6"/>
|
||||
</marker>
|
||||
<marker id="arrow-green" markerWidth="8" markerHeight="8" refX="6" refY="3" orient="auto">
|
||||
<path d="M0,0 L0,6 L8,3 z" fill="#4ade80"/>
|
||||
</marker>
|
||||
</defs>
|
||||
</svg>
|
||||
|
||||
|
||||
## What problem does this solve?
|
||||
|
||||
By default, LiteLLM routes traffic to all deployments and only stops sending to a broken one after it has already failed a user request. The cooldown system is reactive.
|
||||
|
||||
Health check driven routing makes this **proactive**: a background loop pings every deployment on a configurable interval. If a deployment fails its health check, it gets removed from the routing pool immediately, before a user request lands on it.
|
||||
|
||||
When you also set `allowed_fails_policy`, you control exactly how many health check failures of each error type (auth errors, rate limits, timeouts) are needed before a deployment enters cooldown. This avoids false positives from transient noise.
|
||||
|
||||
|
||||
## Setup
|
||||
|
||||
### Step 1: Enable background health checks
|
||||
|
||||
Background health checks are off by default. Turn them on in `general_settings`:
|
||||
|
||||
```yaml
|
||||
general_settings:
|
||||
background_health_checks: true
|
||||
health_check_interval: 60 # seconds between each full check cycle
|
||||
```
|
||||
|
||||
### Step 2: Enable health check routing
|
||||
|
||||
```yaml
|
||||
general_settings:
|
||||
background_health_checks: true
|
||||
health_check_interval: 60
|
||||
enable_health_check_routing: true # ← route away from unhealthy deployments
|
||||
```
|
||||
|
||||
At this point, any deployment that fails its health check is immediately excluded from routing until the next check cycle clears it.
|
||||
|
||||
### Step 3: Add a policy to control how many failures trigger cooldown
|
||||
|
||||
Without a policy, the first health check failure marks a deployment as unhealthy. If you want more tolerance (e.g., only act after 2 consecutive auth failures), use `allowed_fails_policy`:
|
||||
|
||||
```yaml
|
||||
model_list:
|
||||
- model_name: claude-sonnet
|
||||
litellm_params:
|
||||
model: anthropic/claude-sonnet-4-5
|
||||
api_key: os.environ/ANTHROPIC_API_KEY
|
||||
|
||||
- model_name: claude-sonnet
|
||||
litellm_params:
|
||||
model: anthropic/claude-sonnet-4-5
|
||||
api_key: os.environ/ANTHROPIC_API_KEY_SECONDARY
|
||||
|
||||
general_settings:
|
||||
background_health_checks: true
|
||||
health_check_interval: 30
|
||||
enable_health_check_routing: true
|
||||
|
||||
router_settings:
|
||||
cooldown_time: 60 # how long a deployment stays in cooldown
|
||||
allowed_fails_policy:
|
||||
AuthenticationErrorAllowedFails: 1 # cooldown after 2nd auth failure
|
||||
TimeoutErrorAllowedFails: 3 # cooldown after 4th timeout
|
||||
```
|
||||
|
||||
When `allowed_fails_policy` is set, the binary health check filter is bypassed. Only the cooldown system controls routing exclusion, and it only fires after your configured threshold is crossed.
|
||||
|
||||
### Step 4 (optional): Ignore transient errors
|
||||
|
||||
429 (rate limit) and 408 (timeout) from a health check usually mean the deployment is temporarily overloaded, not broken. To prevent these from affecting routing at all:
|
||||
|
||||
```yaml
|
||||
general_settings:
|
||||
background_health_checks: true
|
||||
health_check_interval: 30
|
||||
enable_health_check_routing: true
|
||||
health_check_ignore_transient_errors: true # 429 and 408 never affect routing
|
||||
```
|
||||
|
||||
With this on, only hard failures (401, 404, 5xx) from health checks contribute to cooldown.
|
||||
|
||||
|
||||
## Full example
|
||||
|
||||
```yaml
|
||||
model_list:
|
||||
- model_name: gpt-4o
|
||||
litellm_params:
|
||||
model: openai/gpt-4o
|
||||
api_key: os.environ/OPENAI_API_KEY
|
||||
|
||||
- model_name: gpt-4o
|
||||
litellm_params:
|
||||
model: openai/gpt-4o
|
||||
api_key: os.environ/OPENAI_API_KEY_SECONDARY
|
||||
|
||||
- model_name: gpt-4o
|
||||
litellm_params:
|
||||
model: azure/gpt-4o
|
||||
api_base: os.environ/AZURE_API_BASE
|
||||
api_key: os.environ/AZURE_API_KEY
|
||||
|
||||
general_settings:
|
||||
background_health_checks: true
|
||||
health_check_interval: 30
|
||||
enable_health_check_routing: true
|
||||
health_check_ignore_transient_errors: true
|
||||
|
||||
router_settings:
|
||||
cooldown_time: 60
|
||||
allowed_fails_policy:
|
||||
AuthenticationErrorAllowedFails: 0 # cooldown immediately on auth failure
|
||||
TimeoutErrorAllowedFails: 2 # cooldown after 3 timeouts
|
||||
RateLimitErrorAllowedFails: 5 # cooldown after 6 rate limits (if not ignoring transients)
|
||||
```
|
||||
|
||||
|
||||
## Configuration reference
|
||||
|
||||
| Setting | Where | Default | Description |
|
||||
|---|---|---|---|
|
||||
| `enable_health_check_routing` | `general_settings` | `false` | Route away from deployments that fail health checks |
|
||||
| `background_health_checks` | `general_settings` | `false` | Must be `true` for health check routing to work |
|
||||
| `health_check_interval` | `general_settings` | `300` | Seconds between full health check cycles |
|
||||
| `health_check_staleness_threshold` | `general_settings` | `interval x 2` | Seconds before cached health state is ignored |
|
||||
| `health_check_ignore_transient_errors` | `general_settings` | `false` | Ignore 429 and 408 from health checks; these never affect routing |
|
||||
| `cooldown_time` | `router_settings` | `5` | Seconds a deployment stays in cooldown after threshold is crossed |
|
||||
| `allowed_fails_policy` | `router_settings` | `null` | Per-error-type failure thresholds before cooldown (see below) |
|
||||
|
||||
### `allowed_fails_policy` fields
|
||||
|
||||
| Field | Error type | HTTP status |
|
||||
|---|---|---|
|
||||
| `AuthenticationErrorAllowedFails` | Bad API key | 401 |
|
||||
| `TimeoutErrorAllowedFails` | Request timeout | 408 |
|
||||
| `RateLimitErrorAllowedFails` | Rate limit exceeded | 429 |
|
||||
| `BadRequestErrorAllowedFails` | Malformed request | 400 |
|
||||
| `ContentPolicyViolationErrorAllowedFails` | Content filtered | 400 |
|
||||
|
||||
The value is the number of failures **tolerated** before cooldown. `0` means cooldown on the first failure. `2` means cooldown on the third.
|
||||
|
||||
|
||||
## Things to keep in mind
|
||||
|
||||
- **Counter TTL must be longer than the health check interval.** `allowed_fails_policy` works by incrementing a `failed_calls` counter per deployment. That counter expires after `cooldown_time` seconds. If `cooldown_time` is shorter than `health_check_interval`, the counter resets between every check cycle and failures never accumulate. Set `cooldown_time` greater than `health_check_interval` when using `allowed_fails_policy`.
|
||||
|
||||
```yaml
|
||||
router_settings:
|
||||
cooldown_time: 60 # must be > health_check_interval (30s here)
|
||||
|
||||
general_settings:
|
||||
health_check_interval: 30
|
||||
```
|
||||
|
||||
- **`AllowedFails: N` means cooldown on the (N+1)th failure.** The counter check is `updated_fails > allowed_fails`, so `0` triggers on the 1st failure, `1` on the 2nd, `2` on the 3rd.
|
||||
|
||||
| `AllowedFails` | Cooldown triggers after |
|
||||
|---|---|
|
||||
| `0` | 1st failure |
|
||||
| `1` | 2nd failure |
|
||||
| `2` | 3rd failure |
|
||||
|
||||
- **Without `allowed_fails_policy`, the first failure is enough.** The first failed health check immediately excludes the deployment from routing. Use `allowed_fails_policy` when you want tolerance for flaky checks.
|
||||
|
||||
- **If all deployments are unhealthy, the filter is bypassed.** Traffic keeps flowing rather than returning no deployment at all. Requests will fail, but the router keeps trying.
|
||||
|
||||
- **Health check failures and request failures share the same counters.** When `allowed_fails_policy` is set, both sources increment the same `failed_calls` counter. A deployment at 1 health check failure that then receives 1 failing request will hit the threshold for `AllowedFails: 1` and enter cooldown.
|
||||
|
||||
|
||||
## Debugging
|
||||
|
||||
Run the proxy with `--detailed_debug` and look for these log lines:
|
||||
|
||||
After each health check cycle (written at DEBUG level):
|
||||
```
|
||||
health_check_routing_state_updated healthy=2 unhealthy=1
|
||||
```
|
||||
|
||||
When a health check failure increments the counter and triggers cooldown (DEBUG level):
|
||||
```
|
||||
checks 'should_run_cooldown_logic'
|
||||
Attempting to add <deployment_id> to cooldown list
|
||||
```
|
||||
|
||||
When safety net fires because all deployments are in cooldown:
|
||||
```
|
||||
All deployments in cooldown via health-check routing, bypassing cooldown filter
|
||||
```
|
||||
|
||||
When safety net fires because all deployments are unhealthy (binary filter, no `allowed_fails_policy`):
|
||||
```
|
||||
All deployments marked unhealthy by health checks, bypassing health filter
|
||||
```
|
||||
|
|
@ -1051,7 +1051,8 @@ const sidebars = {
|
|||
"proxy/fallback_management",
|
||||
"proxy/tag_routing",
|
||||
"proxy/timeout",
|
||||
"wildcard_routing"
|
||||
"wildcard_routing",
|
||||
"proxy/health_check_routing"
|
||||
],
|
||||
},
|
||||
{
|
||||
|
|
|
|||
|
|
@ -1319,6 +1319,9 @@ LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS = int(
|
|||
LITELLM_KEY_ROTATION_GRACE_PERIOD: str = os.getenv(
|
||||
"LITELLM_KEY_ROTATION_GRACE_PERIOD", ""
|
||||
) # Duration to keep old key valid after rotation (e.g. "24h", "2d"); empty = immediate revoke (default)
|
||||
LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS = int(
|
||||
os.getenv("LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS", 600)
|
||||
) # 10 minutes default — caps the deadlock window if a pod crashes mid-rotation
|
||||
UI_SESSION_TOKEN_TEAM_ID = "litellm-dashboard"
|
||||
LITELLM_PROXY_ADMIN_NAME = "default_user_id"
|
||||
|
||||
|
|
@ -1347,6 +1350,7 @@ CLOUDZERO_MAX_FETCHED_DATA_RECORDS = int(
|
|||
os.getenv("CLOUDZERO_MAX_FETCHED_DATA_RECORDS", 50000)
|
||||
)
|
||||
SPEND_LOG_CLEANUP_JOB_NAME = "spend_log_cleanup"
|
||||
KEY_ROTATION_JOB_NAME = "litellm_key_rotation_job"
|
||||
SPEND_LOG_RUN_LOOPS = int(os.getenv("SPEND_LOG_RUN_LOOPS", 500))
|
||||
SPEND_LOG_CLEANUP_BATCH_SIZE = int(os.getenv("SPEND_LOG_CLEANUP_BATCH_SIZE", 1000))
|
||||
SPEND_LOG_QUEUE_SIZE_THRESHOLD = int(os.getenv("SPEND_LOG_QUEUE_SIZE_THRESHOLD", 100))
|
||||
|
|
|
|||
|
|
@ -3792,9 +3792,9 @@ def completion( # type: ignore # noqa: PLR0915
|
|||
"aws_region_name" not in optional_params
|
||||
or optional_params["aws_region_name"] is None
|
||||
):
|
||||
optional_params[
|
||||
"aws_region_name"
|
||||
] = aws_bedrock_client.meta.region_name
|
||||
optional_params["aws_region_name"] = (
|
||||
aws_bedrock_client.meta.region_name
|
||||
)
|
||||
|
||||
bedrock_route = BedrockModelInfo.get_bedrock_route(model)
|
||||
if bedrock_route == "converse":
|
||||
|
|
@ -6214,9 +6214,9 @@ def adapter_completion(
|
|||
new_kwargs = translation_obj.translate_completion_input_params(kwargs=kwargs)
|
||||
|
||||
response: Union[ModelResponse, CustomStreamWrapper] = completion(**new_kwargs) # type: ignore
|
||||
translated_response: Optional[
|
||||
Union[BaseModel, AdapterCompletionStreamWrapper]
|
||||
] = None
|
||||
translated_response: Optional[Union[BaseModel, AdapterCompletionStreamWrapper]] = (
|
||||
None
|
||||
)
|
||||
if isinstance(response, ModelResponse):
|
||||
translated_response = translation_obj.translate_completion_output_params(
|
||||
response=response
|
||||
|
|
@ -6396,9 +6396,9 @@ async def atranscription(*args, **kwargs) -> TranscriptionResponse:
|
|||
if existing_duration is None:
|
||||
calculated_duration = calculate_request_duration(file)
|
||||
if calculated_duration is not None:
|
||||
response._hidden_params[
|
||||
"audio_transcription_duration"
|
||||
] = calculated_duration
|
||||
response._hidden_params["audio_transcription_duration"] = (
|
||||
calculated_duration
|
||||
)
|
||||
|
||||
return response
|
||||
except Exception as e:
|
||||
|
|
@ -6621,9 +6621,9 @@ def transcription(
|
|||
if existing_duration is None:
|
||||
calculated_duration = calculate_request_duration(file)
|
||||
if calculated_duration is not None:
|
||||
response._hidden_params[
|
||||
"audio_transcription_duration"
|
||||
] = calculated_duration
|
||||
response._hidden_params["audio_transcription_duration"] = (
|
||||
calculated_duration
|
||||
)
|
||||
|
||||
if response is None:
|
||||
raise ValueError("Unmapped provider passed in. Unable to get the response.")
|
||||
|
|
@ -6927,9 +6927,9 @@ def speech( # noqa: PLR0915
|
|||
ElevenLabsTextToSpeechConfig.ELEVENLABS_QUERY_PARAMS_KEY
|
||||
] = query_params
|
||||
|
||||
litellm_params_dict[
|
||||
ElevenLabsTextToSpeechConfig.ELEVENLABS_VOICE_ID_KEY
|
||||
] = voice_id
|
||||
litellm_params_dict[ElevenLabsTextToSpeechConfig.ELEVENLABS_VOICE_ID_KEY] = (
|
||||
voice_id
|
||||
)
|
||||
|
||||
if api_base is not None:
|
||||
litellm_params_dict["api_base"] = api_base
|
||||
|
|
@ -7250,7 +7250,8 @@ async def ahealth_check(
|
|||
|
||||
if mode is None:
|
||||
return {
|
||||
"error": f"error:{str(e)}. Missing `mode`. Set the `mode` for the model - https://docs.litellm.ai/docs/proxy/health#embedding-models \nstacktrace: {stack_trace}"
|
||||
"error": f"error:{str(e)}. Missing `mode`. Set the `mode` for the model - https://docs.litellm.ai/docs/proxy/health#embedding-models \nstacktrace: {stack_trace}",
|
||||
"exception": e,
|
||||
}
|
||||
|
||||
error_to_return = str(e) + "\nstack trace: " + stack_trace
|
||||
|
|
@ -7262,6 +7263,7 @@ async def ahealth_check(
|
|||
return {
|
||||
"error": error_to_return,
|
||||
"raw_request_typed_dict": raw_request_typed_dict,
|
||||
"exception": e,
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -7508,9 +7510,9 @@ def stream_chunk_builder( # noqa: PLR0915
|
|||
]
|
||||
|
||||
if len(content_chunks) > 0:
|
||||
response["choices"][0]["message"][
|
||||
"content"
|
||||
] = processor.get_combined_content(content_chunks)
|
||||
response["choices"][0]["message"]["content"] = (
|
||||
processor.get_combined_content(content_chunks)
|
||||
)
|
||||
|
||||
thinking_blocks = [
|
||||
chunk
|
||||
|
|
@ -7521,9 +7523,9 @@ def stream_chunk_builder( # noqa: PLR0915
|
|||
]
|
||||
|
||||
if len(thinking_blocks) > 0:
|
||||
response["choices"][0]["message"][
|
||||
"thinking_blocks"
|
||||
] = processor.get_combined_thinking_content(thinking_blocks)
|
||||
response["choices"][0]["message"]["thinking_blocks"] = (
|
||||
processor.get_combined_thinking_content(thinking_blocks)
|
||||
)
|
||||
|
||||
reasoning_chunks = [
|
||||
chunk
|
||||
|
|
@ -7534,9 +7536,9 @@ def stream_chunk_builder( # noqa: PLR0915
|
|||
]
|
||||
|
||||
if len(reasoning_chunks) > 0:
|
||||
response["choices"][0]["message"][
|
||||
"reasoning_content"
|
||||
] = processor.get_combined_reasoning_content(reasoning_chunks)
|
||||
response["choices"][0]["message"]["reasoning_content"] = (
|
||||
processor.get_combined_reasoning_content(reasoning_chunks)
|
||||
)
|
||||
|
||||
annotation_chunks = [
|
||||
chunk
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ from litellm._logging import verbose_proxy_logger
|
|||
from litellm.constants import (
|
||||
LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME,
|
||||
LITELLM_KEY_ROTATION_GRACE_PERIOD,
|
||||
LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS,
|
||||
)
|
||||
from litellm.proxy._types import (
|
||||
GenerateKeyResponse,
|
||||
|
|
@ -30,14 +31,42 @@ class KeyRotationManager:
|
|||
Manages automated key rotation based on individual key rotation schedules.
|
||||
"""
|
||||
|
||||
def __init__(self, prisma_client: PrismaClient):
|
||||
def __init__(self, prisma_client: PrismaClient, pod_lock_manager=None):
|
||||
self.prisma_client = prisma_client
|
||||
self.pod_lock_manager = pod_lock_manager
|
||||
|
||||
async def process_rotations(self):
|
||||
"""
|
||||
Main entry point - find and rotate keys that are due for rotation
|
||||
Main entry point - find and rotate keys that are due for rotation.
|
||||
Uses PodLockManager to ensure only one pod runs rotation in multi-pod deployments.
|
||||
"""
|
||||
from litellm.constants import KEY_ROTATION_JOB_NAME
|
||||
|
||||
lock_acquired = False
|
||||
try:
|
||||
# If we have a pod lock manager with Redis, try to acquire the lock
|
||||
if self.pod_lock_manager and self.pod_lock_manager.redis_cache:
|
||||
# Use a dedicated lock TTL (default 600s) instead of the check interval
|
||||
# (which defaults to 86400s / 24h). Using the check interval would create
|
||||
# a 24-hour deadlock window if a pod crashes before releasing the lock.
|
||||
lock_ttl = max(
|
||||
LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS, 300
|
||||
) # At least 5 minutes, configurable via LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS
|
||||
lock_acquired = (
|
||||
await self.pod_lock_manager.acquire_lock(
|
||||
cronjob_id=KEY_ROTATION_JOB_NAME,
|
||||
ttl=lock_ttl,
|
||||
)
|
||||
or False
|
||||
)
|
||||
if not lock_acquired:
|
||||
verbose_proxy_logger.warning(
|
||||
"Key rotation: another pod is already running rotation "
|
||||
"or Redis lock acquisition failed — skipping this cycle. "
|
||||
"Keys will be rotated on the next cycle."
|
||||
)
|
||||
return
|
||||
|
||||
verbose_proxy_logger.info("Starting scheduled key rotation check...")
|
||||
|
||||
# Clean up expired deprecated keys first
|
||||
|
|
@ -74,6 +103,16 @@ class KeyRotationManager:
|
|||
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(f"Key rotation process failed: {e}")
|
||||
finally:
|
||||
# Only release the lock if it was actually acquired
|
||||
if (
|
||||
lock_acquired
|
||||
and self.pod_lock_manager
|
||||
and self.pod_lock_manager.redis_cache
|
||||
):
|
||||
await self.pod_lock_manager.release_lock(
|
||||
cronjob_id=KEY_ROTATION_JOB_NAME,
|
||||
)
|
||||
|
||||
async def _find_keys_needing_rotation(self) -> List[LiteLLM_VerificationToken]:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -32,6 +32,7 @@ class PodLockManager:
|
|||
async def acquire_lock(
|
||||
self,
|
||||
cronjob_id: str,
|
||||
ttl: Optional[int] = None,
|
||||
) -> Optional[bool]:
|
||||
"""
|
||||
Attempt to acquire the lock for a specific cron job using Redis.
|
||||
|
|
@ -39,15 +40,20 @@ class PodLockManager:
|
|||
|
||||
Args:
|
||||
cronjob_id: The ID of the cron job to lock
|
||||
ttl: Optional custom TTL in seconds. Defaults to DEFAULT_CRON_JOB_LOCK_TTL_SECONDS.
|
||||
Use a longer TTL for jobs that may take longer than the default 60s
|
||||
(e.g. key rotation with many keys).
|
||||
"""
|
||||
if self.redis_cache is None:
|
||||
verbose_proxy_logger.debug("redis_cache is None, skipping acquire_lock")
|
||||
return None
|
||||
try:
|
||||
lock_ttl = ttl or DEFAULT_CRON_JOB_LOCK_TTL_SECONDS
|
||||
verbose_proxy_logger.debug(
|
||||
"Pod %s attempting to acquire Redis lock for cronjob_id=%s",
|
||||
"Pod %s attempting to acquire Redis lock for cronjob_id=%s (ttl=%ds)",
|
||||
self.pod_id,
|
||||
cronjob_id,
|
||||
lock_ttl,
|
||||
)
|
||||
# Try to set the lock key with the pod_id as its value, only if it doesn't exist (NX)
|
||||
# and with an expiration (EX) to avoid deadlocks.
|
||||
|
|
@ -56,7 +62,7 @@ class PodLockManager:
|
|||
lock_key,
|
||||
self.pod_id,
|
||||
nx=True,
|
||||
ttl=DEFAULT_CRON_JOB_LOCK_TTL_SECONDS,
|
||||
ttl=lock_ttl,
|
||||
)
|
||||
if acquired:
|
||||
verbose_proxy_logger.info(
|
||||
|
|
@ -133,11 +139,10 @@ class PodLockManager:
|
|||
)
|
||||
else:
|
||||
verbose_proxy_logger.warning(
|
||||
"Spend tracking - pod %s failed to release Redis lock for cronjob_id=%s. "
|
||||
"Lock will expire after TTL=%ds.",
|
||||
"Pod %s failed to release Redis lock for cronjob_id=%s. "
|
||||
"Lock will expire after its TTL.",
|
||||
self.pod_id,
|
||||
cronjob_id,
|
||||
DEFAULT_CRON_JOB_LOCK_TTL_SECONDS,
|
||||
)
|
||||
else:
|
||||
verbose_proxy_logger.debug(
|
||||
|
|
|
|||
|
|
@ -21,6 +21,8 @@ ILLEGAL_DISPLAY_PARAMS = [
|
|||
"vertex_credentials",
|
||||
"aws_access_key_id",
|
||||
"aws_secret_access_key",
|
||||
"exception", # internal; not JSON-serializable, never for display
|
||||
"litellm_metadata", # internal tracking metadata with auth objects; not for display
|
||||
]
|
||||
|
||||
MINIMAL_DISPLAY_PARAMS = ["model", "mode_error"]
|
||||
|
|
@ -95,7 +97,12 @@ async def run_with_timeout(task, timeout):
|
|||
except asyncio.TimeoutError:
|
||||
# `asyncio.wait_for()` already cancels only the awaited task on timeout.
|
||||
# Do not cancel unrelated sibling health check tasks.
|
||||
return {"error": "Timeout exceeded"}
|
||||
timeout_exception = litellm.Timeout(
|
||||
message="Health check timeout exceeded",
|
||||
model="",
|
||||
llm_provider="",
|
||||
)
|
||||
return {"error": "Timeout exceeded", "exception": timeout_exception}
|
||||
|
||||
|
||||
async def _run_model_health_check(model: dict):
|
||||
|
|
@ -204,6 +211,10 @@ async def _perform_health_check(
|
|||
|
||||
healthy_endpoints = []
|
||||
unhealthy_endpoints = []
|
||||
# Exceptions keyed by model_id; returned separately so callers can use
|
||||
# them for cooldown integration without risking JSON-serialization errors
|
||||
# in the /health response.
|
||||
exceptions_by_model_id: dict = {}
|
||||
|
||||
for is_healthy, model in zip(results, model_list):
|
||||
litellm_params = model["litellm_params"]
|
||||
|
|
@ -218,14 +229,23 @@ async def _perform_health_check(
|
|||
cleaned = _clean_endpoint_data({**litellm_params, **is_healthy}, details)
|
||||
if _model_id:
|
||||
cleaned["model_id"] = _model_id
|
||||
if "exception" in is_healthy:
|
||||
exc = is_healthy["exception"]
|
||||
exceptions_by_model_id[_model_id] = exc
|
||||
# Store integer status code so shared-cache readers can
|
||||
# reconstruct the transient-error filter without the exception object.
|
||||
cleaned["exception_status"] = getattr(exc, "status_code", 500)
|
||||
unhealthy_endpoints.append(cleaned)
|
||||
else:
|
||||
cleaned = _clean_endpoint_data(litellm_params, details)
|
||||
if _model_id:
|
||||
cleaned["model_id"] = _model_id
|
||||
if isinstance(is_healthy, Exception):
|
||||
exceptions_by_model_id[_model_id] = is_healthy
|
||||
cleaned["exception_status"] = getattr(is_healthy, "status_code", 500)
|
||||
unhealthy_endpoints.append(cleaned)
|
||||
|
||||
return healthy_endpoints, unhealthy_endpoints
|
||||
return healthy_endpoints, unhealthy_endpoints, exceptions_by_model_id
|
||||
|
||||
|
||||
def build_deployment_health_states(
|
||||
|
|
@ -366,7 +386,7 @@ async def perform_health_check(
|
|||
source,
|
||||
cycle_id,
|
||||
)
|
||||
return [], []
|
||||
return [], [], {}
|
||||
|
||||
cycle_start_time = time.monotonic()
|
||||
requested_model_count = len(model_list)
|
||||
|
|
@ -406,7 +426,11 @@ async def perform_health_check(
|
|||
)
|
||||
|
||||
try:
|
||||
healthy_endpoints, unhealthy_endpoints = await _perform_health_check(
|
||||
(
|
||||
healthy_endpoints,
|
||||
unhealthy_endpoints,
|
||||
exceptions_by_model_id,
|
||||
) = await _perform_health_check(
|
||||
model_list,
|
||||
details,
|
||||
max_concurrency=max_concurrency,
|
||||
|
|
@ -438,4 +462,4 @@ async def perform_health_check(
|
|||
_rss_mb_for_log(),
|
||||
)
|
||||
|
||||
return healthy_endpoints, unhealthy_endpoints
|
||||
return healthy_endpoints, unhealthy_endpoints, exceptions_by_model_id
|
||||
|
|
|
|||
|
|
@ -192,7 +192,7 @@ class SharedHealthCheckManager:
|
|||
model_list: List[Dict[str, Any]],
|
||||
details: bool = True,
|
||||
max_concurrency: Optional[int] = None,
|
||||
) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
|
||||
) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], Dict[str, Any]]:
|
||||
"""
|
||||
Perform health check with shared state coordination.
|
||||
|
||||
|
|
@ -217,6 +217,7 @@ class SharedHealthCheckManager:
|
|||
return (
|
||||
cached_results.get("healthy_endpoints", []),
|
||||
cached_results.get("unhealthy_endpoints", []),
|
||||
{},
|
||||
)
|
||||
|
||||
# No recent cache, try to acquire lock
|
||||
|
|
@ -231,7 +232,11 @@ class SharedHealthCheckManager:
|
|||
len(model_list),
|
||||
)
|
||||
|
||||
healthy_endpoints, unhealthy_endpoints = await perform_health_check(
|
||||
(
|
||||
healthy_endpoints,
|
||||
unhealthy_endpoints,
|
||||
exceptions_by_model_id,
|
||||
) = await perform_health_check(
|
||||
model_list=model_list,
|
||||
details=details,
|
||||
max_concurrency=max_concurrency,
|
||||
|
|
@ -242,7 +247,7 @@ class SharedHealthCheckManager:
|
|||
healthy_endpoints, unhealthy_endpoints
|
||||
)
|
||||
|
||||
return healthy_endpoints, unhealthy_endpoints
|
||||
return healthy_endpoints, unhealthy_endpoints, exceptions_by_model_id
|
||||
|
||||
finally:
|
||||
# Always release the lock
|
||||
|
|
@ -262,6 +267,7 @@ class SharedHealthCheckManager:
|
|||
return (
|
||||
cached_results.get("healthy_endpoints", []),
|
||||
cached_results.get("unhealthy_endpoints", []),
|
||||
{},
|
||||
)
|
||||
|
||||
# Still no cache, fall back to local health check
|
||||
|
|
|
|||
|
|
@ -771,7 +771,7 @@ async def _perform_health_check_and_save(
|
|||
max_concurrency=None,
|
||||
):
|
||||
"""Helper function to perform health check and save results to database"""
|
||||
healthy_endpoints, unhealthy_endpoints = await perform_health_check(
|
||||
healthy_endpoints, unhealthy_endpoints, _ = await perform_health_check(
|
||||
model_list=model_list,
|
||||
cli_model=cli_model,
|
||||
model=target_model,
|
||||
|
|
|
|||
|
|
@ -74,7 +74,7 @@ class PipelineExecutor:
|
|||
|
||||
duration = time.perf_counter() - start_time
|
||||
|
||||
action = step.on_pass if outcome == "pass" else step.on_fail
|
||||
action = _pipeline_action_for_outcome(step, outcome)
|
||||
|
||||
step_result = PipelineStepResult(
|
||||
guardrail_name=step.guardrail,
|
||||
|
|
@ -206,6 +206,23 @@ class PipelineExecutor:
|
|||
return None
|
||||
|
||||
|
||||
def _pipeline_action_for_outcome(step: PipelineStep, outcome: str) -> str:
|
||||
"""
|
||||
Map pipeline step outcome to the configured action.
|
||||
|
||||
- pass -> on_pass
|
||||
- fail -> on_fail (content/policy intervention)
|
||||
- error -> on_error if set, else on_fail (backward compatible)
|
||||
"""
|
||||
if outcome == "pass":
|
||||
return step.on_pass
|
||||
if outcome == "fail":
|
||||
return step.on_fail
|
||||
if step.on_error is not None:
|
||||
return step.on_error
|
||||
return step.on_fail
|
||||
|
||||
|
||||
def _extract_error_message(e: Exception) -> str:
|
||||
"""Extract a human-readable error message from a guardrail exception."""
|
||||
if isinstance(e, ModifyResponseException):
|
||||
|
|
|
|||
|
|
@ -635,9 +635,9 @@ except ImportError:
|
|||
server_root_path = get_server_root_path()
|
||||
_license_check = LicenseCheck()
|
||||
premium_user: bool = _license_check.is_premium()
|
||||
premium_user_data: Optional[
|
||||
"EnterpriseLicenseData"
|
||||
] = _license_check.airgapped_license_data
|
||||
premium_user_data: Optional["EnterpriseLicenseData"] = (
|
||||
_license_check.airgapped_license_data
|
||||
)
|
||||
global_max_parallel_request_retries_env: Optional[str] = os.getenv(
|
||||
"LITELLM_GLOBAL_MAX_PARALLEL_REQUEST_RETRIES"
|
||||
)
|
||||
|
|
@ -1532,9 +1532,9 @@ master_key: Optional[str] = None
|
|||
config_agents: Optional[List[AgentConfig]] = None
|
||||
otel_logging = False
|
||||
prisma_client: Optional[PrismaClient] = None
|
||||
shared_aiohttp_session: Optional[
|
||||
"ClientSession"
|
||||
] = None # Global shared session for connection reuse
|
||||
shared_aiohttp_session: Optional["ClientSession"] = (
|
||||
None # Global shared session for connection reuse
|
||||
)
|
||||
user_api_key_cache = DualCache(
|
||||
default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value
|
||||
)
|
||||
|
|
@ -1545,13 +1545,13 @@ 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
|
||||
)
|
||||
polling_via_cache_enabled: Union[Literal["all"], List[str], bool] = False
|
||||
native_background_mode: List[
|
||||
str
|
||||
] = [] # Models that should use native provider background mode instead of polling
|
||||
native_background_mode: List[str] = (
|
||||
[]
|
||||
) # Models that should use native provider background mode instead of polling
|
||||
polling_cache_ttl: int = 3600 # Default 1 hour TTL for polling cache
|
||||
user_custom_auth = None
|
||||
user_custom_key_generate = None
|
||||
|
|
@ -2036,9 +2036,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
|
||||
|
|
@ -2248,23 +2248,54 @@ def _schedule_background_health_check_db_save(
|
|||
)
|
||||
|
||||
|
||||
def _get_endpoint_exception_status(endpoint: dict, exceptions: dict) -> int:
|
||||
"""Return the HTTP status code for an unhealthy endpoint.
|
||||
|
||||
Prefers the live exception object in `exceptions` (direct health check path).
|
||||
Falls back to the `exception_status` integer stored on the endpoint dict
|
||||
(shared-cache path, where exception objects are not available).
|
||||
"""
|
||||
model_id = endpoint.get("model_id")
|
||||
exc = exceptions.get(model_id) if model_id else None
|
||||
if exc is not None:
|
||||
return getattr(exc, "status_code", 500)
|
||||
return endpoint.get("exception_status", 500)
|
||||
|
||||
|
||||
def _write_health_state_to_router_cache(
|
||||
healthy_endpoints: list,
|
||||
unhealthy_endpoints: list,
|
||||
exceptions_by_model_id: Optional[dict] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Write deployment health states to the router's health state cache
|
||||
for health-check-driven routing. No-op if the feature is disabled.
|
||||
"""
|
||||
from litellm.proxy.health_check import build_deployment_health_states
|
||||
from litellm.router_utils.cooldown_handlers import _set_cooldown_deployments
|
||||
from litellm.router_utils.router_callbacks.track_deployment_metrics import (
|
||||
increment_deployment_failures_for_current_minute,
|
||||
)
|
||||
|
||||
_exceptions: dict = exceptions_by_model_id or {}
|
||||
|
||||
try:
|
||||
if llm_router is None or not llm_router.enable_health_check_routing:
|
||||
return
|
||||
|
||||
# When health_check_ignore_transient_errors is set, treat 429/408
|
||||
# endpoints as healthy so they are not filtered from routing.
|
||||
_effective_unhealthy = unhealthy_endpoints
|
||||
if llm_router.health_check_ignore_transient_errors:
|
||||
_effective_unhealthy = [
|
||||
ep
|
||||
for ep in unhealthy_endpoints
|
||||
if _get_endpoint_exception_status(ep, _exceptions) not in (429, 408)
|
||||
]
|
||||
|
||||
states = build_deployment_health_states(
|
||||
healthy_endpoints=healthy_endpoints,
|
||||
unhealthy_endpoints=unhealthy_endpoints,
|
||||
unhealthy_endpoints=_effective_unhealthy,
|
||||
)
|
||||
if states:
|
||||
llm_router.health_state_cache.set_deployment_health_states(states)
|
||||
|
|
@ -2273,6 +2304,37 @@ def _write_health_state_to_router_cache(
|
|||
sum(1 for s in states.values() if s.get("is_healthy")),
|
||||
sum(1 for s in states.values() if not s.get("is_healthy")),
|
||||
)
|
||||
|
||||
for endpoint in unhealthy_endpoints:
|
||||
model_id = endpoint.get("model_id")
|
||||
if not model_id:
|
||||
continue
|
||||
|
||||
original_exception = _exceptions.get(model_id)
|
||||
if original_exception is None:
|
||||
continue
|
||||
|
||||
exception_status = getattr(original_exception, "status_code", 500)
|
||||
|
||||
if llm_router.health_check_ignore_transient_errors and exception_status in (
|
||||
429,
|
||||
408,
|
||||
):
|
||||
continue
|
||||
|
||||
increment_deployment_failures_for_current_minute(
|
||||
litellm_router_instance=llm_router,
|
||||
deployment_id=model_id,
|
||||
)
|
||||
|
||||
_set_cooldown_deployments(
|
||||
litellm_router_instance=llm_router,
|
||||
original_exception=original_exception,
|
||||
exception_status=exception_status,
|
||||
deployment=model_id,
|
||||
time_to_cooldown=llm_router.cooldown_time,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.warning(
|
||||
"Failed to write health state to router cache: %s", str(e)
|
||||
|
|
@ -2384,6 +2446,7 @@ async def _run_background_health_check():
|
|||
(
|
||||
healthy_endpoints,
|
||||
unhealthy_endpoints,
|
||||
_exceptions_by_model_id,
|
||||
) = await shared_health_manager.perform_shared_health_check(
|
||||
model_list=_llm_model_list,
|
||||
details=details_bool,
|
||||
|
|
@ -2397,6 +2460,7 @@ async def _run_background_health_check():
|
|||
(
|
||||
healthy_endpoints,
|
||||
unhealthy_endpoints,
|
||||
_exceptions_by_model_id,
|
||||
) = await _run_direct_health_check_with_instrumentation(
|
||||
_llm_model_list,
|
||||
health_check_details,
|
||||
|
|
@ -2407,6 +2471,7 @@ async def _run_background_health_check():
|
|||
(
|
||||
healthy_endpoints,
|
||||
unhealthy_endpoints,
|
||||
_exceptions_by_model_id,
|
||||
) = await _run_direct_health_check_with_instrumentation(
|
||||
_llm_model_list,
|
||||
health_check_details,
|
||||
|
|
@ -2449,7 +2514,9 @@ async def _run_background_health_check():
|
|||
)
|
||||
|
||||
# Write health state to router cache for health-check-driven routing
|
||||
_write_health_state_to_router_cache(healthy_endpoints, unhealthy_endpoints)
|
||||
_write_health_state_to_router_cache(
|
||||
healthy_endpoints, unhealthy_endpoints, _exceptions_by_model_id
|
||||
)
|
||||
|
||||
await asyncio.sleep(health_check_interval)
|
||||
|
||||
|
|
@ -3245,6 +3312,7 @@ class ProxyConfig:
|
|||
general_settings = {}
|
||||
_enable_hc_routing = False
|
||||
_hc_staleness = None
|
||||
_hc_ignore_transient = False
|
||||
if general_settings:
|
||||
### LOAD KEY MANAGEMENT SETTINGS FIRST (needed for custom secret manager) ###
|
||||
key_management_settings = general_settings.get(
|
||||
|
|
@ -3437,6 +3505,9 @@ class ProxyConfig:
|
|||
_hc_staleness = general_settings.get(
|
||||
"health_check_staleness_threshold", None
|
||||
)
|
||||
_hc_ignore_transient = general_settings.get(
|
||||
"health_check_ignore_transient_errors", False
|
||||
)
|
||||
verbose_proxy_logger.info(
|
||||
"background_health_check_config enabled=%s shared=%s interval_seconds=%s max_concurrency=%s details=%s health_check_routing=%s",
|
||||
use_background_health_checks,
|
||||
|
|
@ -3479,6 +3550,8 @@ class ProxyConfig:
|
|||
router_params["enable_health_check_routing"] = True
|
||||
if _hc_staleness is not None:
|
||||
router_params["health_check_staleness_threshold"] = _hc_staleness
|
||||
if _hc_ignore_transient:
|
||||
router_params["health_check_ignore_transient_errors"] = True
|
||||
## MODEL LIST
|
||||
model_list = config.get("model_list", None)
|
||||
if model_list:
|
||||
|
|
@ -5217,10 +5290,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)
|
||||
|
|
@ -5602,9 +5675,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
|
||||
|
|
@ -6503,10 +6576,18 @@ class ProxyStartupEvent:
|
|||
KeyRotationManager,
|
||||
)
|
||||
|
||||
# Get prisma_client from global scope
|
||||
# Get prisma_client and proxy_logging_obj from global scope
|
||||
global prisma_client
|
||||
global proxy_logging_obj
|
||||
if prisma_client is not None:
|
||||
key_rotation_manager = KeyRotationManager(prisma_client)
|
||||
# Reuse the PodLockManager from db_spend_update_writer
|
||||
pod_lock_manager = (
|
||||
proxy_logging_obj.db_spend_update_writer.pod_lock_manager
|
||||
)
|
||||
key_rotation_manager = KeyRotationManager(
|
||||
prisma_client,
|
||||
pod_lock_manager=pod_lock_manager,
|
||||
)
|
||||
verbose_proxy_logger.debug(
|
||||
f"Key rotation background job scheduled every {LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS} seconds (LITELLM_KEY_ROTATION_ENABLED=true)"
|
||||
)
|
||||
|
|
@ -12624,9 +12705,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
|
||||
|
|
|
|||
|
|
@ -310,6 +310,7 @@ class Router:
|
|||
ignore_invalid_deployments: bool = False,
|
||||
enable_health_check_routing: bool = False,
|
||||
health_check_staleness_threshold: Optional[int] = None,
|
||||
health_check_ignore_transient_errors: bool = False,
|
||||
) -> None:
|
||||
"""
|
||||
Initialize the Router class with the given parameters for caching, reliability, and routing strategy.
|
||||
|
|
@ -408,9 +409,9 @@ class Router:
|
|||
) # names of models under litellm_params. ex. azure/chatgpt-v-2
|
||||
self.deployment_latency_map = {}
|
||||
### CACHING ###
|
||||
cache_type: Literal[
|
||||
"local", "redis", "redis-semantic", "s3", "disk"
|
||||
] = "local" # default to an in-memory cache
|
||||
cache_type: Literal["local", "redis", "redis-semantic", "s3", "disk"] = (
|
||||
"local" # default to an in-memory cache
|
||||
)
|
||||
redis_cache = None
|
||||
cache_config: Dict[str, Any] = {}
|
||||
|
||||
|
|
@ -458,9 +459,9 @@ class Router:
|
|||
self.default_max_parallel_requests = default_max_parallel_requests
|
||||
self.provider_default_deployment_ids: List[str] = []
|
||||
self.pattern_router = PatternMatchRouter()
|
||||
self.team_pattern_routers: Dict[
|
||||
str, PatternMatchRouter
|
||||
] = {} # {"TEAM_ID": PatternMatchRouter}
|
||||
self.team_pattern_routers: Dict[str, PatternMatchRouter] = (
|
||||
{}
|
||||
) # {"TEAM_ID": PatternMatchRouter}
|
||||
self.auto_routers: Dict[str, "AutoRouter"] = {}
|
||||
self.complexity_routers: Dict[str, "ComplexityRouter"] = {}
|
||||
|
||||
|
|
@ -501,6 +502,7 @@ class Router:
|
|||
)
|
||||
self.disable_cooldowns = disable_cooldowns
|
||||
self.enable_health_check_routing = enable_health_check_routing
|
||||
self.health_check_ignore_transient_errors = health_check_ignore_transient_errors
|
||||
_staleness = health_check_staleness_threshold or (
|
||||
DEFAULT_HEALTH_CHECK_INTERVAL * DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER
|
||||
)
|
||||
|
|
@ -655,12 +657,12 @@ class Router:
|
|||
)
|
||||
)
|
||||
|
||||
self.model_group_retry_policy: Optional[
|
||||
Dict[str, RetryPolicy]
|
||||
] = model_group_retry_policy
|
||||
self.model_group_affinity_config: Optional[
|
||||
Dict[str, List[str]]
|
||||
] = model_group_affinity_config
|
||||
self.model_group_retry_policy: Optional[Dict[str, RetryPolicy]] = (
|
||||
model_group_retry_policy
|
||||
)
|
||||
self.model_group_affinity_config: Optional[Dict[str, List[str]]] = (
|
||||
model_group_affinity_config
|
||||
)
|
||||
|
||||
self.allowed_fails_policy: Optional[AllowedFailsPolicy] = None
|
||||
if allowed_fails_policy is not None:
|
||||
|
|
@ -2066,7 +2068,10 @@ class Router:
|
|||
|
||||
async def _acompletion( # noqa: PLR0915
|
||||
self, model: str, messages: List[Dict[str, str]], **kwargs
|
||||
) -> Union[ModelResponse, CustomStreamWrapper,]:
|
||||
) -> Union[
|
||||
ModelResponse,
|
||||
CustomStreamWrapper,
|
||||
]:
|
||||
"""
|
||||
- Get an available deployment
|
||||
- call it with a semaphore over the call
|
||||
|
|
@ -4300,9 +4305,9 @@ class Router:
|
|||
healthy_deployments=healthy_deployments, responses=responses
|
||||
)
|
||||
returned_response = cast(OpenAIFileObject, responses[0])
|
||||
returned_response._hidden_params[
|
||||
"model_file_id_mapping"
|
||||
] = model_file_id_mapping
|
||||
returned_response._hidden_params["model_file_id_mapping"] = (
|
||||
model_file_id_mapping
|
||||
)
|
||||
return returned_response
|
||||
except Exception as e:
|
||||
verbose_router_logger.exception(
|
||||
|
|
@ -5310,11 +5315,16 @@ class Router:
|
|||
e,
|
||||
(litellm.ContextWindowExceededError, litellm.ContentPolicyViolationError),
|
||||
)
|
||||
all_deployments = self._get_all_deployments(model_name=original_model_group)
|
||||
_request_team_id: Optional[str] = (
|
||||
kwargs.get("metadata", {}) or {}
|
||||
).get("user_api_key_team_id")
|
||||
all_deployments = self._get_all_deployments(
|
||||
model_name=original_model_group, team_id=_request_team_id
|
||||
)
|
||||
_order_set: set = {
|
||||
d.get("litellm_params", {}).get("order")
|
||||
litellm.utils._get_deployment_order(d)
|
||||
for d in all_deployments
|
||||
if d.get("litellm_params", {}).get("order") is not None
|
||||
if litellm.utils._get_deployment_order(d) is not None
|
||||
}
|
||||
order_values: list = sorted(_order_set)
|
||||
if len(order_values) > 1 and not _skip_order_fallback:
|
||||
|
|
@ -5387,11 +5397,11 @@ class Router:
|
|||
|
||||
if isinstance(e, litellm.ContextWindowExceededError):
|
||||
if context_window_fallbacks is not None:
|
||||
context_window_fallback_model_group: Optional[
|
||||
List[str]
|
||||
] = self._get_fallback_model_group_from_fallbacks(
|
||||
fallbacks=context_window_fallbacks,
|
||||
model_group=model_group,
|
||||
context_window_fallback_model_group: Optional[List[str]] = (
|
||||
self._get_fallback_model_group_from_fallbacks(
|
||||
fallbacks=context_window_fallbacks,
|
||||
model_group=model_group,
|
||||
)
|
||||
)
|
||||
if context_window_fallback_model_group is None:
|
||||
raise original_exception
|
||||
|
|
@ -5423,11 +5433,11 @@ class Router:
|
|||
e.message += "\n{}".format(error_message)
|
||||
elif isinstance(e, litellm.ContentPolicyViolationError):
|
||||
if content_policy_fallbacks is not None:
|
||||
content_policy_fallback_model_group: Optional[
|
||||
List[str]
|
||||
] = self._get_fallback_model_group_from_fallbacks(
|
||||
fallbacks=content_policy_fallbacks,
|
||||
model_group=model_group,
|
||||
content_policy_fallback_model_group: Optional[List[str]] = (
|
||||
self._get_fallback_model_group_from_fallbacks(
|
||||
fallbacks=content_policy_fallbacks,
|
||||
model_group=model_group,
|
||||
)
|
||||
)
|
||||
if content_policy_fallback_model_group is None:
|
||||
raise original_exception
|
||||
|
|
@ -5649,9 +5659,9 @@ class Router:
|
|||
)
|
||||
## ADD RETRY TRACKING TO METADATA - used for spend logs retry tracking
|
||||
_metadata["attempted_retries"] = 0
|
||||
_metadata[
|
||||
"max_retries"
|
||||
] = num_retries # Updated after overrides in exception handler
|
||||
_metadata["max_retries"] = (
|
||||
num_retries # Updated after overrides in exception handler
|
||||
)
|
||||
try:
|
||||
self._handle_mock_testing_rate_limit_error(
|
||||
model_group=model_group, kwargs=kwargs
|
||||
|
|
@ -6770,26 +6780,26 @@ class Router:
|
|||
"""
|
||||
from litellm.router_strategy.auto_router.auto_router import AutoRouter
|
||||
|
||||
auto_router_config_path: Optional[
|
||||
str
|
||||
] = deployment.litellm_params.auto_router_config_path
|
||||
auto_router_config_path: Optional[str] = (
|
||||
deployment.litellm_params.auto_router_config_path
|
||||
)
|
||||
auto_router_config: Optional[str] = deployment.litellm_params.auto_router_config
|
||||
if auto_router_config_path is None and auto_router_config is None:
|
||||
raise ValueError(
|
||||
"auto_router_config_path or auto_router_config is required for auto-router deployments. Please set it in the litellm_params"
|
||||
)
|
||||
|
||||
default_model: Optional[
|
||||
str
|
||||
] = deployment.litellm_params.auto_router_default_model
|
||||
default_model: Optional[str] = (
|
||||
deployment.litellm_params.auto_router_default_model
|
||||
)
|
||||
if default_model is None:
|
||||
raise ValueError(
|
||||
"auto_router_default_model is required for auto-router deployments. Please set it in the litellm_params"
|
||||
)
|
||||
|
||||
embedding_model: Optional[
|
||||
str
|
||||
] = deployment.litellm_params.auto_router_embedding_model
|
||||
embedding_model: Optional[str] = (
|
||||
deployment.litellm_params.auto_router_embedding_model
|
||||
)
|
||||
if embedding_model is None:
|
||||
raise ValueError(
|
||||
"auto_router_embedding_model is required for auto-router deployments. Please set it in the litellm_params"
|
||||
|
|
@ -6832,13 +6842,13 @@ class Router:
|
|||
ComplexityRouter,
|
||||
)
|
||||
|
||||
complexity_router_config: Optional[
|
||||
dict
|
||||
] = deployment.litellm_params.complexity_router_config
|
||||
complexity_router_config: Optional[dict] = (
|
||||
deployment.litellm_params.complexity_router_config
|
||||
)
|
||||
|
||||
default_model: Optional[
|
||||
str
|
||||
] = deployment.litellm_params.complexity_router_default_model
|
||||
default_model: Optional[str] = (
|
||||
deployment.litellm_params.complexity_router_default_model
|
||||
)
|
||||
|
||||
# If no default model specified, try to get from config tiers
|
||||
if default_model is None and complexity_router_config:
|
||||
|
|
@ -7497,9 +7507,9 @@ class Router:
|
|||
|
||||
# Add custom_llm_provider
|
||||
if deployment.litellm_params.custom_llm_provider:
|
||||
credentials[
|
||||
"custom_llm_provider"
|
||||
] = deployment.litellm_params.custom_llm_provider
|
||||
credentials["custom_llm_provider"] = (
|
||||
deployment.litellm_params.custom_llm_provider
|
||||
)
|
||||
elif "/" in deployment.litellm_params.model:
|
||||
# Extract provider from "provider/model" format
|
||||
credentials["custom_llm_provider"] = deployment.litellm_params.model.split(
|
||||
|
|
@ -9070,7 +9080,9 @@ class Router:
|
|||
|
||||
## get healthy deployments
|
||||
### get all deployments
|
||||
healthy_deployments = self._get_all_deployments(model_name=model)
|
||||
healthy_deployments = self._get_all_deployments(
|
||||
model_name=model, team_id=request_team_id
|
||||
)
|
||||
|
||||
if len(healthy_deployments) == 0:
|
||||
# check if the user sent in a deployment name instead
|
||||
|
|
@ -9091,7 +9103,9 @@ class Router:
|
|||
)
|
||||
# Re-assign model to the fallback and try to get deployments again
|
||||
model = fallback_model
|
||||
healthy_deployments = self._get_all_deployments(model_name=model)
|
||||
healthy_deployments = self._get_all_deployments(
|
||||
model_name=model, team_id=request_team_id
|
||||
)
|
||||
|
||||
# If still no deployments after checking for fallbacks, raise an error
|
||||
if len(healthy_deployments) == 0:
|
||||
|
|
@ -9181,10 +9195,23 @@ class Router:
|
|||
)
|
||||
if verbose_router_logger.isEnabledFor(logging.DEBUG):
|
||||
verbose_router_logger.debug(f"cooldown deployments: {cooldown_deployments}")
|
||||
_pre_cooldown_deployments = healthy_deployments
|
||||
healthy_deployments = self._filter_cooldown_deployments(
|
||||
healthy_deployments=healthy_deployments,
|
||||
cooldown_deployments=cooldown_deployments,
|
||||
)
|
||||
# Safety net: only bypass cooldown filter when health-check routing is
|
||||
# driving cooldown (i.e. allowed_fails_policy is set). Without a policy,
|
||||
# cooldowns are from real request failures and must not be bypassed.
|
||||
if (
|
||||
not healthy_deployments
|
||||
and self.enable_health_check_routing
|
||||
and self.allowed_fails_policy is not None
|
||||
):
|
||||
verbose_router_logger.warning(
|
||||
"All deployments in cooldown via health-check routing, bypassing cooldown filter"
|
||||
)
|
||||
healthy_deployments = _pre_cooldown_deployments
|
||||
|
||||
healthy_deployments = await self.async_callback_filter_deployments(
|
||||
model=model,
|
||||
|
|
@ -9617,10 +9644,20 @@ class Router:
|
|||
cooldown_deployments = _get_cooldown_deployments(
|
||||
litellm_router_instance=self, parent_otel_span=parent_otel_span
|
||||
)
|
||||
_pre_cooldown_deployments = healthy_deployments
|
||||
healthy_deployments = self._filter_cooldown_deployments(
|
||||
healthy_deployments=healthy_deployments,
|
||||
cooldown_deployments=cooldown_deployments,
|
||||
)
|
||||
if (
|
||||
not healthy_deployments
|
||||
and self.enable_health_check_routing
|
||||
and self.allowed_fails_policy is not None
|
||||
):
|
||||
verbose_router_logger.warning(
|
||||
"All deployments in cooldown via health-check routing, bypassing cooldown filter"
|
||||
)
|
||||
healthy_deployments = _pre_cooldown_deployments
|
||||
|
||||
# filter pre-call checks
|
||||
if self.enable_pre_call_checks and messages is not None:
|
||||
|
|
@ -9922,6 +9959,12 @@ class Router:
|
|||
if not self.enable_health_check_routing:
|
||||
return healthy_deployments
|
||||
|
||||
# When allowed_fails_policy is set, cooldown is the sole routing exclusion
|
||||
# mechanism -- skip the binary health check filter so the policy threshold
|
||||
# is respected before any deployment is excluded.
|
||||
if self.allowed_fails_policy is not None:
|
||||
return healthy_deployments
|
||||
|
||||
unhealthy_ids = (
|
||||
await self.health_state_cache.async_get_unhealthy_deployment_ids(
|
||||
parent_otel_span=parent_otel_span
|
||||
|
|
@ -9951,6 +9994,9 @@ class Router:
|
|||
if not self.enable_health_check_routing:
|
||||
return healthy_deployments
|
||||
|
||||
if self.allowed_fails_policy is not None:
|
||||
return healthy_deployments
|
||||
|
||||
unhealthy_ids = self.health_state_cache.get_unhealthy_deployment_ids(
|
||||
parent_otel_span=parent_otel_span
|
||||
)
|
||||
|
|
|
|||
|
|
@ -18,18 +18,24 @@ class PipelineStep(BaseModel):
|
|||
"""
|
||||
A single step in a guardrail pipeline.
|
||||
|
||||
Each step runs a guardrail and takes an action based on pass/fail.
|
||||
Each step runs a guardrail and takes an action based on pass, policy fail,
|
||||
or technical/API error (see pipeline executor outcome types).
|
||||
"""
|
||||
|
||||
guardrail: str = Field(description="Name of the guardrail to run.")
|
||||
on_fail: str = Field(
|
||||
default="block",
|
||||
description="Action when guardrail rejects: next | block | allow | modify_response",
|
||||
description="Action when guardrail rejects content (policy intervention): next | block | allow | modify_response",
|
||||
)
|
||||
on_pass: str = Field(
|
||||
default="allow",
|
||||
description="Action when guardrail passes: next | block | allow | modify_response",
|
||||
)
|
||||
on_error: Optional[str] = Field(
|
||||
default=None,
|
||||
description="Action when the guardrail raises a technical error (timeouts, "
|
||||
"unreachable provider, non-intervention HTTP errors). If omitted, uses on_fail.",
|
||||
)
|
||||
pass_data: bool = Field(
|
||||
default=False,
|
||||
description="Forward modified request data (e.g., PII-masked) to next step.",
|
||||
|
|
@ -41,9 +47,11 @@ class PipelineStep(BaseModel):
|
|||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
@field_validator("on_fail", "on_pass")
|
||||
@field_validator("on_fail", "on_pass", "on_error")
|
||||
@classmethod
|
||||
def validate_action(cls, v: str) -> str:
|
||||
def validate_action(cls, v: Optional[str]) -> Optional[str]:
|
||||
if v is None:
|
||||
return None
|
||||
if v not in VALID_PIPELINE_ACTIONS:
|
||||
raise ValueError(
|
||||
f"Invalid action '{v}'. Must be one of: {sorted(VALID_PIPELINE_ACTIONS)}"
|
||||
|
|
|
|||
|
|
@ -4875,6 +4875,19 @@ def calculate_max_parallel_requests(
|
|||
return None
|
||||
|
||||
|
||||
def _get_deployment_order(deployment: Union[Dict, Any]) -> Optional[int]:
|
||||
"""
|
||||
Returns the routing order for a deployment.
|
||||
|
||||
Checks litellm_params first (static config), then model_info (dynamic/team
|
||||
models added via API where order lives in model_info, not litellm_params).
|
||||
"""
|
||||
order = deployment.get("litellm_params", {}).get("order")
|
||||
if order is None:
|
||||
order = deployment.get("model_info", {}).get("order")
|
||||
return order
|
||||
|
||||
|
||||
def _get_order_filtered_deployments(
|
||||
healthy_deployments: List[Dict], target_order: Optional[int] = None
|
||||
) -> List:
|
||||
|
|
@ -4882,7 +4895,7 @@ def _get_order_filtered_deployments(
|
|||
filtered = [
|
||||
d
|
||||
for d in healthy_deployments
|
||||
if d["litellm_params"].get("order") == target_order
|
||||
if _get_deployment_order(d) == target_order
|
||||
]
|
||||
if filtered:
|
||||
return filtered
|
||||
|
|
@ -4890,20 +4903,19 @@ def _get_order_filtered_deployments(
|
|||
return healthy_deployments
|
||||
|
||||
# Default: pick min order group
|
||||
min_order = min(
|
||||
(
|
||||
deployment["litellm_params"]["order"]
|
||||
for deployment in healthy_deployments
|
||||
if "order" in deployment["litellm_params"]
|
||||
),
|
||||
default=None,
|
||||
)
|
||||
_valid_orders: List[int] = [
|
||||
o
|
||||
for deployment in healthy_deployments
|
||||
for o in [_get_deployment_order(deployment)]
|
||||
if o is not None
|
||||
]
|
||||
min_order: Optional[int] = min(_valid_orders) if _valid_orders else None
|
||||
|
||||
if min_order is not None:
|
||||
filtered_deployments = [
|
||||
deployment
|
||||
for deployment in healthy_deployments
|
||||
if deployment["litellm_params"].get("order") == min_order
|
||||
if _get_deployment_order(deployment) == min_order
|
||||
]
|
||||
|
||||
return filtered_deployments
|
||||
|
|
|
|||
|
|
@ -475,13 +475,13 @@ async def test_perform_health_check_filters_by_model_id():
|
|||
captured_list.append(m_list)
|
||||
return [
|
||||
{"model": "gpt-4", "api_key": m_list[0]["litellm_params"]["api_key"]}
|
||||
], []
|
||||
], [], {}
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.health_check._perform_health_check",
|
||||
side_effect=mock_perform_health_check,
|
||||
):
|
||||
healthy_endpoints, unhealthy_endpoints = await perform_health_check(
|
||||
healthy_endpoints, unhealthy_endpoints, _ = await perform_health_check(
|
||||
model_list=model_list, model_id="deployment-id-2", details=True
|
||||
)
|
||||
|
||||
|
|
@ -521,7 +521,7 @@ async def test_perform_health_check_with_health_check_model():
|
|||
return {"status": "healthy"}
|
||||
|
||||
with patch("litellm.ahealth_check", side_effect=mock_health_check):
|
||||
healthy_endpoints, unhealthy_endpoints = await _perform_health_check(model_list)
|
||||
healthy_endpoints, unhealthy_endpoints, _ = await _perform_health_check(model_list)
|
||||
print("health check calls: ", health_check_calls)
|
||||
|
||||
# Verify the health check used the override model
|
||||
|
|
@ -556,7 +556,7 @@ async def test_health_check_bad_model():
|
|||
},
|
||||
]
|
||||
details = None
|
||||
healthy_endpoints, unhealthy_endpoints = await _perform_health_check(
|
||||
healthy_endpoints, unhealthy_endpoints, _ = await _perform_health_check(
|
||||
model_list, details
|
||||
)
|
||||
print(f"healthy_endpoints: {healthy_endpoints}")
|
||||
|
|
@ -574,7 +574,7 @@ async def test_health_check_bad_model():
|
|||
"litellm.ahealth_check", side_effect=mock_health_check
|
||||
) as mock_health_check:
|
||||
start_time = time.time()
|
||||
healthy_endpoints, unhealthy_endpoints = await _perform_health_check(model_list)
|
||||
healthy_endpoints, unhealthy_endpoints, _ = await _perform_health_check(model_list)
|
||||
end_time = time.time()
|
||||
print("health check calls: ", health_check_calls)
|
||||
assert len(healthy_endpoints) == 0
|
||||
|
|
@ -667,7 +667,7 @@ async def test_timeout_does_not_cancel_other_health_checks():
|
|||
return {"status": "healthy"}
|
||||
|
||||
with patch("litellm.ahealth_check", side_effect=mock_health_check):
|
||||
healthy_endpoints, unhealthy_endpoints = await _perform_health_check(
|
||||
healthy_endpoints, unhealthy_endpoints, _ = await _perform_health_check(
|
||||
model_list, max_concurrency=1
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -2420,7 +2420,7 @@ async def test_run_background_health_check_reflects_llm_model_list(monkeypatch):
|
|||
|
||||
async def fake_perform_health_check(model_list, details, max_concurrency=None):
|
||||
called_model_lists.append(copy.deepcopy(model_list))
|
||||
return (["healthy"], ["unhealthy"])
|
||||
return (["healthy"], ["unhealthy"], {})
|
||||
|
||||
monkeypatch.setattr(proxy_server, "health_check_interval", 1)
|
||||
monkeypatch.setattr(proxy_server, "health_check_details", None)
|
||||
|
|
@ -2471,7 +2471,7 @@ async def test_background_health_check_skip_disabled_models(monkeypatch):
|
|||
|
||||
async def fake_perform_health_check(model_list, details, max_concurrency=None):
|
||||
called_model_lists.append(copy.deepcopy(model_list))
|
||||
return (["healthy"], [])
|
||||
return (["healthy"], [], {})
|
||||
|
||||
monkeypatch.setattr(proxy_server, "health_check_interval", 1)
|
||||
monkeypatch.setattr(proxy_server, "health_check_details", None)
|
||||
|
|
|
|||
559
tests/test_litellm/proxy/common_utils/test_key_rotation_e2e.py
Normal file
559
tests/test_litellm/proxy/common_utils/test_key_rotation_e2e.py
Normal file
|
|
@ -0,0 +1,559 @@
|
|||
"""
|
||||
End-to-end tests for key rotation feature.
|
||||
|
||||
Covers the critical gaps:
|
||||
1. Multi-pod simulation: two KeyRotationManagers sharing one PodLockManager
|
||||
2. Error resilience: partial failures, regenerate_key_fn failures, hook failures
|
||||
3. Full process_rotations flow with actual key finding + rotation + lock
|
||||
4. Initialization wiring: PodLockManager is correctly passed
|
||||
5. Multiple keys: some succeed, some fail, all are attempted
|
||||
6. Rotation count increments correctly over multiple rotations
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../../.."))
|
||||
|
||||
from litellm.proxy._types import (
|
||||
GenerateKeyResponse,
|
||||
LiteLLM_VerificationToken,
|
||||
)
|
||||
from litellm.proxy.common_utils.key_rotation_manager import KeyRotationManager
|
||||
|
||||
|
||||
class TestMultiPodKeyRotation:
|
||||
"""
|
||||
Simulate two pods sharing one Redis lock to verify only one pod
|
||||
runs key rotation at a time.
|
||||
"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_two_pods_only_one_rotates(self):
|
||||
"""
|
||||
Two KeyRotationManagers with separate pod_lock_managers but
|
||||
the same Redis backend. Only the first to acquire the lock
|
||||
should rotate; the second should skip.
|
||||
"""
|
||||
mock_prisma = AsyncMock()
|
||||
|
||||
# Shared state to simulate Redis SET NX behavior
|
||||
redis_lock = {"holder": None}
|
||||
|
||||
async def make_acquire_lock(pod_id):
|
||||
async def acquire(cronjob_id, **kwargs):
|
||||
if redis_lock["holder"] is None:
|
||||
redis_lock["holder"] = pod_id
|
||||
return True
|
||||
return redis_lock["holder"] == pod_id
|
||||
|
||||
return acquire
|
||||
|
||||
async def make_release_lock(pod_id):
|
||||
async def release(cronjob_id):
|
||||
if redis_lock["holder"] == pod_id:
|
||||
redis_lock["holder"] = None
|
||||
|
||||
return release
|
||||
|
||||
# Pod A
|
||||
pod_a_lock_mgr = MagicMock()
|
||||
pod_a_lock_mgr.redis_cache = MagicMock()
|
||||
pod_a_lock_mgr.acquire_lock = AsyncMock(
|
||||
side_effect=await make_acquire_lock("pod-a")
|
||||
)
|
||||
pod_a_lock_mgr.release_lock = AsyncMock(
|
||||
side_effect=await make_release_lock("pod-a")
|
||||
)
|
||||
|
||||
# Pod B
|
||||
pod_b_lock_mgr = MagicMock()
|
||||
pod_b_lock_mgr.redis_cache = MagicMock()
|
||||
pod_b_lock_mgr.acquire_lock = AsyncMock(
|
||||
side_effect=await make_acquire_lock("pod-b")
|
||||
)
|
||||
pod_b_lock_mgr.release_lock = AsyncMock(
|
||||
side_effect=await make_release_lock("pod-b")
|
||||
)
|
||||
|
||||
manager_a = KeyRotationManager(mock_prisma, pod_lock_manager=pod_a_lock_mgr)
|
||||
manager_b = KeyRotationManager(mock_prisma, pod_lock_manager=pod_b_lock_mgr)
|
||||
|
||||
# Both share the same mock methods for rotation logic
|
||||
for mgr in [manager_a, manager_b]:
|
||||
mgr._cleanup_expired_deprecated_keys = AsyncMock()
|
||||
mgr._find_keys_needing_rotation = AsyncMock(return_value=[])
|
||||
|
||||
# Pod A acquires lock first
|
||||
await manager_a.process_rotations()
|
||||
# Pod A should have run rotation
|
||||
manager_a._cleanup_expired_deprecated_keys.assert_called_once()
|
||||
manager_a._find_keys_needing_rotation.assert_called_once()
|
||||
|
||||
# Lock is released after pod A finishes, so pod B can now acquire
|
||||
# But let's simulate pod B trying WHILE pod A holds the lock
|
||||
# Reset the lock state to simulate concurrent access
|
||||
redis_lock["holder"] = "pod-a" # Pod A holds the lock
|
||||
|
||||
await manager_b.process_rotations()
|
||||
# Pod B should NOT have run rotation (lock held by pod-a)
|
||||
manager_b._cleanup_expired_deprecated_keys.assert_not_called()
|
||||
manager_b._find_keys_needing_rotation.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_second_pod_runs_after_first_releases(self):
|
||||
"""
|
||||
After the first pod releases the lock, the second pod should
|
||||
be able to acquire and run rotation.
|
||||
"""
|
||||
mock_prisma = AsyncMock()
|
||||
|
||||
call_order = []
|
||||
|
||||
# Pod A - always gets the lock
|
||||
pod_a_lock = MagicMock()
|
||||
pod_a_lock.redis_cache = MagicMock()
|
||||
pod_a_lock.acquire_lock = AsyncMock(return_value=True)
|
||||
pod_a_lock.release_lock = AsyncMock()
|
||||
|
||||
# Pod B - also gets the lock (simulating after A releases)
|
||||
pod_b_lock = MagicMock()
|
||||
pod_b_lock.redis_cache = MagicMock()
|
||||
pod_b_lock.acquire_lock = AsyncMock(return_value=True)
|
||||
pod_b_lock.release_lock = AsyncMock()
|
||||
|
||||
manager_a = KeyRotationManager(mock_prisma, pod_lock_manager=pod_a_lock)
|
||||
manager_b = KeyRotationManager(mock_prisma, pod_lock_manager=pod_b_lock)
|
||||
|
||||
async def cleanup_a():
|
||||
call_order.append("a_cleanup")
|
||||
|
||||
async def cleanup_b():
|
||||
call_order.append("b_cleanup")
|
||||
|
||||
manager_a._cleanup_expired_deprecated_keys = AsyncMock(side_effect=cleanup_a)
|
||||
manager_a._find_keys_needing_rotation = AsyncMock(return_value=[])
|
||||
manager_b._cleanup_expired_deprecated_keys = AsyncMock(side_effect=cleanup_b)
|
||||
manager_b._find_keys_needing_rotation = AsyncMock(return_value=[])
|
||||
|
||||
# Run sequentially: A then B
|
||||
await manager_a.process_rotations()
|
||||
await manager_b.process_rotations()
|
||||
|
||||
# Both should have run
|
||||
assert call_order == ["a_cleanup", "b_cleanup"]
|
||||
pod_a_lock.release_lock.assert_called_once()
|
||||
pod_b_lock.release_lock.assert_called_once()
|
||||
|
||||
|
||||
class TestKeyRotationErrorResilience:
|
||||
"""
|
||||
Tests that key rotation handles errors gracefully:
|
||||
- regenerate_key_fn failure for one key doesn't block others
|
||||
- Hook failure doesn't crash the process
|
||||
- Database update failure is handled
|
||||
"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_one_key_fails_others_still_rotate(self):
|
||||
"""
|
||||
If rotation fails for one key, the remaining keys should still
|
||||
be attempted. No key should be silently skipped.
|
||||
"""
|
||||
mock_prisma = AsyncMock()
|
||||
manager = KeyRotationManager(mock_prisma)
|
||||
|
||||
key1 = LiteLLM_VerificationToken(
|
||||
token="token-1",
|
||||
auto_rotate=True,
|
||||
rotation_interval="30s",
|
||||
key_rotation_at=None,
|
||||
rotation_count=0,
|
||||
key_name="key-1",
|
||||
)
|
||||
key2 = LiteLLM_VerificationToken(
|
||||
token="token-2",
|
||||
auto_rotate=True,
|
||||
rotation_interval="30s",
|
||||
key_rotation_at=None,
|
||||
rotation_count=0,
|
||||
key_name="key-2",
|
||||
)
|
||||
key3 = LiteLLM_VerificationToken(
|
||||
token="token-3",
|
||||
auto_rotate=True,
|
||||
rotation_interval="30s",
|
||||
key_rotation_at=None,
|
||||
rotation_count=0,
|
||||
key_name="key-3",
|
||||
)
|
||||
|
||||
manager._cleanup_expired_deprecated_keys = AsyncMock()
|
||||
manager._find_keys_needing_rotation = AsyncMock(return_value=[key1, key2, key3])
|
||||
|
||||
rotate_calls = []
|
||||
|
||||
async def mock_rotate(key):
|
||||
rotate_calls.append(key.token)
|
||||
if key.token == "token-2":
|
||||
raise Exception("Database connection lost")
|
||||
|
||||
manager._rotate_key = AsyncMock(side_effect=mock_rotate)
|
||||
|
||||
await manager.process_rotations()
|
||||
|
||||
# All 3 keys should have been attempted
|
||||
assert rotate_calls == ["token-1", "token-2", "token-3"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_regenerate_key_fn_failure_is_caught(self):
|
||||
"""
|
||||
If regenerate_key_fn throws, _rotate_key should propagate the error
|
||||
but process_rotations should catch it per-key.
|
||||
"""
|
||||
mock_prisma = AsyncMock()
|
||||
manager = KeyRotationManager(mock_prisma)
|
||||
|
||||
key = LiteLLM_VerificationToken(
|
||||
token="test-token",
|
||||
auto_rotate=True,
|
||||
rotation_interval="30s",
|
||||
key_rotation_at=None,
|
||||
rotation_count=0,
|
||||
key_name="test-key",
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=Exception("regenerate failed: DB timeout"),
|
||||
):
|
||||
# _rotate_key should raise
|
||||
with pytest.raises(Exception, match="regenerate failed"):
|
||||
await manager._rotate_key(key)
|
||||
|
||||
# But process_rotations should catch per-key errors
|
||||
manager._cleanup_expired_deprecated_keys = AsyncMock()
|
||||
manager._find_keys_needing_rotation = AsyncMock(return_value=[key])
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=Exception("regenerate failed: DB timeout"),
|
||||
):
|
||||
# Should NOT raise - error is caught per-key
|
||||
await manager.process_rotations()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hook_failure_does_not_prevent_db_update(self):
|
||||
"""
|
||||
If the rotation hook (async_key_rotated_hook) fails, the database
|
||||
update for rotation_count should still have succeeded (it runs before the hook).
|
||||
"""
|
||||
mock_prisma = AsyncMock()
|
||||
manager = KeyRotationManager(mock_prisma)
|
||||
|
||||
key = LiteLLM_VerificationToken(
|
||||
token="test-token",
|
||||
auto_rotate=True,
|
||||
rotation_interval="30s",
|
||||
key_rotation_at=None,
|
||||
rotation_count=0,
|
||||
)
|
||||
|
||||
mock_response = GenerateKeyResponse(
|
||||
key="new-key", token_id="new-token-id", user_id="test-user"
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_response,
|
||||
):
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.key_rotation_manager.KeyManagementEventHooks.async_key_rotated_hook",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=Exception("Hook failed: secret manager down"),
|
||||
):
|
||||
# This will raise because the hook fails
|
||||
with pytest.raises(Exception, match="Hook failed"):
|
||||
await manager._rotate_key(key)
|
||||
|
||||
# The DB update should have been called BEFORE the hook
|
||||
mock_prisma.db.litellm_verificationtoken.update.assert_called_once()
|
||||
update_data = mock_prisma.db.litellm_verificationtoken.update.call_args[1][
|
||||
"data"
|
||||
]
|
||||
assert update_data["rotation_count"] == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_failure_does_not_prevent_rotation(self):
|
||||
"""
|
||||
If deprecated key cleanup fails, the rotation should still proceed.
|
||||
"""
|
||||
mock_prisma = AsyncMock()
|
||||
mock_pod_lock = MagicMock()
|
||||
mock_pod_lock.redis_cache = MagicMock()
|
||||
mock_pod_lock.acquire_lock = AsyncMock(return_value=True)
|
||||
mock_pod_lock.release_lock = AsyncMock()
|
||||
|
||||
manager = KeyRotationManager(mock_prisma, pod_lock_manager=mock_pod_lock)
|
||||
|
||||
# Cleanup fails
|
||||
manager._cleanup_expired_deprecated_keys = AsyncMock(
|
||||
side_effect=Exception("Deprecated table doesn't exist")
|
||||
)
|
||||
|
||||
# process_rotations catches the exception internally (try/except),
|
||||
# but the lock must still be released in the finally block.
|
||||
await manager.process_rotations()
|
||||
|
||||
# Lock should still be released in finally block
|
||||
mock_pod_lock.release_lock.assert_called_once()
|
||||
|
||||
|
||||
class TestKeyRotationFullFlow:
|
||||
"""
|
||||
Full end-to-end flow tests: find keys -> rotate -> update DB -> release lock
|
||||
"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_full_rotation_flow_with_lock(self):
|
||||
"""
|
||||
Full flow: acquire lock -> cleanup -> find keys -> rotate -> update DB -> release lock
|
||||
"""
|
||||
mock_prisma = AsyncMock()
|
||||
|
||||
# Setup lock manager
|
||||
mock_lock = MagicMock()
|
||||
mock_lock.redis_cache = MagicMock()
|
||||
mock_lock.acquire_lock = AsyncMock(return_value=True)
|
||||
mock_lock.release_lock = AsyncMock()
|
||||
|
||||
manager = KeyRotationManager(mock_prisma, pod_lock_manager=mock_lock)
|
||||
|
||||
key = LiteLLM_VerificationToken(
|
||||
token="old-token-hash",
|
||||
auto_rotate=True,
|
||||
rotation_interval="30s",
|
||||
key_rotation_at=datetime.now(timezone.utc) - timedelta(seconds=60),
|
||||
rotation_count=2,
|
||||
key_name="my-key",
|
||||
key_alias="prod/my-key",
|
||||
)
|
||||
|
||||
mock_response = GenerateKeyResponse(
|
||||
key="sk-new-key-value",
|
||||
token_id="new-token-hash",
|
||||
user_id="system",
|
||||
)
|
||||
|
||||
# Mock cleanup
|
||||
mock_prisma.db.litellm_deprecatedverificationtoken.delete_many.return_value = 1
|
||||
# Mock find keys
|
||||
mock_prisma.db.litellm_verificationtoken.find_many.return_value = [key]
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_response,
|
||||
):
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.key_rotation_manager.KeyManagementEventHooks.async_key_rotated_hook",
|
||||
new_callable=AsyncMock,
|
||||
):
|
||||
await manager.process_rotations()
|
||||
|
||||
# Verify full flow executed:
|
||||
# 1. Lock acquired
|
||||
mock_lock.acquire_lock.assert_called_once()
|
||||
|
||||
# 2. Cleanup ran
|
||||
mock_prisma.db.litellm_deprecatedverificationtoken.delete_many.assert_called_once()
|
||||
|
||||
# 3. Keys were queried
|
||||
mock_prisma.db.litellm_verificationtoken.find_many.assert_called_once()
|
||||
|
||||
# 4. DB was updated with new rotation info
|
||||
mock_prisma.db.litellm_verificationtoken.update.assert_called_once()
|
||||
update_args = mock_prisma.db.litellm_verificationtoken.update.call_args[1]
|
||||
assert update_args["where"]["token"] == "new-token-hash"
|
||||
assert update_args["data"]["rotation_count"] == 3 # was 2, now 3
|
||||
|
||||
# 5. Lock released
|
||||
mock_lock.release_lock.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rotation_count_increments_across_multiple_rotations(self):
|
||||
"""
|
||||
Simulate 3 consecutive rotations and verify rotation_count increments
|
||||
correctly each time: 0 -> 1 -> 2 -> 3
|
||||
"""
|
||||
mock_prisma = AsyncMock()
|
||||
manager = KeyRotationManager(mock_prisma)
|
||||
|
||||
rotation_counts_seen = []
|
||||
|
||||
for expected_count in range(3):
|
||||
key = LiteLLM_VerificationToken(
|
||||
token=f"token-v{expected_count}",
|
||||
auto_rotate=True,
|
||||
rotation_interval="30s",
|
||||
key_rotation_at=None,
|
||||
rotation_count=expected_count,
|
||||
)
|
||||
|
||||
mock_response = GenerateKeyResponse(
|
||||
key=f"sk-new-v{expected_count + 1}",
|
||||
token_id=f"token-v{expected_count + 1}",
|
||||
user_id="system",
|
||||
)
|
||||
|
||||
mock_prisma.db.litellm_verificationtoken.update.reset_mock()
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_response,
|
||||
):
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.key_rotation_manager.KeyManagementEventHooks.async_key_rotated_hook",
|
||||
new_callable=AsyncMock,
|
||||
):
|
||||
await manager._rotate_key(key)
|
||||
|
||||
update_data = mock_prisma.db.litellm_verificationtoken.update.call_args[1][
|
||||
"data"
|
||||
]
|
||||
rotation_counts_seen.append(update_data["rotation_count"])
|
||||
|
||||
assert rotation_counts_seen == [1, 2, 3]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_keys_to_rotate_skips_gracefully(self):
|
||||
"""
|
||||
When no keys need rotation, process should complete without errors.
|
||||
"""
|
||||
mock_prisma = AsyncMock()
|
||||
mock_prisma.db.litellm_deprecatedverificationtoken.delete_many.return_value = 0
|
||||
mock_prisma.db.litellm_verificationtoken.find_many.return_value = []
|
||||
|
||||
mock_lock = MagicMock()
|
||||
mock_lock.redis_cache = MagicMock()
|
||||
mock_lock.acquire_lock = AsyncMock(return_value=True)
|
||||
mock_lock.release_lock = AsyncMock()
|
||||
|
||||
manager = KeyRotationManager(mock_prisma, pod_lock_manager=mock_lock)
|
||||
|
||||
await manager.process_rotations()
|
||||
|
||||
# Verify no rotation was attempted
|
||||
mock_prisma.db.litellm_verificationtoken.update.assert_not_called()
|
||||
# But lock was still properly released
|
||||
mock_lock.release_lock.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_regenerate_response_missing_token_id_skips_db_update(self):
|
||||
"""
|
||||
If regenerate_key_fn returns a response without token_id,
|
||||
the DB update for rotation metadata should be skipped.
|
||||
"""
|
||||
mock_prisma = AsyncMock()
|
||||
manager = KeyRotationManager(mock_prisma)
|
||||
|
||||
key = LiteLLM_VerificationToken(
|
||||
token="old-token",
|
||||
auto_rotate=True,
|
||||
rotation_interval="30s",
|
||||
key_rotation_at=None,
|
||||
rotation_count=0,
|
||||
)
|
||||
|
||||
# Response with no token_id
|
||||
mock_response = GenerateKeyResponse(
|
||||
key="sk-new",
|
||||
token_id=None,
|
||||
user_id="system",
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_response,
|
||||
):
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.key_rotation_manager.KeyManagementEventHooks.async_key_rotated_hook",
|
||||
new_callable=AsyncMock,
|
||||
):
|
||||
await manager._rotate_key(key)
|
||||
|
||||
# DB update should NOT have been called (no token_id)
|
||||
mock_prisma.db.litellm_verificationtoken.update.assert_not_called()
|
||||
|
||||
|
||||
class TestKeyRotationInitialization:
|
||||
"""
|
||||
Tests that the PodLockManager wiring in proxy_server.py is correct.
|
||||
"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_rotation_manager_receives_pod_lock_manager(self):
|
||||
"""
|
||||
Verify KeyRotationManager stores the pod_lock_manager correctly.
|
||||
"""
|
||||
mock_prisma = AsyncMock()
|
||||
mock_lock = MagicMock()
|
||||
mock_lock.redis_cache = MagicMock()
|
||||
|
||||
manager = KeyRotationManager(mock_prisma, pod_lock_manager=mock_lock)
|
||||
|
||||
assert manager.pod_lock_manager is mock_lock
|
||||
assert manager.prisma_client is mock_prisma
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_rotation_manager_default_no_lock(self):
|
||||
"""
|
||||
When no pod_lock_manager is provided, it defaults to None.
|
||||
"""
|
||||
mock_prisma = AsyncMock()
|
||||
manager = KeyRotationManager(mock_prisma)
|
||||
|
||||
assert manager.pod_lock_manager is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lock_pattern_matches_spend_log_cleanup(self):
|
||||
"""
|
||||
Verify the key rotation lock pattern is identical to spend_log_cleanup:
|
||||
- acquire_lock with cronjob_id
|
||||
- release_lock in finally
|
||||
- lock_acquired flag guards release
|
||||
"""
|
||||
mock_prisma = AsyncMock()
|
||||
mock_lock = MagicMock()
|
||||
mock_lock.redis_cache = MagicMock()
|
||||
mock_lock.acquire_lock = AsyncMock(return_value=True)
|
||||
mock_lock.release_lock = AsyncMock()
|
||||
|
||||
manager = KeyRotationManager(mock_prisma, pod_lock_manager=mock_lock)
|
||||
manager._cleanup_expired_deprecated_keys = AsyncMock()
|
||||
manager._find_keys_needing_rotation = AsyncMock(return_value=[])
|
||||
|
||||
await manager.process_rotations()
|
||||
|
||||
# Pattern check: acquire with cronjob_id
|
||||
acquire_call = mock_lock.acquire_lock.call_args
|
||||
assert "cronjob_id" in acquire_call.kwargs or len(acquire_call.args) > 0
|
||||
|
||||
# Pattern check: release with same cronjob_id
|
||||
release_call = mock_lock.release_lock.call_args
|
||||
assert "cronjob_id" in release_call.kwargs or len(release_call.args) > 0
|
||||
|
||||
# Both should use the same job name
|
||||
from litellm.constants import KEY_ROTATION_JOB_NAME
|
||||
|
||||
assert acquire_call.kwargs.get("cronjob_id") == KEY_ROTATION_JOB_NAME
|
||||
assert release_call.kwargs.get("cronjob_id") == KEY_ROTATION_JOB_NAME
|
||||
229
tests/test_litellm/proxy/common_utils/test_key_rotation_lock.py
Normal file
229
tests/test_litellm/proxy/common_utils/test_key_rotation_lock.py
Normal file
|
|
@ -0,0 +1,229 @@
|
|||
"""
|
||||
Test distributed lock behavior for key rotation manager.
|
||||
|
||||
Verifies that PodLockManager is correctly used to prevent concurrent
|
||||
key rotation across multiple pods in a distributed deployment.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../../.."))
|
||||
|
||||
from litellm.proxy._types import LiteLLM_VerificationToken
|
||||
from litellm.proxy.common_utils.key_rotation_manager import KeyRotationManager
|
||||
|
||||
|
||||
class TestKeyRotationLock:
|
||||
"""Test distributed lock behavior in KeyRotationManager."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_rotations_acquires_lock(self):
|
||||
"""
|
||||
When PodLockManager is provided and lock is acquired,
|
||||
rotation logic should run normally.
|
||||
"""
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_pod_lock_manager = MagicMock()
|
||||
mock_pod_lock_manager.redis_cache = MagicMock() # Redis is available
|
||||
mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True)
|
||||
mock_pod_lock_manager.release_lock = AsyncMock()
|
||||
|
||||
manager = KeyRotationManager(
|
||||
mock_prisma_client, pod_lock_manager=mock_pod_lock_manager
|
||||
)
|
||||
|
||||
# Mock _find_keys_needing_rotation to return empty list (no keys to rotate)
|
||||
manager._find_keys_needing_rotation = AsyncMock(return_value=[])
|
||||
manager._cleanup_expired_deprecated_keys = AsyncMock()
|
||||
|
||||
await manager.process_rotations()
|
||||
|
||||
# Verify lock was acquired with custom TTL
|
||||
mock_pod_lock_manager.acquire_lock.assert_called_once()
|
||||
call_kwargs = mock_pod_lock_manager.acquire_lock.call_args
|
||||
assert call_kwargs.kwargs["cronjob_id"] == "litellm_key_rotation_job"
|
||||
assert call_kwargs.kwargs["ttl"] >= 300 # At least 5 minutes
|
||||
|
||||
# Verify rotation logic ran (cleanup + find keys called)
|
||||
manager._cleanup_expired_deprecated_keys.assert_called_once()
|
||||
manager._find_keys_needing_rotation.assert_called_once()
|
||||
|
||||
# Verify lock was released
|
||||
mock_pod_lock_manager.release_lock.assert_called_once_with(
|
||||
cronjob_id="litellm_key_rotation_job",
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_rotations_skips_when_lock_held(self):
|
||||
"""
|
||||
When lock is held by another pod, process_rotations() should
|
||||
return early without performing any rotation.
|
||||
"""
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_pod_lock_manager = MagicMock()
|
||||
mock_pod_lock_manager.redis_cache = MagicMock()
|
||||
mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=False)
|
||||
mock_pod_lock_manager.release_lock = AsyncMock()
|
||||
|
||||
manager = KeyRotationManager(
|
||||
mock_prisma_client, pod_lock_manager=mock_pod_lock_manager
|
||||
)
|
||||
|
||||
manager._find_keys_needing_rotation = AsyncMock()
|
||||
manager._cleanup_expired_deprecated_keys = AsyncMock()
|
||||
|
||||
await manager.process_rotations()
|
||||
|
||||
# Verify lock was attempted
|
||||
mock_pod_lock_manager.acquire_lock.assert_called_once()
|
||||
|
||||
# Verify rotation logic was NOT executed
|
||||
manager._cleanup_expired_deprecated_keys.assert_not_called()
|
||||
manager._find_keys_needing_rotation.assert_not_called()
|
||||
|
||||
# Verify lock was NOT released (since it was never acquired)
|
||||
mock_pod_lock_manager.release_lock.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_rotations_releases_lock_on_success(self):
|
||||
"""
|
||||
Lock should be released in the finally block after successful rotation.
|
||||
"""
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_pod_lock_manager = MagicMock()
|
||||
mock_pod_lock_manager.redis_cache = MagicMock()
|
||||
mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True)
|
||||
mock_pod_lock_manager.release_lock = AsyncMock()
|
||||
|
||||
manager = KeyRotationManager(
|
||||
mock_prisma_client, pod_lock_manager=mock_pod_lock_manager
|
||||
)
|
||||
|
||||
# Simulate finding and rotating a key successfully
|
||||
mock_key = LiteLLM_VerificationToken(
|
||||
token="test-token",
|
||||
auto_rotate=True,
|
||||
rotation_interval="30s",
|
||||
key_rotation_at=None,
|
||||
rotation_count=0,
|
||||
key_name="test-key",
|
||||
)
|
||||
manager._find_keys_needing_rotation = AsyncMock(return_value=[mock_key])
|
||||
manager._cleanup_expired_deprecated_keys = AsyncMock()
|
||||
manager._rotate_key = AsyncMock()
|
||||
|
||||
await manager.process_rotations()
|
||||
|
||||
# Verify rotation was performed
|
||||
manager._rotate_key.assert_called_once_with(mock_key)
|
||||
|
||||
# Verify lock was released after success
|
||||
mock_pod_lock_manager.release_lock.assert_called_once_with(
|
||||
cronjob_id="litellm_key_rotation_job",
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_rotations_releases_lock_on_error(self):
|
||||
"""
|
||||
Lock should be released in the finally block even if rotation
|
||||
throws an exception.
|
||||
"""
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_pod_lock_manager = MagicMock()
|
||||
mock_pod_lock_manager.redis_cache = MagicMock()
|
||||
mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True)
|
||||
mock_pod_lock_manager.release_lock = AsyncMock()
|
||||
|
||||
manager = KeyRotationManager(
|
||||
mock_prisma_client, pod_lock_manager=mock_pod_lock_manager
|
||||
)
|
||||
|
||||
# Simulate an error during cleanup
|
||||
manager._cleanup_expired_deprecated_keys = AsyncMock(
|
||||
side_effect=Exception("Database connection failed")
|
||||
)
|
||||
|
||||
await manager.process_rotations()
|
||||
|
||||
# Verify lock was still released despite the error
|
||||
mock_pod_lock_manager.release_lock.assert_called_once_with(
|
||||
cronjob_id="litellm_key_rotation_job",
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_rotations_works_without_lock_manager(self):
|
||||
"""
|
||||
When pod_lock_manager=None, rotation should run normally
|
||||
without any lock logic (backward compat / single-pod mode).
|
||||
"""
|
||||
mock_prisma_client = AsyncMock()
|
||||
|
||||
# No pod_lock_manager provided (default None)
|
||||
manager = KeyRotationManager(mock_prisma_client)
|
||||
|
||||
manager._find_keys_needing_rotation = AsyncMock(return_value=[])
|
||||
manager._cleanup_expired_deprecated_keys = AsyncMock()
|
||||
|
||||
await manager.process_rotations()
|
||||
|
||||
# Verify rotation logic ran normally
|
||||
manager._cleanup_expired_deprecated_keys.assert_called_once()
|
||||
manager._find_keys_needing_rotation.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_rotations_works_without_redis_cache(self):
|
||||
"""
|
||||
When pod_lock_manager exists but redis_cache is None (no Redis configured),
|
||||
rotation should run normally without locking.
|
||||
"""
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_pod_lock_manager = MagicMock()
|
||||
mock_pod_lock_manager.redis_cache = None # No Redis available
|
||||
|
||||
manager = KeyRotationManager(
|
||||
mock_prisma_client, pod_lock_manager=mock_pod_lock_manager
|
||||
)
|
||||
|
||||
manager._find_keys_needing_rotation = AsyncMock(return_value=[])
|
||||
manager._cleanup_expired_deprecated_keys = AsyncMock()
|
||||
|
||||
await manager.process_rotations()
|
||||
|
||||
# Verify lock was NOT attempted (no Redis)
|
||||
mock_pod_lock_manager.acquire_lock.assert_not_called()
|
||||
|
||||
# Verify rotation logic still ran
|
||||
manager._cleanup_expired_deprecated_keys.assert_called_once()
|
||||
manager._find_keys_needing_rotation.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_rotations_handles_none_lock_result(self):
|
||||
"""
|
||||
When acquire_lock returns None (edge case), it should be treated
|
||||
as lock NOT acquired, and rotation should be skipped.
|
||||
"""
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_pod_lock_manager = MagicMock()
|
||||
mock_pod_lock_manager.redis_cache = MagicMock()
|
||||
mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=None)
|
||||
mock_pod_lock_manager.release_lock = AsyncMock()
|
||||
|
||||
manager = KeyRotationManager(
|
||||
mock_prisma_client, pod_lock_manager=mock_pod_lock_manager
|
||||
)
|
||||
|
||||
manager._find_keys_needing_rotation = AsyncMock()
|
||||
manager._cleanup_expired_deprecated_keys = AsyncMock()
|
||||
|
||||
await manager.process_rotations()
|
||||
|
||||
# Verify rotation logic was NOT executed (None treated as False via `or False`)
|
||||
manager._cleanup_expired_deprecated_keys.assert_not_called()
|
||||
manager._find_keys_needing_rotation.assert_not_called()
|
||||
|
||||
# Verify lock was NOT released (lock_acquired is False)
|
||||
mock_pod_lock_manager.release_lock.assert_not_called()
|
||||
|
|
@ -46,6 +46,28 @@ class AlwaysFailGuardrail(CustomGuardrail):
|
|||
raise HTTPException(status_code=400, detail="Content policy violation")
|
||||
|
||||
|
||||
class HttpStatusGuardrail(CustomGuardrail):
|
||||
"""Raises HTTPException with a configurable status (e.g. 503 for API outage)."""
|
||||
|
||||
def __init__(self, guardrail_name: str, status_code: int):
|
||||
super().__init__(
|
||||
guardrail_name=guardrail_name,
|
||||
event_hook="pre_call",
|
||||
default_on=True,
|
||||
)
|
||||
self.status_code = status_code
|
||||
self.calls = 0
|
||||
|
||||
def should_run_guardrail(self, data, event_type) -> bool:
|
||||
return True
|
||||
|
||||
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
|
||||
self.calls += 1
|
||||
raise HTTPException(
|
||||
status_code=self.status_code, detail="Simulated HTTP error"
|
||||
)
|
||||
|
||||
|
||||
class AlwaysPassGuardrail(CustomGuardrail):
|
||||
"""Mock guardrail that always passes."""
|
||||
|
||||
|
|
@ -350,6 +372,125 @@ async def test_guardrail_not_found_uses_on_fail():
|
|||
litellm.callbacks = original_callbacks
|
||||
|
||||
|
||||
@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed")
|
||||
@pytest.mark.asyncio
|
||||
async def test_on_error_next_fallback_on_api_outage_on_fail_blocks_content():
|
||||
"""
|
||||
Policy intervention (400) uses on_fail; technical error (503) uses on_error.
|
||||
|
||||
Primary returns 503 -> on_error: next -> fallback runs -> allow.
|
||||
"""
|
||||
primary = HttpStatusGuardrail("primary-mod", status_code=503)
|
||||
fallback = AlwaysPassGuardrail("fallback-filter")
|
||||
|
||||
pipeline = GuardrailPipeline(
|
||||
mode="pre_call",
|
||||
steps=[
|
||||
PipelineStep(
|
||||
guardrail="primary-mod",
|
||||
on_fail="block",
|
||||
on_error="next",
|
||||
on_pass="allow",
|
||||
),
|
||||
PipelineStep(
|
||||
guardrail="fallback-filter",
|
||||
on_fail="block",
|
||||
on_pass="allow",
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
original_callbacks = litellm.callbacks.copy()
|
||||
litellm.callbacks = [primary, fallback]
|
||||
|
||||
try:
|
||||
result = await PipelineExecutor.execute_steps(
|
||||
steps=pipeline.steps,
|
||||
mode=pipeline.mode,
|
||||
data={"messages": [{"role": "user", "content": "any"}]},
|
||||
user_api_key_dict=MagicMock(),
|
||||
call_type="completion",
|
||||
policy_name="mod-fallback",
|
||||
)
|
||||
|
||||
assert primary.calls == 1
|
||||
assert fallback.calls == 1
|
||||
assert result.terminal_action == "allow"
|
||||
assert result.step_results[0].outcome == "error"
|
||||
assert result.step_results[0].action_taken == "next"
|
||||
assert result.step_results[1].outcome == "pass"
|
||||
finally:
|
||||
litellm.callbacks = original_callbacks
|
||||
|
||||
|
||||
@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed")
|
||||
@pytest.mark.asyncio
|
||||
async def test_on_fail_next_on_content_on_error_block_stops_api_fallback():
|
||||
"""
|
||||
Content policy fail (400) uses on_fail: next; API error uses on_error: block (no second step).
|
||||
"""
|
||||
primary_content = AlwaysFailGuardrail("strict-mod")
|
||||
primary_api = HttpStatusGuardrail("strict-mod", status_code=503)
|
||||
fallback = AlwaysPassGuardrail("fallback-filter")
|
||||
|
||||
# Content violation: on_fail next -> would reach fallback if we had two steps
|
||||
pipeline_content = GuardrailPipeline(
|
||||
mode="pre_call",
|
||||
steps=[
|
||||
PipelineStep(
|
||||
guardrail="strict-mod",
|
||||
on_fail="next",
|
||||
on_error="block",
|
||||
on_pass="allow",
|
||||
),
|
||||
PipelineStep(
|
||||
guardrail="fallback-filter",
|
||||
on_fail="block",
|
||||
on_pass="allow",
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
original_callbacks = litellm.callbacks.copy()
|
||||
litellm.callbacks = [primary_content, fallback]
|
||||
|
||||
try:
|
||||
result = await PipelineExecutor.execute_steps(
|
||||
steps=pipeline_content.steps,
|
||||
mode=pipeline_content.mode,
|
||||
data={"messages": [{"role": "user", "content": "bad"}]},
|
||||
user_api_key_dict=MagicMock(),
|
||||
call_type="completion",
|
||||
policy_name="test",
|
||||
)
|
||||
assert result.terminal_action == "allow"
|
||||
assert primary_content.calls == 1
|
||||
assert fallback.calls == 1
|
||||
finally:
|
||||
litellm.callbacks = original_callbacks
|
||||
|
||||
# API outage: on_error block -> do not run fallback
|
||||
fallback.calls = 0
|
||||
original_callbacks = litellm.callbacks.copy()
|
||||
litellm.callbacks = [primary_api, fallback]
|
||||
try:
|
||||
result = await PipelineExecutor.execute_steps(
|
||||
steps=pipeline_content.steps,
|
||||
mode=pipeline_content.mode,
|
||||
data={"messages": [{"role": "user", "content": "ok"}]},
|
||||
user_api_key_dict=MagicMock(),
|
||||
call_type="completion",
|
||||
policy_name="test",
|
||||
)
|
||||
assert result.terminal_action == "block"
|
||||
assert primary_api.calls == 1
|
||||
assert fallback.calls == 0
|
||||
assert result.step_results[0].outcome == "error"
|
||||
assert result.step_results[0].action_taken == "block"
|
||||
finally:
|
||||
litellm.callbacks = original_callbacks
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_guardrail_not_found_with_next_continues():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -481,7 +481,7 @@ async def test_perform_health_check_and_save_passes_model_id_to_perform_health_c
|
|||
unhealthy = []
|
||||
|
||||
async def mock_perform_health_check(model_list, model=None, cli_model=None, details=True, model_id=None, max_concurrency=None):
|
||||
return healthy, unhealthy
|
||||
return healthy, unhealthy, {}
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.health_endpoints._health_endpoints.perform_health_check",
|
||||
|
|
|
|||
|
|
@ -246,7 +246,7 @@ class TestSharedHealthCheckManager:
|
|||
model_list = [{"model_name": "test-model", "litellm_params": {"model": "test-model"}}]
|
||||
|
||||
with patch("litellm.proxy.health_check_utils.shared_health_check_manager.perform_health_check") as mock_perform:
|
||||
healthy, unhealthy = await shared_health_manager.perform_shared_health_check(
|
||||
healthy, unhealthy, _ = await shared_health_manager.perform_shared_health_check(
|
||||
model_list, details=True
|
||||
)
|
||||
|
||||
|
|
@ -268,9 +268,9 @@ class TestSharedHealthCheckManager:
|
|||
expected_unhealthy = []
|
||||
|
||||
with patch("litellm.proxy.health_check_utils.shared_health_check_manager.perform_health_check") as mock_perform:
|
||||
mock_perform.return_value = (expected_healthy, expected_unhealthy)
|
||||
mock_perform.return_value = (expected_healthy, expected_unhealthy, {})
|
||||
|
||||
healthy, unhealthy = await shared_health_manager.perform_shared_health_check(
|
||||
healthy, unhealthy, _ = await shared_health_manager.perform_shared_health_check(
|
||||
model_list, details=True
|
||||
)
|
||||
|
||||
|
|
@ -302,7 +302,7 @@ class TestSharedHealthCheckManager:
|
|||
model_list = [{"model_name": "test-model", "litellm_params": {"model": "test-model"}}]
|
||||
|
||||
with patch("asyncio.sleep") as mock_sleep: # Mock sleep to avoid actual delay
|
||||
healthy, unhealthy = await shared_health_manager.perform_shared_health_check(
|
||||
healthy, unhealthy, _ = await shared_health_manager.perform_shared_health_check(
|
||||
model_list, details=True
|
||||
)
|
||||
|
||||
|
|
@ -324,9 +324,9 @@ class TestSharedHealthCheckManager:
|
|||
|
||||
with patch("asyncio.sleep") as mock_sleep, \
|
||||
patch("litellm.proxy.health_check_utils.shared_health_check_manager.perform_health_check") as mock_perform:
|
||||
mock_perform.return_value = (expected_healthy, expected_unhealthy)
|
||||
mock_perform.return_value = (expected_healthy, expected_unhealthy, {})
|
||||
|
||||
healthy, unhealthy = await shared_health_manager.perform_shared_health_check(
|
||||
healthy, unhealthy, _ = await shared_health_manager.perform_shared_health_check(
|
||||
model_list, details=True
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,789 @@
|
|||
"""
|
||||
Tests for health check failures integrating with allowed_fails_policy cooldown pipeline.
|
||||
|
||||
When enable_health_check_routing is True and a health check fails, the failure
|
||||
should increment the same counters used by allowed_fails_policy, using the
|
||||
actual exception type from the health check error.
|
||||
"""
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.proxy.health_check import run_with_timeout
|
||||
from litellm.router import Router
|
||||
from litellm.types.router import AllowedFailsPolicy
|
||||
|
||||
|
||||
def _make_model(model_id: str, model_name: str = "gpt-4") -> dict:
|
||||
return {
|
||||
"model_name": model_name,
|
||||
"litellm_params": {"model": model_name, "api_key": "fake-key"},
|
||||
"model_info": {"id": model_id},
|
||||
}
|
||||
|
||||
|
||||
class TestAhealthCheckExceptionPreservation:
|
||||
"""Test that ahealth_check() preserves the exception object in its return dict."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_with_timeout_returns_timeout_exception(self):
|
||||
"""run_with_timeout should return a litellm.Timeout in the 'exception' key on timeout."""
|
||||
import asyncio
|
||||
|
||||
async def slow_task():
|
||||
await asyncio.sleep(10)
|
||||
|
||||
result = await run_with_timeout(slow_task(), timeout=0.01)
|
||||
|
||||
assert "error" in result
|
||||
assert "exception" in result
|
||||
assert isinstance(result["exception"], litellm.Timeout)
|
||||
|
||||
|
||||
class TestHealthCheckEndpointExceptionPropagation:
|
||||
"""Test that _perform_health_check returns exceptions via exceptions_by_model_id."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unhealthy_endpoint_dict_exception_in_map(self):
|
||||
"""When ahealth_check returns {"error": ..., "exception": e}, the exception
|
||||
must appear in exceptions_by_model_id keyed by model_id — not in the endpoint dict."""
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
from litellm.proxy.health_check import _perform_health_check
|
||||
|
||||
auth_error = litellm.AuthenticationError(
|
||||
message="Invalid key", llm_provider="openai", model="gpt-4"
|
||||
)
|
||||
model_list = [
|
||||
{
|
||||
"model_name": "gpt-4",
|
||||
"litellm_params": {"model": "gpt-4", "api_key": "fake"},
|
||||
"model_info": {"id": "deploy-abc"},
|
||||
}
|
||||
]
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.health_check.litellm.ahealth_check",
|
||||
new=AsyncMock(return_value={"error": "auth failed", "exception": auth_error}),
|
||||
):
|
||||
healthy, unhealthy, exc_map = await _perform_health_check(model_list)
|
||||
|
||||
assert len(unhealthy) == 1
|
||||
assert "exception" not in unhealthy[0], "exception must not be in endpoint dict"
|
||||
assert exc_map.get("deploy-abc") is auth_error
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_raw_exception_from_gather_in_map(self):
|
||||
"""When asyncio.gather returns a raw Exception, it must appear in
|
||||
exceptions_by_model_id — not in the endpoint dict."""
|
||||
from unittest.mock import patch
|
||||
|
||||
from litellm.proxy.health_check import _perform_health_check
|
||||
|
||||
raw_exc = litellm.RateLimitError(
|
||||
message="Rate limited", llm_provider="openai", model="gpt-4"
|
||||
)
|
||||
model_list = [
|
||||
{
|
||||
"model_name": "gpt-4",
|
||||
"litellm_params": {"model": "gpt-4", "api_key": "fake"},
|
||||
"model_info": {"id": "deploy-xyz"},
|
||||
}
|
||||
]
|
||||
|
||||
# Simulate asyncio.gather returning a raw exception for this task
|
||||
with patch(
|
||||
"litellm.proxy.health_check._run_model_health_check",
|
||||
side_effect=raw_exc,
|
||||
):
|
||||
healthy, unhealthy, exc_map = await _perform_health_check(model_list)
|
||||
|
||||
assert len(unhealthy) == 1
|
||||
assert "exception" not in unhealthy[0], "exception must not be in endpoint dict"
|
||||
assert exc_map.get("deploy-xyz") is raw_exc
|
||||
|
||||
|
||||
class TestGetAllowedFailsFromPolicyWithHealthCheckExceptions:
|
||||
"""Test that get_allowed_fails_from_policy correctly resolves thresholds for health-check exceptions."""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"exception_type, policy_field, threshold",
|
||||
[
|
||||
(litellm.Timeout, "TimeoutErrorAllowedFails", 5),
|
||||
(litellm.AuthenticationError, "AuthenticationErrorAllowedFails", 3),
|
||||
(litellm.RateLimitError, "RateLimitErrorAllowedFails", 10),
|
||||
(
|
||||
litellm.ContentPolicyViolationError,
|
||||
"ContentPolicyViolationErrorAllowedFails",
|
||||
2,
|
||||
),
|
||||
(litellm.BadRequestError, "BadRequestErrorAllowedFails", 7),
|
||||
],
|
||||
)
|
||||
def test_policy_resolves_for_health_check_exception_types(
|
||||
self, exception_type, policy_field, threshold
|
||||
):
|
||||
"""Each exception type from a health check should resolve to its policy threshold."""
|
||||
policy = AllowedFailsPolicy(**{policy_field: threshold})
|
||||
router = Router(
|
||||
model_list=[_make_model("d1")],
|
||||
allowed_fails_policy=policy,
|
||||
)
|
||||
exception = exception_type(
|
||||
message="health check failed", llm_provider="openai", model="gpt-4"
|
||||
)
|
||||
result = router.get_allowed_fails_from_policy(exception=exception)
|
||||
assert result == threshold
|
||||
|
||||
def test_policy_returns_none_for_unmatched_exception(self):
|
||||
"""When no policy field matches the exception type, return None (fall back to allowed_fails)."""
|
||||
policy = AllowedFailsPolicy(TimeoutErrorAllowedFails=5)
|
||||
router = Router(
|
||||
model_list=[_make_model("d1")],
|
||||
allowed_fails_policy=policy,
|
||||
)
|
||||
# Use a generic Exception that doesn't match any policy field
|
||||
result = router.get_allowed_fails_from_policy(exception=Exception("generic"))
|
||||
assert result is None
|
||||
|
||||
|
||||
class TestHealthCheckCooldownIntegration:
|
||||
"""Test that health check failures trigger cooldown via _set_cooldown_deployments."""
|
||||
|
||||
def test_health_check_failure_increments_failed_calls(self):
|
||||
"""Health check failure should increment the failed_calls counter."""
|
||||
from litellm.router_utils.cooldown_handlers import (
|
||||
should_cooldown_based_on_allowed_fails_policy,
|
||||
)
|
||||
|
||||
router = Router(
|
||||
model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")],
|
||||
allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=3),
|
||||
)
|
||||
|
||||
timeout_exc = litellm.Timeout(
|
||||
message="Health check timeout", model="gpt-4", llm_provider="openai"
|
||||
)
|
||||
|
||||
# First call: should not cooldown (1 <= 3)
|
||||
result = should_cooldown_based_on_allowed_fails_policy(
|
||||
litellm_router_instance=router,
|
||||
deployment="deploy-1",
|
||||
original_exception=timeout_exc,
|
||||
)
|
||||
assert result is False
|
||||
|
||||
# Check counter was incremented
|
||||
current_fails = router.failed_calls.get_cache(key="deploy-1")
|
||||
assert current_fails == 1
|
||||
|
||||
def test_health_check_failure_triggers_cooldown_at_threshold(self):
|
||||
"""After exceeding allowed_fails threshold, deployment should enter cooldown."""
|
||||
from litellm.router_utils.cooldown_handlers import (
|
||||
should_cooldown_based_on_allowed_fails_policy,
|
||||
)
|
||||
|
||||
router = Router(
|
||||
model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")],
|
||||
allowed_fails_policy=AllowedFailsPolicy(AuthenticationErrorAllowedFails=2),
|
||||
)
|
||||
|
||||
auth_exc = litellm.AuthenticationError(
|
||||
message="Invalid key", model="gpt-4", llm_provider="openai"
|
||||
)
|
||||
|
||||
# Fails 1 and 2: should not cooldown
|
||||
for _ in range(2):
|
||||
result = should_cooldown_based_on_allowed_fails_policy(
|
||||
litellm_router_instance=router,
|
||||
deployment="deploy-1",
|
||||
original_exception=auth_exc,
|
||||
)
|
||||
assert result is False
|
||||
|
||||
# Fail 3: should trigger cooldown (3 > 2)
|
||||
result = should_cooldown_based_on_allowed_fails_policy(
|
||||
litellm_router_instance=router,
|
||||
deployment="deploy-1",
|
||||
original_exception=auth_exc,
|
||||
)
|
||||
assert result is True
|
||||
|
||||
def test_health_check_failure_falls_back_to_allowed_fails(self):
|
||||
"""When policy has no matching field, fall back to generic allowed_fails."""
|
||||
from litellm.router_utils.cooldown_handlers import (
|
||||
should_cooldown_based_on_allowed_fails_policy,
|
||||
)
|
||||
|
||||
router = Router(
|
||||
model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")],
|
||||
allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=10),
|
||||
allowed_fails=1,
|
||||
)
|
||||
|
||||
# Use an exception that doesn't match TimeoutErrorAllowedFails
|
||||
# InternalServerError is not checked by get_allowed_fails_from_policy
|
||||
# so it will fall back to allowed_fails=1
|
||||
generic_exc = Exception("Some internal error")
|
||||
|
||||
# Fail 1: should not cooldown (1 <= 1)
|
||||
result = should_cooldown_based_on_allowed_fails_policy(
|
||||
litellm_router_instance=router,
|
||||
deployment="deploy-1",
|
||||
original_exception=generic_exc,
|
||||
)
|
||||
assert result is False
|
||||
|
||||
# Fail 2: should trigger cooldown (2 > 1)
|
||||
result = should_cooldown_based_on_allowed_fails_policy(
|
||||
litellm_router_instance=router,
|
||||
deployment="deploy-1",
|
||||
original_exception=generic_exc,
|
||||
)
|
||||
assert result is True
|
||||
|
||||
def test_healthy_endpoints_do_not_trigger_cooldown(self):
|
||||
"""Healthy endpoints should not increment any failure counters."""
|
||||
from litellm.router_utils.cooldown_handlers import _set_cooldown_deployments
|
||||
|
||||
router = Router(
|
||||
model_list=[_make_model("deploy-1")],
|
||||
allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=1),
|
||||
enable_health_check_routing=True,
|
||||
)
|
||||
|
||||
# Simulate healthy endpoint -- no exception, no cooldown call
|
||||
healthy_endpoint = {"model_id": "deploy-1"}
|
||||
# Should have no exception key
|
||||
assert "exception" not in healthy_endpoint
|
||||
|
||||
# Verify failed_calls counter is untouched
|
||||
current_fails = router.failed_calls.get_cache(key="deploy-1")
|
||||
assert current_fails is None
|
||||
|
||||
def test_disable_cooldowns_prevents_health_check_cooldown(self):
|
||||
"""When disable_cooldowns=True, health check failures should not trigger cooldown."""
|
||||
from litellm.router_utils.cooldown_handlers import _set_cooldown_deployments
|
||||
|
||||
router = Router(
|
||||
model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")],
|
||||
allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=0),
|
||||
enable_health_check_routing=True,
|
||||
disable_cooldowns=True,
|
||||
)
|
||||
|
||||
timeout_exc = litellm.Timeout(
|
||||
message="Health check timeout", model="gpt-4", llm_provider="openai"
|
||||
)
|
||||
|
||||
result = _set_cooldown_deployments(
|
||||
litellm_router_instance=router,
|
||||
original_exception=timeout_exc,
|
||||
exception_status=500,
|
||||
deployment="deploy-1",
|
||||
time_to_cooldown=router.cooldown_time,
|
||||
)
|
||||
assert result is False
|
||||
|
||||
|
||||
class TestWriteHealthStateIntegration:
|
||||
"""Test _write_health_state_to_router_cache integrates with cooldown pipeline."""
|
||||
|
||||
def test_unhealthy_endpoint_triggers_set_cooldown(self):
|
||||
"""_write_health_state_to_router_cache should call _set_cooldown_deployments for unhealthy endpoints."""
|
||||
import litellm.proxy.proxy_server as proxy_module
|
||||
from litellm.proxy.proxy_server import _write_health_state_to_router_cache
|
||||
|
||||
router = Router(
|
||||
model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")],
|
||||
allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=5),
|
||||
enable_health_check_routing=True,
|
||||
)
|
||||
|
||||
timeout_exc = litellm.Timeout(
|
||||
message="Health check timeout", model="", llm_provider=""
|
||||
)
|
||||
|
||||
unhealthy_endpoints = [
|
||||
{"model_id": "deploy-1", "error": "timeout"},
|
||||
]
|
||||
healthy_endpoints = [
|
||||
{"model_id": "deploy-2"},
|
||||
]
|
||||
|
||||
with patch.object(proxy_module, "llm_router", router):
|
||||
with patch(
|
||||
"litellm.router_utils.cooldown_handlers._set_cooldown_deployments"
|
||||
) as mock_cooldown:
|
||||
_write_health_state_to_router_cache(
|
||||
healthy_endpoints=healthy_endpoints,
|
||||
unhealthy_endpoints=unhealthy_endpoints,
|
||||
exceptions_by_model_id={"deploy-1": timeout_exc},
|
||||
)
|
||||
mock_cooldown.assert_called_once_with(
|
||||
litellm_router_instance=router,
|
||||
original_exception=timeout_exc,
|
||||
exception_status=408, # Timeout has status_code 408
|
||||
deployment="deploy-1",
|
||||
time_to_cooldown=router.cooldown_time,
|
||||
)
|
||||
|
||||
def test_unhealthy_endpoint_without_exception_skips_cooldown(self):
|
||||
"""Unhealthy endpoints without an exception key should not trigger cooldown."""
|
||||
import litellm.proxy.proxy_server as proxy_module
|
||||
from litellm.proxy.proxy_server import _write_health_state_to_router_cache
|
||||
|
||||
router = Router(
|
||||
model_list=[_make_model("deploy-1")],
|
||||
allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=5),
|
||||
enable_health_check_routing=True,
|
||||
)
|
||||
|
||||
unhealthy_endpoints = [
|
||||
{"model_id": "deploy-1", "error": "unknown failure"},
|
||||
]
|
||||
|
||||
with patch.object(proxy_module, "llm_router", router):
|
||||
with patch(
|
||||
"litellm.router_utils.cooldown_handlers._set_cooldown_deployments"
|
||||
) as mock_cooldown:
|
||||
_write_health_state_to_router_cache(
|
||||
healthy_endpoints=[],
|
||||
unhealthy_endpoints=unhealthy_endpoints,
|
||||
# no exceptions_by_model_id → cooldown should not fire
|
||||
)
|
||||
mock_cooldown.assert_not_called()
|
||||
|
||||
def test_unhealthy_endpoint_increments_failure_counter(self):
|
||||
"""Unhealthy endpoints should call increment_deployment_failures_for_current_minute."""
|
||||
import litellm.proxy.proxy_server as proxy_module
|
||||
from litellm.proxy.proxy_server import _write_health_state_to_router_cache
|
||||
|
||||
router = Router(
|
||||
model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")],
|
||||
allowed_fails_policy=AllowedFailsPolicy(RateLimitErrorAllowedFails=10),
|
||||
enable_health_check_routing=True,
|
||||
)
|
||||
|
||||
rate_exc = litellm.RateLimitError(
|
||||
message="Rate limited", model="gpt-4", llm_provider="openai"
|
||||
)
|
||||
|
||||
unhealthy_endpoints = [
|
||||
{"model_id": "deploy-1", "error": "rate limited"},
|
||||
]
|
||||
|
||||
with patch.object(proxy_module, "llm_router", router):
|
||||
with patch(
|
||||
"litellm.router_utils.router_callbacks.track_deployment_metrics.increment_deployment_failures_for_current_minute"
|
||||
) as mock_increment:
|
||||
with patch(
|
||||
"litellm.router_utils.cooldown_handlers._set_cooldown_deployments"
|
||||
):
|
||||
_write_health_state_to_router_cache(
|
||||
healthy_endpoints=[],
|
||||
unhealthy_endpoints=unhealthy_endpoints,
|
||||
exceptions_by_model_id={"deploy-1": rate_exc},
|
||||
)
|
||||
mock_increment.assert_called_once_with(
|
||||
litellm_router_instance=router,
|
||||
deployment_id="deploy-1",
|
||||
)
|
||||
|
||||
|
||||
class TestHealthCheckFilterBypassWithPolicy:
|
||||
"""
|
||||
When allowed_fails_policy is set, the binary health check filter should be
|
||||
bypassed so cooldown is the sole routing exclusion mechanism.
|
||||
"""
|
||||
|
||||
def test_filter_bypassed_when_policy_set(self):
|
||||
"""Binary health check filter is a no-op when allowed_fails_policy is configured."""
|
||||
import time
|
||||
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.router_utils.health_state_cache import DeploymentHealthCache
|
||||
|
||||
router = Router(
|
||||
model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")],
|
||||
allowed_fails_policy=AllowedFailsPolicy(AuthenticationErrorAllowedFails=3),
|
||||
enable_health_check_routing=True,
|
||||
)
|
||||
|
||||
# Mark deploy-1 as unhealthy in the health state cache
|
||||
cache = DualCache()
|
||||
health_cache = DeploymentHealthCache(cache=cache, staleness_threshold=60.0)
|
||||
health_cache.set_deployment_health_states(
|
||||
{
|
||||
"deploy-1": {
|
||||
"is_healthy": False,
|
||||
"timestamp": time.time(),
|
||||
"reason": "test",
|
||||
},
|
||||
}
|
||||
)
|
||||
router.health_state_cache = health_cache
|
||||
|
||||
deployments = [_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")]
|
||||
|
||||
# Filter should pass all through because policy is set
|
||||
result = router._filter_health_check_unhealthy_deployments(deployments)
|
||||
assert (
|
||||
len(result) == 2
|
||||
), "Binary filter should be bypassed when allowed_fails_policy is set"
|
||||
|
||||
def test_filter_active_when_no_policy(self):
|
||||
"""Binary health check filter still works when no allowed_fails_policy is configured."""
|
||||
import time
|
||||
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.router_utils.health_state_cache import DeploymentHealthCache
|
||||
|
||||
router = Router(
|
||||
model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")],
|
||||
enable_health_check_routing=True,
|
||||
)
|
||||
|
||||
cache = DualCache()
|
||||
health_cache = DeploymentHealthCache(cache=cache, staleness_threshold=60.0)
|
||||
health_cache.set_deployment_health_states(
|
||||
{
|
||||
"deploy-1": {
|
||||
"is_healthy": False,
|
||||
"timestamp": time.time(),
|
||||
"reason": "test",
|
||||
},
|
||||
}
|
||||
)
|
||||
router.health_state_cache = health_cache
|
||||
|
||||
deployments = [_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")]
|
||||
|
||||
result = router._filter_health_check_unhealthy_deployments(deployments)
|
||||
assert len(result) == 1
|
||||
assert result[0]["model_info"]["id"] == "deploy-2"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_filter_bypassed_when_policy_set(self):
|
||||
"""Async version also bypasses when allowed_fails_policy is set."""
|
||||
import time
|
||||
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.router_utils.health_state_cache import DeploymentHealthCache
|
||||
|
||||
router = Router(
|
||||
model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")],
|
||||
allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=2),
|
||||
enable_health_check_routing=True,
|
||||
)
|
||||
|
||||
cache = DualCache()
|
||||
health_cache = DeploymentHealthCache(cache=cache, staleness_threshold=60.0)
|
||||
health_cache.set_deployment_health_states(
|
||||
{
|
||||
"deploy-1": {
|
||||
"is_healthy": False,
|
||||
"timestamp": time.time(),
|
||||
"reason": "test",
|
||||
},
|
||||
}
|
||||
)
|
||||
router.health_state_cache = health_cache
|
||||
|
||||
deployments = [_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")]
|
||||
|
||||
result = await router._async_filter_health_check_unhealthy_deployments(
|
||||
deployments
|
||||
)
|
||||
assert len(result) == 2
|
||||
|
||||
|
||||
class TestAllDeploymentsInCooldownSafetyNet:
|
||||
"""
|
||||
When enable_health_check_routing=True and ALL deployments enter cooldown,
|
||||
the async routing path should bypass the cooldown filter and return all
|
||||
deployments rather than blocking all traffic.
|
||||
"""
|
||||
|
||||
def test_raw_cooldown_filter_returns_empty_when_all_cooled(self):
|
||||
"""The raw _filter_cooldown_deployments has no safety net -- it returns empty."""
|
||||
router = Router(
|
||||
model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")],
|
||||
enable_health_check_routing=True,
|
||||
)
|
||||
deployments = [_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")]
|
||||
result = router._filter_cooldown_deployments(
|
||||
healthy_deployments=deployments,
|
||||
cooldown_deployments=["deploy-1", "deploy-2"],
|
||||
)
|
||||
assert result == [] # raw filter has no safety net
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_routing_path_bypasses_all_cooldown(self):
|
||||
"""In the async routing path, all-in-cooldown with enable_health_check_routing
|
||||
returns the full list instead of empty (safety net)."""
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from litellm.router_utils.cooldown_handlers import (
|
||||
_async_get_cooldown_deployments,
|
||||
)
|
||||
|
||||
router = Router(
|
||||
model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")],
|
||||
allowed_fails_policy=AllowedFailsPolicy(AuthenticationErrorAllowedFails=0),
|
||||
enable_health_check_routing=True,
|
||||
)
|
||||
|
||||
deployments = [_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")]
|
||||
|
||||
# Simulate all deployments in cooldown
|
||||
with patch(
|
||||
"litellm.router._async_get_cooldown_deployments",
|
||||
new=AsyncMock(return_value=["deploy-1", "deploy-2"]),
|
||||
):
|
||||
# The safety net in async_get_available_deployment should restore
|
||||
# all deployments when the cooldown filter empties the list
|
||||
_pre = deployments.copy()
|
||||
filtered = router._filter_cooldown_deployments(
|
||||
healthy_deployments=deployments,
|
||||
cooldown_deployments=["deploy-1", "deploy-2"],
|
||||
)
|
||||
# If filtered is empty and enable_health_check_routing is True,
|
||||
# the routing path restores _pre_cooldown_deployments
|
||||
if not filtered and router.enable_health_check_routing:
|
||||
filtered = _pre
|
||||
|
||||
assert (
|
||||
len(filtered) == 2
|
||||
), "Safety net should return all deployments when all are in cooldown"
|
||||
|
||||
|
||||
class TestHealthCheckIgnoreTransientErrors:
|
||||
"""
|
||||
When health_check_ignore_transient_errors=True, health check failures with
|
||||
429 or 408 status codes should NOT increment failure counters or trigger cooldown.
|
||||
401, 404, and 5xx errors should still be processed normally.
|
||||
"""
|
||||
|
||||
def test_429_skipped_when_flag_enabled(self):
|
||||
"""429 from health check does not trigger cooldown when flag is set."""
|
||||
import litellm.proxy.proxy_server as proxy_module
|
||||
from litellm.proxy.proxy_server import _write_health_state_to_router_cache
|
||||
|
||||
router = Router(
|
||||
model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")],
|
||||
allowed_fails_policy=AllowedFailsPolicy(RateLimitErrorAllowedFails=0),
|
||||
enable_health_check_routing=True,
|
||||
health_check_ignore_transient_errors=True,
|
||||
)
|
||||
|
||||
rate_exc = litellm.RateLimitError(
|
||||
message="Rate limited", model="gpt-4", llm_provider="openai"
|
||||
)
|
||||
assert getattr(rate_exc, "status_code", None) == 429
|
||||
|
||||
unhealthy_endpoints = [
|
||||
{"model_id": "deploy-1", "error": "rate limited"},
|
||||
]
|
||||
|
||||
with patch.object(proxy_module, "llm_router", router):
|
||||
with patch(
|
||||
"litellm.router_utils.cooldown_handlers._set_cooldown_deployments"
|
||||
) as mock_cooldown:
|
||||
with patch(
|
||||
"litellm.router_utils.router_callbacks.track_deployment_metrics.increment_deployment_failures_for_current_minute"
|
||||
) as mock_increment:
|
||||
_write_health_state_to_router_cache(
|
||||
healthy_endpoints=[],
|
||||
unhealthy_endpoints=unhealthy_endpoints,
|
||||
exceptions_by_model_id={"deploy-1": rate_exc},
|
||||
)
|
||||
mock_cooldown.assert_not_called()
|
||||
mock_increment.assert_not_called()
|
||||
|
||||
def test_408_skipped_when_flag_enabled(self):
|
||||
"""408 from health check does not trigger cooldown when flag is set."""
|
||||
import litellm.proxy.proxy_server as proxy_module
|
||||
from litellm.proxy.proxy_server import _write_health_state_to_router_cache
|
||||
|
||||
router = Router(
|
||||
model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")],
|
||||
allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=0),
|
||||
enable_health_check_routing=True,
|
||||
health_check_ignore_transient_errors=True,
|
||||
)
|
||||
|
||||
timeout_exc = litellm.Timeout(
|
||||
message="Health check timeout exceeded", model="", llm_provider=""
|
||||
)
|
||||
|
||||
unhealthy_endpoints = [
|
||||
{"model_id": "deploy-1", "error": "timeout"},
|
||||
]
|
||||
|
||||
with patch.object(proxy_module, "llm_router", router):
|
||||
with patch(
|
||||
"litellm.router_utils.cooldown_handlers._set_cooldown_deployments"
|
||||
) as mock_cooldown:
|
||||
_write_health_state_to_router_cache(
|
||||
healthy_endpoints=[],
|
||||
unhealthy_endpoints=unhealthy_endpoints,
|
||||
exceptions_by_model_id={"deploy-1": timeout_exc},
|
||||
)
|
||||
mock_cooldown.assert_not_called()
|
||||
|
||||
def test_401_still_triggers_cooldown_when_flag_enabled(self):
|
||||
"""Auth errors (401) still trigger cooldown even when flag is set."""
|
||||
import litellm.proxy.proxy_server as proxy_module
|
||||
from litellm.proxy.proxy_server import _write_health_state_to_router_cache
|
||||
|
||||
router = Router(
|
||||
model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")],
|
||||
allowed_fails_policy=AllowedFailsPolicy(AuthenticationErrorAllowedFails=0),
|
||||
enable_health_check_routing=True,
|
||||
health_check_ignore_transient_errors=True,
|
||||
)
|
||||
|
||||
auth_exc = litellm.AuthenticationError(
|
||||
message="Invalid key", model="gpt-4", llm_provider="openai"
|
||||
)
|
||||
|
||||
unhealthy_endpoints = [
|
||||
{"model_id": "deploy-1", "error": "auth failed"},
|
||||
]
|
||||
|
||||
with patch.object(proxy_module, "llm_router", router):
|
||||
with patch(
|
||||
"litellm.router_utils.cooldown_handlers._set_cooldown_deployments"
|
||||
) as mock_cooldown:
|
||||
_write_health_state_to_router_cache(
|
||||
healthy_endpoints=[],
|
||||
unhealthy_endpoints=unhealthy_endpoints,
|
||||
exceptions_by_model_id={"deploy-1": auth_exc},
|
||||
)
|
||||
mock_cooldown.assert_called_once()
|
||||
|
||||
def test_429_not_written_to_health_state_cache_when_flag_enabled(self):
|
||||
"""429 endpoint is excluded from health state cache when flag is set,
|
||||
so the binary health check filter does not mark it as unhealthy."""
|
||||
import litellm.proxy.proxy_server as proxy_module
|
||||
from litellm.proxy.proxy_server import _write_health_state_to_router_cache
|
||||
|
||||
router = Router(
|
||||
model_list=[_make_model("deploy-1")],
|
||||
enable_health_check_routing=True,
|
||||
health_check_ignore_transient_errors=True,
|
||||
)
|
||||
|
||||
rate_exc = litellm.RateLimitError(
|
||||
message="Rate limited", model="gpt-4", llm_provider="openai"
|
||||
)
|
||||
|
||||
unhealthy_endpoints = [
|
||||
{"model_id": "deploy-1", "error": "rate limited"},
|
||||
]
|
||||
|
||||
with patch.object(proxy_module, "llm_router", router):
|
||||
_write_health_state_to_router_cache(
|
||||
healthy_endpoints=[],
|
||||
unhealthy_endpoints=unhealthy_endpoints,
|
||||
exceptions_by_model_id={"deploy-1": rate_exc},
|
||||
)
|
||||
|
||||
# Health state cache should have NO entry for deploy-1
|
||||
# (429 was ignored, not written as unhealthy)
|
||||
unhealthy_ids = router.health_state_cache.get_unhealthy_deployment_ids()
|
||||
assert "deploy-1" not in unhealthy_ids
|
||||
|
||||
def test_429_triggers_cooldown_when_flag_disabled(self):
|
||||
"""When flag is False (default), 429 still triggers cooldown."""
|
||||
import litellm.proxy.proxy_server as proxy_module
|
||||
from litellm.proxy.proxy_server import _write_health_state_to_router_cache
|
||||
|
||||
router = Router(
|
||||
model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")],
|
||||
allowed_fails_policy=AllowedFailsPolicy(RateLimitErrorAllowedFails=0),
|
||||
enable_health_check_routing=True,
|
||||
health_check_ignore_transient_errors=False,
|
||||
)
|
||||
|
||||
rate_exc = litellm.RateLimitError(
|
||||
message="Rate limited", model="gpt-4", llm_provider="openai"
|
||||
)
|
||||
|
||||
unhealthy_endpoints = [
|
||||
{"model_id": "deploy-1", "error": "rate limited"},
|
||||
]
|
||||
|
||||
with patch.object(proxy_module, "llm_router", router):
|
||||
with patch(
|
||||
"litellm.router_utils.cooldown_handlers._set_cooldown_deployments"
|
||||
) as mock_cooldown:
|
||||
_write_health_state_to_router_cache(
|
||||
healthy_endpoints=[],
|
||||
unhealthy_endpoints=unhealthy_endpoints,
|
||||
exceptions_by_model_id={"deploy-1": rate_exc},
|
||||
)
|
||||
mock_cooldown.assert_called_once()
|
||||
|
||||
|
||||
class TestSharedCacheTransientErrorFilter:
|
||||
"""
|
||||
When SharedHealthCheckManager returns cached results, exceptions_by_model_id
|
||||
is always {}. The filter must fall back to the 'exception_status' field stored
|
||||
on each endpoint dict so 429/408 endpoints are still excluded correctly.
|
||||
"""
|
||||
|
||||
def test_cached_429_excluded_via_exception_status_field(self):
|
||||
"""Cache-hit path: endpoint with exception_status=429 is excluded from health state."""
|
||||
import litellm.proxy.proxy_server as proxy_module
|
||||
from litellm.proxy.proxy_server import _write_health_state_to_router_cache
|
||||
|
||||
router = Router(
|
||||
model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")],
|
||||
enable_health_check_routing=True,
|
||||
health_check_ignore_transient_errors=True,
|
||||
)
|
||||
|
||||
# Simulate a cache-hit endpoint: exception_status stored as int, no exceptions dict
|
||||
unhealthy_endpoints = [
|
||||
{"model_id": "deploy-1", "error": "rate limited", "exception_status": 429},
|
||||
]
|
||||
|
||||
with patch.object(proxy_module, "llm_router", router):
|
||||
_write_health_state_to_router_cache(
|
||||
healthy_endpoints=[],
|
||||
unhealthy_endpoints=unhealthy_endpoints,
|
||||
exceptions_by_model_id={},
|
||||
)
|
||||
|
||||
# deploy-1 should NOT be marked unhealthy (429 was filtered)
|
||||
unhealthy_ids = router.health_state_cache.get_unhealthy_deployment_ids()
|
||||
assert "deploy-1" not in unhealthy_ids
|
||||
|
||||
def test_cached_401_still_marked_unhealthy(self):
|
||||
"""Cache-hit path: endpoint with exception_status=401 is still written as unhealthy."""
|
||||
import litellm.proxy.proxy_server as proxy_module
|
||||
from litellm.proxy.proxy_server import _write_health_state_to_router_cache
|
||||
|
||||
router = Router(
|
||||
model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")],
|
||||
enable_health_check_routing=True,
|
||||
health_check_ignore_transient_errors=True,
|
||||
)
|
||||
|
||||
unhealthy_endpoints = [
|
||||
{"model_id": "deploy-1", "error": "auth failed", "exception_status": 401},
|
||||
]
|
||||
|
||||
with patch.object(proxy_module, "llm_router", router):
|
||||
_write_health_state_to_router_cache(
|
||||
healthy_endpoints=[],
|
||||
unhealthy_endpoints=unhealthy_endpoints,
|
||||
exceptions_by_model_id={},
|
||||
)
|
||||
|
||||
unhealthy_ids = router.health_state_cache.get_unhealthy_deployment_ids()
|
||||
assert "deploy-1" in unhealthy_ids
|
||||
|
|
@ -50,6 +50,7 @@ class TestFilterHealthCheckUnhealthyDeployments:
|
|||
def __init__(self):
|
||||
self.enable_health_check_routing = enable
|
||||
self.health_state_cache = health_cache
|
||||
self.allowed_fails_policy = None
|
||||
|
||||
# Import the actual method and bind it
|
||||
from litellm.router import Router
|
||||
|
|
@ -125,6 +126,7 @@ class TestAsyncFilterHealthCheckUnhealthyDeployments:
|
|||
def __init__(self):
|
||||
self.enable_health_check_routing = enable
|
||||
self.health_state_cache = health_cache
|
||||
self.allowed_fails_policy = None
|
||||
|
||||
fake = FakeRouter()
|
||||
fake._async_filter_health_check_unhealthy_deployments = (
|
||||
|
|
|
|||
|
|
@ -5,7 +5,6 @@ import sys
|
|||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
sys.path.insert(
|
||||
0, os.path.abspath("../../..")
|
||||
|
|
@ -13,7 +12,6 @@ sys.path.insert(
|
|||
|
||||
|
||||
import litellm
|
||||
from litellm.router_utils.fallback_event_handlers import run_async_fallback
|
||||
|
||||
|
||||
def test_update_kwargs_does_not_mutate_defaults_and_merges_metadata():
|
||||
|
|
@ -127,7 +125,7 @@ async def test_async_router_acreate_file():
|
|||
"""
|
||||
Write to all deployments of a model
|
||||
"""
|
||||
from unittest.mock import MagicMock, call, patch
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
|
|
@ -747,7 +745,7 @@ async def test_router_ageneric_api_call_with_fallbacks_helper():
|
|||
"""
|
||||
Test the _ageneric_api_call_with_fallbacks_helper method with various scenarios
|
||||
"""
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
from unittest.mock import patch
|
||||
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
|
|
@ -1134,10 +1132,9 @@ def test_get_model_access_groups_cache_invalidation_upsert_deployment():
|
|||
@pytest.mark.asyncio
|
||||
async def test_acompletion_streaming_iterator():
|
||||
"""Test _acompletion_streaming_iterator for normal streaming and fallback behavior."""
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from litellm.exceptions import MidStreamFallbackError
|
||||
from litellm.types.utils import ModelResponseStream
|
||||
|
||||
# Helper class for creating async iterators
|
||||
class AsyncIterator:
|
||||
|
|
@ -2847,3 +2844,286 @@ def test_combine_fallback_usage():
|
|||
assert chunk.usage.prompt_tokens == 10
|
||||
assert chunk.usage.completion_tokens == 5
|
||||
assert chunk.usage.total_tokens == 15
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_scoped_model_fallback():
|
||||
"""
|
||||
Test that fallback works correctly for team-scoped models.
|
||||
|
||||
When a team-scoped model fails and the fallback model is also team-scoped,
|
||||
the router should find the fallback deployment by matching team_public_model_name.
|
||||
"""
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "team-a-primary-internal",
|
||||
"litellm_params": {"model": "gpt-3.5-turbo", "api_key": "fake"},
|
||||
"model_info": {
|
||||
"team_id": "team-a",
|
||||
"team_public_model_name": "primary-model",
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "team-a-fallback-internal",
|
||||
"litellm_params": {
|
||||
"model": "gpt-4",
|
||||
"api_key": "fake",
|
||||
"mock_response": "fallback success from team-a",
|
||||
},
|
||||
"model_info": {
|
||||
"team_id": "team-a",
|
||||
"team_public_model_name": "fallback-model",
|
||||
},
|
||||
},
|
||||
],
|
||||
fallbacks=[{"primary-model": ["fallback-model"]}],
|
||||
)
|
||||
|
||||
response = await router.acompletion(
|
||||
model="primary-model",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
metadata={"user_api_key_team_id": "team-a"},
|
||||
mock_testing_fallbacks=True,
|
||||
)
|
||||
assert response is not None
|
||||
assert response.choices[0].message.content == "fallback success from team-a"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_scoped_model_fallback_to_global():
|
||||
"""
|
||||
Test that a team-scoped model can fall back to a global (non-team) model.
|
||||
|
||||
Global models (no team_id on deployment) should be accessible as fallback
|
||||
targets for team-scoped requests.
|
||||
"""
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "team-a-primary-internal",
|
||||
"litellm_params": {"model": "gpt-3.5-turbo", "api_key": "fake"},
|
||||
"model_info": {
|
||||
"team_id": "team-a",
|
||||
"team_public_model_name": "primary-model",
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "global-fallback",
|
||||
"litellm_params": {
|
||||
"model": "gpt-4",
|
||||
"api_key": "fake",
|
||||
"mock_response": "global fallback success",
|
||||
},
|
||||
},
|
||||
],
|
||||
fallbacks=[{"primary-model": ["global-fallback"]}],
|
||||
)
|
||||
|
||||
response = await router.acompletion(
|
||||
model="primary-model",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
metadata={"user_api_key_team_id": "team-a"},
|
||||
mock_testing_fallbacks=True,
|
||||
)
|
||||
assert response is not None
|
||||
assert response.choices[0].message.content == "global fallback success"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_scoped_model_fallback_cross_team_blocked():
|
||||
"""
|
||||
Test that cross-team fallback is correctly blocked.
|
||||
|
||||
When team-a's model fails and the fallback target is scoped to team-b,
|
||||
the router should NOT use it (team isolation).
|
||||
"""
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "team-a-primary-internal",
|
||||
"litellm_params": {"model": "gpt-3.5-turbo", "api_key": "fake"},
|
||||
"model_info": {
|
||||
"team_id": "team-a",
|
||||
"team_public_model_name": "primary-model",
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "team-b-fallback-internal",
|
||||
"litellm_params": {
|
||||
"model": "gpt-4",
|
||||
"api_key": "fake",
|
||||
"mock_response": "team-b response - should not reach here",
|
||||
},
|
||||
"model_info": {
|
||||
"team_id": "team-b",
|
||||
"team_public_model_name": "fallback-model",
|
||||
},
|
||||
},
|
||||
],
|
||||
fallbacks=[{"primary-model": ["fallback-model"]}],
|
||||
)
|
||||
|
||||
with pytest.raises(Exception):
|
||||
await router.acompletion(
|
||||
model="primary-model",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
metadata={"user_api_key_team_id": "team-a"},
|
||||
mock_testing_fallbacks=True,
|
||||
)
|
||||
|
||||
|
||||
def test_get_all_deployments_with_team_id():
|
||||
"""
|
||||
Test that _get_all_deployments with team_id can find deployments
|
||||
by team_public_model_name when the model_name is not in the index.
|
||||
"""
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "internal-team-deployment",
|
||||
"litellm_params": {"model": "gpt-4", "api_key": "fake"},
|
||||
"model_info": {
|
||||
"team_id": "team-x",
|
||||
"team_public_model_name": "gpt-4",
|
||||
},
|
||||
},
|
||||
],
|
||||
)
|
||||
|
||||
# Without team_id: "gpt-4" is not in the model_name index (internal name is different)
|
||||
deployments = router._get_all_deployments(model_name="gpt-4")
|
||||
assert len(deployments) == 0
|
||||
|
||||
# With correct team_id: should find via O(n) scan matching team_public_model_name
|
||||
deployments = router._get_all_deployments(model_name="gpt-4", team_id="team-x")
|
||||
assert len(deployments) == 1
|
||||
assert deployments[0]["model_name"] == "internal-team-deployment"
|
||||
|
||||
# With wrong team_id: should find nothing
|
||||
deployments = router._get_all_deployments(model_name="gpt-4", team_id="team-y")
|
||||
assert len(deployments) == 0
|
||||
|
||||
|
||||
def test_multiregion_team_deployments_unique_model_names():
|
||||
"""
|
||||
Simulates athenahealth's exact setup: unique model_names per deployment,
|
||||
same team_public_model_name, multiple regions.
|
||||
|
||||
Verifies that _get_all_deployments returns ALL regional deployments
|
||||
for a team when queried by team_public_model_name.
|
||||
"""
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "metis-claude-us-east-1",
|
||||
"litellm_params": {
|
||||
"model": "bedrock/anthropic.claude-3-sonnet",
|
||||
"aws_region_name": "us-east-1",
|
||||
"api_key": "fake",
|
||||
},
|
||||
"model_info": {
|
||||
"team_id": "metis-team",
|
||||
"team_public_model_name": "claude-sonnet",
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "metis-claude-us-west-2",
|
||||
"litellm_params": {
|
||||
"model": "bedrock/anthropic.claude-3-sonnet",
|
||||
"aws_region_name": "us-west-2",
|
||||
"api_key": "fake",
|
||||
},
|
||||
"model_info": {
|
||||
"team_id": "metis-team",
|
||||
"team_public_model_name": "claude-sonnet",
|
||||
},
|
||||
},
|
||||
],
|
||||
)
|
||||
|
||||
# "claude-sonnet" is NOT in the model_name index
|
||||
assert "claude-sonnet" not in router.model_names
|
||||
|
||||
# Without team_id: returns nothing (no model_name="claude-sonnet" in index, no O(n) scan)
|
||||
deployments = router._get_all_deployments(model_name="claude-sonnet")
|
||||
assert len(deployments) == 0
|
||||
|
||||
# With team_id: O(n) scan finds BOTH regional deployments
|
||||
deployments = router._get_all_deployments(
|
||||
model_name="claude-sonnet", team_id="metis-team"
|
||||
)
|
||||
assert len(deployments) == 2
|
||||
deployment_names = {d["model_name"] for d in deployments}
|
||||
assert deployment_names == {"metis-claude-us-east-1", "metis-claude-us-west-2"}
|
||||
|
||||
# Each deployment has a unique ID (critical for cooldown/retry to work)
|
||||
deployment_ids = {d["model_info"]["id"] for d in deployments}
|
||||
assert len(deployment_ids) == 2, "Each deployment must have a unique ID for cooldown tracking"
|
||||
|
||||
# Wrong team: returns nothing
|
||||
deployments = router._get_all_deployments(
|
||||
model_name="claude-sonnet", team_id="other-team"
|
||||
)
|
||||
assert len(deployments) == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multiregion_team_failover_between_regions():
|
||||
"""
|
||||
Simulates athenahealth's multiregion failover scenario:
|
||||
- Two Bedrock deployments (us-east-1 and us-west-2) with unique model_names
|
||||
- Same team_public_model_name ("claude-sonnet")
|
||||
- Primary region fails → router should failover to second region
|
||||
|
||||
This is the exact scenario Sean Glover from athenahealth will demonstrate.
|
||||
"""
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "metis-claude-us-east-1",
|
||||
"litellm_params": {
|
||||
"model": "bedrock/anthropic.claude-3-sonnet",
|
||||
"api_key": "fake",
|
||||
"mock_response": "response from us-east-1",
|
||||
},
|
||||
"model_info": {
|
||||
"team_id": "metis-team",
|
||||
"team_public_model_name": "claude-sonnet",
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "metis-claude-us-west-2",
|
||||
"litellm_params": {
|
||||
"model": "bedrock/anthropic.claude-3-sonnet",
|
||||
"api_key": "fake",
|
||||
"mock_response": "response from us-west-2",
|
||||
},
|
||||
"model_info": {
|
||||
"team_id": "metis-team",
|
||||
"team_public_model_name": "claude-sonnet",
|
||||
},
|
||||
},
|
||||
],
|
||||
num_retries=1,
|
||||
)
|
||||
|
||||
# Verify the router finds both deployments for the team
|
||||
deployments = router._get_all_deployments(
|
||||
model_name="claude-sonnet", team_id="metis-team"
|
||||
)
|
||||
assert len(deployments) == 2, (
|
||||
"Router must find both regional deployments by team_public_model_name"
|
||||
)
|
||||
|
||||
# Make a normal request — should succeed from one of the regions
|
||||
response = await router.acompletion(
|
||||
model="claude-sonnet",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
metadata={"user_api_key_team_id": "metis-team"},
|
||||
)
|
||||
assert response is not None
|
||||
assert response.choices[0].message.content in [
|
||||
"response from us-east-1",
|
||||
"response from us-west-2",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ def test_pipeline_step_defaults():
|
|||
step = PipelineStep(guardrail="my-guard")
|
||||
assert step.on_fail == "block"
|
||||
assert step.on_pass == "allow"
|
||||
assert step.on_error is None
|
||||
assert step.pass_data is False
|
||||
assert step.modify_response_message is None
|
||||
|
||||
|
|
@ -33,9 +34,10 @@ def test_pipeline_step_valid_actions():
|
|||
|
||||
def test_pipeline_step_all_action_types():
|
||||
for action in ("allow", "block", "next", "modify_response"):
|
||||
step = PipelineStep(guardrail="g", on_fail=action, on_pass=action)
|
||||
step = PipelineStep(guardrail="g", on_fail=action, on_pass=action, on_error=action)
|
||||
assert step.on_fail == action
|
||||
assert step.on_pass == action
|
||||
assert step.on_error == action
|
||||
|
||||
|
||||
def test_pipeline_step_invalid_action_rejected():
|
||||
|
|
@ -48,6 +50,16 @@ def test_pipeline_step_invalid_on_pass_rejected():
|
|||
PipelineStep(guardrail="my-guard", on_pass="skip")
|
||||
|
||||
|
||||
def test_pipeline_step_on_error_valid():
|
||||
step = PipelineStep(guardrail="g", on_error="next", on_fail="block", on_pass="allow")
|
||||
assert step.on_error == "next"
|
||||
|
||||
|
||||
def test_pipeline_step_invalid_on_error_rejected():
|
||||
with pytest.raises(ValidationError):
|
||||
PipelineStep(guardrail="my-guard", on_error="invalid")
|
||||
|
||||
|
||||
def test_pipeline_requires_at_least_one_step():
|
||||
with pytest.raises(ValidationError):
|
||||
GuardrailPipeline(mode="pre_call", steps=[])
|
||||
|
|
|
|||
|
|
@ -155,6 +155,14 @@ const FailIcon: React.FC = () => (
|
|||
</svg>
|
||||
);
|
||||
|
||||
const ApiFailureIcon: React.FC = () => (
|
||||
<svg width="14" height="14" viewBox="0 0 24 24" fill="none" stroke="#d97706" strokeWidth="2" strokeLinecap="round" strokeLinejoin="round" style={{ flexShrink: 0 }}>
|
||||
<path d="M10.29 3.86L1.82 18a2 2 0 0 0 1.71 3h16.94a2 2 0 0 0 1.71-3L13.71 3.86a2 2 0 0 0-3.42 0z" />
|
||||
<line x1="12" y1="9" x2="12" y2="13" />
|
||||
<line x1="12" y1="17" x2="12.01" y2="17" />
|
||||
</svg>
|
||||
);
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────────────
|
||||
// Connector
|
||||
// ─────────────────────────────────────────────────────────────────────────────
|
||||
|
|
@ -349,6 +357,41 @@ const StepCard: React.FC<StepCardProps> = ({
|
|||
</div>
|
||||
)}
|
||||
</div>
|
||||
|
||||
{/* ON API FAILURE (technical / provider outage) — optional; defaults to ON FAIL */}
|
||||
<div style={{ borderTop: "1px solid #f0f0f0", padding: "14px 20px" }}>
|
||||
<div className="flex items-center gap-2" style={{ marginBottom: 8 }}>
|
||||
<ApiFailureIcon />
|
||||
<span style={{ fontSize: 13, fontWeight: 600, color: "#374151" }}>ON API FAILURE</span>
|
||||
</div>
|
||||
<label style={{ fontSize: 12, fontWeight: 500, color: "#6b7280", display: "block", marginBottom: 6 }}>
|
||||
Action
|
||||
</label>
|
||||
<Select
|
||||
style={{ width: "100%" }}
|
||||
placeholder="Same as ON FAIL"
|
||||
allowClear
|
||||
value={step.on_error ?? undefined}
|
||||
onChange={(value) =>
|
||||
onChange({
|
||||
on_error: value === undefined || value === null ? undefined : (value as PipelineStep["on_error"]),
|
||||
})
|
||||
}
|
||||
options={ACTION_OPTIONS}
|
||||
/>
|
||||
{step.on_error === "modify_response" && step.on_fail !== "modify_response" && (
|
||||
<div style={{ marginTop: 8 }}>
|
||||
<label style={{ fontSize: 12, fontWeight: 500, color: "#6b7280", display: "block", marginBottom: 6 }}>
|
||||
Custom Response Message
|
||||
</label>
|
||||
<TextInput
|
||||
placeholder="Enter custom response..."
|
||||
value={step.modify_response_message || ""}
|
||||
onChange={(e) => onChange({ modify_response_message: e.target.value || null })}
|
||||
/>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
|
@ -566,13 +609,19 @@ export const PipelineInfoDisplay: React.FC<PipelineInfoDisplayProps> = ({ pipeli
|
|||
{/* Divider */}
|
||||
<div style={{ borderTop: "1px solid #f3f4f6", marginBottom: 10 }} />
|
||||
|
||||
{/* Pass / Fail */}
|
||||
<div className="flex items-center gap-6" style={{ fontSize: 13, color: "#374151" }}>
|
||||
{/* Pass / Fail / API failure */}
|
||||
<div className="flex flex-col gap-2" style={{ fontSize: 13, color: "#374151" }}>
|
||||
<span className="flex items-center gap-1.5">
|
||||
<PassIcon /> Pass → {ACTION_LABELS[step.on_pass] || step.on_pass}
|
||||
</span>
|
||||
<span className="flex items-center gap-1.5">
|
||||
<FailIcon /> Fail → {ACTION_LABELS[step.on_fail] || step.on_fail}
|
||||
<FailIcon /> On fail → {ACTION_LABELS[step.on_fail] || step.on_fail}
|
||||
</span>
|
||||
<span className="flex items-center gap-1.5">
|
||||
<ApiFailureIcon /> On API failure →{" "}
|
||||
{step.on_error != null
|
||||
? ACTION_LABELS[step.on_error] || step.on_error
|
||||
: `${ACTION_LABELS[step.on_fail] || step.on_fail} (same as on fail)`}
|
||||
</span>
|
||||
</div>
|
||||
</div>
|
||||
|
|
|
|||
|
|
@ -24,6 +24,8 @@ export interface PipelineStep {
|
|||
guardrail: string;
|
||||
on_fail: "block" | "allow" | "next" | "modify_response";
|
||||
on_pass: "allow" | "block" | "next" | "modify_response";
|
||||
/** If unset, API/technical errors use on_fail (backward compatible). */
|
||||
on_error?: "block" | "allow" | "next" | "modify_response" | null;
|
||||
pass_data?: boolean;
|
||||
modify_response_message?: string | null;
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue