backend/tests/test_provider_rate_limits.py

#!/usr/bin/env python
"""
Script to test rate limits of Hugging Face Inference API providers.
Spams requests to a model/provider and collects error messages.

Usage: python test_provider_rate_limits.py --model "model_name" --provider "provider_name" --requests 50
"""

import argparse
import json
import time
import os
import requests
import sys
import logging
from concurrent.futures import ThreadPoolExecutor
from collections import Counter
from typing import Dict, List, Tuple
from dotenv import load_dotenv

# Add parent directory to path to import from tasks
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from tasks.get_available_model_provider import prioritize_providers

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger("rate_limit_test")

# Default model to test
DEFAULT_MODEL = "meta-llama/Llama-3.3-70B-Instruct"

def send_request(model: str, provider: str, token: str, request_id: int) -> Dict:
    """
    Send a single request to the model with the given provider.
    
    Args:
        model: Model name
        provider: Provider name
        token: HF token
        request_id: ID for this request
        
    Returns:
        Dictionary with request info and result
    """
    headers = {
        "Authorization": f"Bearer {token}",
        "Content-Type": "application/json"
    }
    
    payload = {
        "inputs": f"Request {request_id}: Hello, what do you thing about the future of AI? And divide me 10 by {request_id}",
        "parameters": {
            "max_new_tokens": 10000,
            "provider": provider
        }
    }
    
    api_url = f"https://api-inference.huggingface.co/models/{model}"
    
    start_time = time.time()
    try:
        response = requests.post(api_url, headers=headers, json=payload, timeout=15)
        end_time = time.time()
        
        result = {
            "request_id": request_id,
            "status_code": response.status_code,
            "time_taken": end_time - start_time,
            "headers": dict(response.headers),
            "success": response.status_code == 200,
        }
        
        if response.status_code != 200:
            try:
                error_data = response.json()
                if isinstance(error_data, dict) and "error" in error_data:
                    result["error_message"] = error_data["error"]
                else:
                    result["error_message"] = str(error_data)
            except:
                result["error_message"] = response.text
        
        return result
    
    except Exception as e:
        end_time = time.time()
        return {
            "request_id": request_id,
            "status_code": 0,
            "time_taken": end_time - start_time,
            "success": False,
            "error_message": str(e)
        }

def run_rate_limit_test(model: str, provider: str = None, num_requests: int = 50, 
                       max_workers: int = 10, delay: float = 0.1) -> List[Dict]:
    """
    Run a rate limit test by sending multiple requests to the specified model/provider.
    
    Args:
        model: Model to test
        provider: Provider to test (if None, will use first available)
        num_requests: Number of requests to send
        max_workers: Maximum number of concurrent workers
        delay: Delay between batches of requests
        
    Returns:
        List of results for each request
    """
    # Load environment variables
    load_dotenv()
    
    # Get HF token
    hf_token = os.environ.get("HF_TOKEN")
    if not hf_token:
        logger.error("HF_TOKEN not defined in environment")
        return []
    
    # If provider not specified, get first available
    if not provider:
        from tasks.get_available_model_provider import get_available_model_provider
        provider = get_available_model_provider(model)
        if not provider:
            logger.error(f"No available provider found for {model}")
            return []
    
    logger.info(f"Testing rate limits for {model} with provider: {provider}")
    logger.info(f"Sending {num_requests} requests with {max_workers} concurrent workers")
    
    # Send requests in parallel
    results = []
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        future_to_id = {
            executor.submit(send_request, model, provider, hf_token, i): i 
            for i in range(num_requests)
        }
        
        completed = 0
        for future in future_to_id:
            result = future.result()
            results.append(result)
            
            completed += 1
            if completed % 10 == 0:
                logger.info(f"Completed {completed}/{num_requests} requests")
            
            # Add a small delay periodically to avoid overwhelming the API
            if completed % max_workers == 0:
                time.sleep(delay)
    
    return results

def analyze_results(results: List[Dict]) -> Dict:
    """
    Analyze the results of the rate limit test.
    
    Args:
        results: List of request results
        
    Returns:
        Dictionary with analysis
    """
    total_requests = len(results)
    successful = sum(1 for r in results if r["success"])
    failed = total_requests - successful
    
    # Count different error messages
    error_messages = Counter(r.get("error_message") for r in results if not r["success"])
    
    # Calculate timing statistics
    times = [r["time_taken"] for r in results]
    avg_time = sum(times) / len(times) if times else 0
    
    # Check for rate limiting headers
    rate_limit_headers = set()
    for r in results:
        if "headers" in r:
            for header in r["headers"]:
                if "rate" in header.lower() or "limit" in header.lower():
                    rate_limit_headers.add(header)
    
    return {
        "total_requests": total_requests,
        "successful_requests": successful,
        "failed_requests": failed,
        "success_rate": successful / total_requests if total_requests > 0 else 0,
        "average_time": avg_time,
        "error_messages": dict(error_messages),
        "rate_limit_headers": list(rate_limit_headers)
    }

def display_results(results: List[Dict], analysis: Dict) -> None:
    """
    Display the results of the rate limit test.
    
    Args:
        results: List of request results
        analysis: Analysis of results
    """
    print("\n" + "="*80)
    print(f"RATE LIMIT TEST RESULTS")
    print("="*80)
    
    print(f"\nTotal Requests: {analysis['total_requests']}")
    print(f"Successful: {analysis['successful_requests']} ({analysis['success_rate']*100:.1f}%)")
    print(f"Failed: {analysis['failed_requests']}")
    print(f"Average Time: {analysis['average_time']:.3f} seconds")
    
    if analysis["rate_limit_headers"]:
        print("\nRate Limit Headers Found:")
        for header in analysis["rate_limit_headers"]:
            print(f"  - {header}")
    
    if analysis["error_messages"]:
        print("\nError Messages:")
        for msg, count in analysis["error_messages"].items():
            print(f"  - [{count} occurrences] {msg}")
    
    # Print sample of headers from a failed request
    failed_requests = [r for r in results if not r["success"]]
    if failed_requests:
        print("\nSample Headers from a Failed Request:")
        for header, value in failed_requests[0].get("headers", {}).items():
            print(f"  {header}: {value}")

def main():
    """
    Main entry point for the script.
    """
    parser = argparse.ArgumentParser(description="Test rate limits of Hugging Face Inference API providers.")
    parser.add_argument("--model", type=str, default=DEFAULT_MODEL, help="Name of the model to test")
    parser.add_argument("--provider", type=str, help="Name of the provider to test (if not specified, will use first available)")
    parser.add_argument("--requests", type=int, default=50, help="Number of requests to send")
    parser.add_argument("--workers", type=int, default=10, help="Maximum number of concurrent workers")
    parser.add_argument("--delay", type=float, default=0.1, help="Delay between batches of requests")
    parser.add_argument("--output", type=str, help="Path to save results as JSON (optional)")
    
    args = parser.parse_args()
    
    # Run the test
    results = run_rate_limit_test(
        model=args.model,
        provider=args.provider,
        num_requests=args.requests,
        max_workers=args.workers,
        delay=args.delay
    )
    
    if not results:
        logger.error("Test failed to run properly")
        return
    
    # Analyze the results
    analysis = analyze_results(results)
    
    # Display the results
    display_results(results, analysis)
    
    # Save results if requested
    if args.output:
        with open(args.output, "w") as f:
            json.dump({
                "results": results,
                "analysis": analysis
            }, f, indent=2)
        logger.info(f"Results saved to {args.output}")

if __name__ == "__main__":
    main()