#!/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()