diff --git a/scripts/health_check/health_check_client.py b/scripts/health_check/health_check_client.py index 4848735bb7..873727b453 100644 --- a/scripts/health_check/health_check_client.py +++ b/scripts/health_check/health_check_client.py @@ -20,6 +20,17 @@ from typing import Dict, List, Optional, Tuple import httpx import yaml +# Default prompt for health checks - exactly 100k characters +# Generate a repeating pattern to reach exactly 100,000 characters +_base_text = "This is a health check test prompt for LiteLLM proxy. " +_repeat_count = (100000 // len(_base_text)) + 1 +_DEFAULT_COMPLETION_PROMPT = (_base_text * _repeat_count)[:100000] + +# Default embedding text - also exactly 100k characters +_embedding_base_text = "This is a test for vectorization. " +_embedding_repeat_count = (100000 // len(_embedding_base_text)) + 1 +_DEFAULT_EMBEDDING_TEXT = (_embedding_base_text * _embedding_repeat_count)[:100000] + class LiteLLMHealthCheckClient: """Client for health checking LiteLLM proxy models.""" @@ -29,8 +40,9 @@ class LiteLLMHealthCheckClient: base_url: str, api_key: str, timeout: int = 120, # Match Go implementation's 120s timeout - completion_prompt: str = "Say this is a test", # Match Go implementation - embedding_text: str = "This is a test for vectorization.", # Match Go implementation + completion_prompt: str = _DEFAULT_COMPLETION_PROMPT, # Default ~100k chars + embedding_text: str = _DEFAULT_EMBEDDING_TEXT, # Default ~100k chars + custom_auth_header: Optional[str] = None, ): """ Initialize the health check client. @@ -41,16 +53,34 @@ class LiteLLMHealthCheckClient: timeout: Request timeout in seconds (default: 120, matching Go implementation) completion_prompt: Test prompt for chat/completion models embedding_text: Test text for embedding models + custom_auth_header: Optional custom header name for authentication (e.g., "x-ifood-requester-service"). + If provided, uses this header instead of standard "Authorization" header. """ self.base_url = base_url.rstrip("/") self.api_key = api_key self.timeout = timeout self.completion_prompt = completion_prompt self.embedding_text = embedding_text - self.headers = { - "Authorization": f"Bearer {api_key}", - "Content-Type": "application/json", - } + + # Debug: Print prompt/text lengths + print(f"DEBUG: Completion prompt length: {len(self.completion_prompt)} characters", file=sys.stderr) + print(f"DEBUG: Embedding text length: {len(self.embedding_text)} characters", file=sys.stderr) + + # Support custom auth header for proxies with custom authentication + # Handle both None and empty string + if custom_auth_header and custom_auth_header.strip(): + custom_auth_header = custom_auth_header.strip() + self.headers = { + custom_auth_header: f"Bearer {api_key}", + "Content-Type": "application/json", + } + print(f"Using custom auth header: {custom_auth_header}", file=sys.stderr) + else: + self.headers = { + "Authorization": f"Bearer {api_key}", + "Content-Type": "application/json", + } + print("Using standard Authorization header", file=sys.stderr) def load_models_from_yaml(self, yaml_path: str) -> List[Dict]: """ @@ -182,6 +212,8 @@ class LiteLLMHealthCheckClient: if is_embedding: # Test embedding endpoint (matching Go implementation) + embedding_text_length = len(self.embedding_text) + print(f"DEBUG: Sending embedding text of length {embedding_text_length} chars to model {model_id}", file=sys.stderr) embedding_response = await client.post( f"{self.base_url}/v1/embeddings", headers=self.headers, @@ -202,6 +234,8 @@ class LiteLLMHealthCheckClient: result["dimensions"] = dimensions else: # Test chat completion endpoint (matching Go implementation) + prompt_length = len(self.completion_prompt) + print(f"DEBUG: Sending prompt of length {prompt_length} chars to model {model_id}", file=sys.stderr) completion_response = await client.post( f"{self.base_url}/v1/chat/completions", headers=self.headers, @@ -358,6 +392,11 @@ async def main(): base_url = os.environ.get("LITELLM_BASE_URL", "http://localhost:4000") api_key = os.environ.get("LITELLM_API_KEY", "sk-1234") yaml_path = os.environ.get("LITELLM_MODELS_YAML") + custom_auth_header = os.environ.get("LITELLM_CUSTOM_AUTH_HEADER") # e.g., "x-ifood-requester-service" + + # Debug: Print custom auth header value if set + if custom_auth_header: + print(f"Custom auth header from env: '{custom_auth_header}'", file=sys.stderr) if not base_url: print("Error: LITELLM_BASE_URL environment variable not set", file=sys.stderr) @@ -369,10 +408,10 @@ async def main(): timeout = int(os.environ.get("LITELLM_TIMEOUT", "120")) # Match Go's 120s default completion_prompt = os.environ.get( - "LITELLM_COMPLETION_PROMPT", "Say this is a test" + "LITELLM_COMPLETION_PROMPT", _DEFAULT_COMPLETION_PROMPT ) embedding_text = os.environ.get( - "LITELLM_EMBEDDING_TEXT", "This is a test for vectorization." + "LITELLM_EMBEDDING_TEXT", _DEFAULT_EMBEDDING_TEXT ) json_output = os.environ.get("LITELLM_JSON_OUTPUT", "").lower() == "true" # Optional: only health-check these model IDs (comma-separated). E.g.: @@ -386,6 +425,7 @@ async def main(): timeout=timeout, completion_prompt=completion_prompt, embedding_text=embedding_text, + custom_auth_header=custom_auth_header, ) # Load models from YAML if provided, otherwise fetch from API diff --git a/scripts/health_check/health_check_client_README.md b/scripts/health_check/health_check_client_README.md index e3132499e0..41c7ab7c35 100644 --- a/scripts/health_check/health_check_client_README.md +++ b/scripts/health_check/health_check_client_README.md @@ -30,6 +30,14 @@ export LITELLM_MODELS_YAML="/path/to/config.yaml" python scripts/health_check/health_check_client.py ``` +**Option 3: Use custom authentication header** +```bash +export LITELLM_BASE_URL="https://litellm.example.com" +export LITELLM_API_KEY="your-api-key" +export LITELLM_CUSTOM_AUTH_HEADER="x-custom-auth-header" +python scripts/health_check/health_check_client.py +``` + ### As a Docker Container 1. Build the Docker image: @@ -47,6 +55,16 @@ docker run --rm \ litellm/litellm-health-check:latest ``` +3. Run with custom authentication header: + +```bash +docker run --rm \ + -e LITELLM_BASE_URL="https://litellm.example.com" \ + -e LITELLM_API_KEY="your-api-key" \ + -e LITELLM_CUSTOM_AUTH_HEADER="x-custom-auth-header" \ + litellm/litellm-health-check:latest +``` + ### Parallel Execution (Stress Testing) Run multiple health check containers in parallel: @@ -65,6 +83,30 @@ export LITELLM_API_KEY="your-api-key" ./scripts/health_check/run_parallel_health_checks.sh 16 ``` +**With Custom Auth Header:** +```powershell +$env:LITELLM_BASE_URL="https://litellm.example.com" +$env:LITELLM_API_KEY="your-api-key" +$env:LITELLM_CUSTOM_AUTH_HEADER="x-custom-auth-header" +.\scripts\health_check\run_parallel_health_checks.ps1 16 +``` + +**With Custom Docker Image:** +```powershell +$env:LITELLM_BASE_URL="https://litellm.example.com" +$env:LITELLM_API_KEY="your-api-key" +$env:LITELLM_CUSTOM_AUTH_HEADER="x-custom-auth-header" +.\scripts\health_check\run_parallel_health_checks.ps1 -NumParallelJobs 16 -ImageName "your-registry/your-image:tag" +``` + +**Bash with Custom Image:** +```bash +export LITELLM_BASE_URL="https://litellm.example.com" +export LITELLM_API_KEY="your-api-key" +export LITELLM_CUSTOM_AUTH_HEADER="x-custom-auth-header" +./scripts/health_check/run_parallel_health_checks.sh 16 "your-registry/your-image:tag" +``` + ## Configuration @@ -73,14 +115,29 @@ export LITELLM_API_KEY="your-api-key" - `LITELLM_BASE_URL` (required): Base URL of the LiteLLM proxy - Example: `https://litellm.example.com` - `LITELLM_API_KEY` (required): API key for authentication +- `LITELLM_CUSTOM_AUTH_HEADER` (optional): Custom header name for authentication + - Use this when your LiteLLM proxy uses a custom authentication header instead of the standard `Authorization` header + - Example: `x-custom-auth-header` (the API key will be sent as `Bearer ` in this header) - `LITELLM_MODELS_YAML` (optional): Path to YAML config file with model_list - If provided, reads models from YAML instead of fetching from API - Example: `/path/to/config.yaml` - `LITELLM_TIMEOUT` (optional): Request timeout in seconds (default: 120) -- `LITELLM_COMPLETION_PROMPT` (optional): Test prompt for chat/completion models (default: "Say this is a test") -- `LITELLM_EMBEDDING_TEXT` (optional): Test text for embedding models (default: "This is a test for vectorization.") +- `LITELLM_COMPLETION_PROMPT` (optional): Test prompt for chat/completion models (default: ~100k characters) +- `LITELLM_EMBEDDING_TEXT` (optional): Test text for embedding models (default: ~100k characters) - `LITELLM_JSON_OUTPUT` (optional): Output results as JSON (default: false) +### Parallel Script Parameters + +**PowerShell (`run_parallel_health_checks.ps1`):** +- `-NumParallelJobs` (optional): Number of parallel containers to run (default: 16) +- `-ImageName` (optional): Docker image to use (default: `litellm/litellm-health-check:latest`) +- `-ContainerRuntime` (optional): Container runtime to use (default: `docker`) + +**Bash (`run_parallel_health_checks.sh`):** +- `[num_parallel_jobs]` (optional): Number of parallel containers to run (default: 16) +- `[image_name]` (optional): Docker image to use (default: `litellm/litellm-health-check:latest`) +- `[container_runtime]` (optional): Container runtime to use (default: `docker`) + ## Output ### Standard Output (Human-Readable) @@ -166,7 +223,20 @@ Run multiple health checks in parallel: **PowerShell:** ```powershell +# Using default image .\scripts\health_check\run_parallel_health_checks.ps1 16 + +# Using custom image +.\scripts\health_check\run_parallel_health_checks.ps1 -NumParallelJobs 16 -ImageName "your-registry/your-image:tag" +``` + +**Bash:** +```bash +# Using default image +./scripts/health_check/run_parallel_health_checks.sh 16 + +# Using custom image +./scripts/health_check/run_parallel_health_checks.sh 16 "your-registry/your-image:tag" ``` ### 3. CI/CD Integration diff --git a/scripts/health_check/run_parallel_health_checks.ps1 b/scripts/health_check/run_parallel_health_checks.ps1 index 856e7f20ec..850fc6baab 100644 --- a/scripts/health_check/run_parallel_health_checks.ps1 +++ b/scripts/health_check/run_parallel_health_checks.ps1 @@ -49,6 +49,11 @@ Write-Host " 2. Set LITELLM_BASE_URL to the correct URL (e.g., http://host.do Write-Host " 3. On Linux, you may need to use the host IP instead of host.docker.internal" -ForegroundColor Yellow Write-Host "" +# Capture environment variables in parent scope for use in parallel block +$baseUrl = $env:LITELLM_BASE_URL +$apiKey = $env:LITELLM_API_KEY +$customAuthHeader = $env:LITELLM_CUSTOM_AUTH_HEADER + # Run parallel health checks # This creates an infinite loop that keeps spawning containers # Each container tests all models, then exits, and a new one starts @@ -57,13 +62,20 @@ while ($true) { 1..$NumParallelJobs | ForEach-Object -Parallel { $runtime = $using:ContainerRuntime $imageName = $using:ImageName - $baseUrl = $env:LITELLM_BASE_URL - $apiKey = $env:LITELLM_API_KEY + $baseUrl = $using:baseUrl + $apiKey = $using:apiKey + $customAuthHeader = $using:customAuthHeader - & $runtime run --rm ` - -e LITELLM_BASE_URL="$baseUrl" ` - -e LITELLM_API_KEY="$apiKey" ` - -e LITELLM_JSON_OUTPUT="true" ` - $imageName + $envVars = @( + "-e", "LITELLM_BASE_URL=$baseUrl", + "-e", "LITELLM_API_KEY=$apiKey", + "-e", "LITELLM_JSON_OUTPUT=true" + ) + + if ($customAuthHeader) { + $envVars += "-e", "LITELLM_CUSTOM_AUTH_HEADER=$customAuthHeader" + } + + & $runtime run --rm $envVars $imageName } -ThrottleLimit $NumParallelJobs } diff --git a/scripts/health_check/run_parallel_health_checks.sh b/scripts/health_check/run_parallel_health_checks.sh index 9b6c5d9f39..d1913e3fdd 100644 --- a/scripts/health_check/run_parallel_health_checks.sh +++ b/scripts/health_check/run_parallel_health_checks.sh @@ -54,11 +54,18 @@ echo "" # Function to run a single health check container run_health_check() { - "$CONTAINER_RUNTIME" run --rm \ - -e LITELLM_BASE_URL="$LITELLM_BASE_URL" \ - -e LITELLM_API_KEY="$LITELLM_API_KEY" \ - -e LITELLM_JSON_OUTPUT="true" \ - "$IMAGE_NAME" + local env_vars=( + -e "LITELLM_BASE_URL=$LITELLM_BASE_URL" + -e "LITELLM_API_KEY=$LITELLM_API_KEY" + -e "LITELLM_JSON_OUTPUT=true" + ) + + # Pass through custom auth header if set + if [ -n "$LITELLM_CUSTOM_AUTH_HEADER" ]; then + env_vars+=(-e "LITELLM_CUSTOM_AUTH_HEADER=$LITELLM_CUSTOM_AUTH_HEADER") + fi + + "$CONTAINER_RUNTIME" run --rm "${env_vars[@]}" "$IMAGE_NAME" } # Run parallel health checks