File size: 6,210 Bytes
7c4d825 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 |
#!/usr/bin/env python3
# ----------------------------------------------------------------------
# Test script for GPU quota handling
# ----------------------------------------------------------------------
import requests
import json
import time
import sys
# ----------------------------------------------------------------------
# Configuration
# ----------------------------------------------------------------------
BASE_URL = "http://localhost:7860"
TEST_IMAGE_URL = "https://cdn.shopify.com/s/files/1/0505/0928/3527/files/hugging_face_test_image_shirt_product_type.jpg"
# ----------------------------------------------------------------------
# Test Functions
# ----------------------------------------------------------------------
def test_quota_info():
"""Test the quota info endpoint"""
print("\n=== Testing /api/quota-info ===")
response = requests.get(f"{BASE_URL}/api/quota-info")
print(f"Status Code: {response.status_code}")
if response.status_code == 200:
data = response.json()
print(f"Quota Management: {data.get('quota_management')}")
print(f"User Type Quotas:")
for user_type, info in data.get('user_type_quotas', {}).items():
print(f" - {user_type}: {info.get('quota_seconds')}s ({info.get('description')})")
if 'last_quota_error' in data:
print(f"\nLast Quota Error:")
error_info = data['last_quota_error']
print(f" - Time ago: {error_info.get('time_ago_seconds')}s")
print(f" - Retry after: {error_info.get('retry_after')}s")
print(f" - Estimated recovery: {error_info.get('estimated_recovery')}s")
if 'usage_stats' in data:
print(f"\nUsage Stats (Last Hour):")
stats = data['usage_stats']['last_hour']
print(f" - Total requests: {stats.get('total_requests')}")
print(f" - Successful: {stats.get('successful')}")
print(f" - Failed: {stats.get('failed')}")
print(f" - GPU seconds used: {stats.get('total_gpu_seconds', 0):.2f}s")
print(f" - Average duration: {stats.get('average_duration', 0):.2f}s")
def test_process_image():
"""Test image processing and quota error handling"""
print("\n=== Testing Image Processing ===")
payload = {
"data": [
[{"url": TEST_IMAGE_URL}],
"Shirt"
]
}
response = requests.post(
f"{BASE_URL}/api/rb_and_crop",
json=payload,
timeout=180
)
print(f"Status Code: {response.status_code}")
# Check if we got a quota error
if response.status_code == 429:
print("Got expected 429 (Too Many Requests) status for quota exceeded")
# Check for Retry-After header
retry_after = response.headers.get('Retry-After')
if retry_after:
print(f"Retry-After header: {retry_after}s")
else:
print("WARNING: No Retry-After header found")
# Parse error response
data = response.json()
print(f"Error Type: {data.get('error_type')}")
print(f"Error Message: {data.get('error_message')[:100]}...")
error_details = data.get('error_details', {})
if 'retry_after' in error_details:
print(f"Retry After (from body): {error_details['retry_after']}s")
if 'quota_info' in error_details:
quota_info = error_details['quota_info']
print(f"\nQuota Info:")
print(f" - Message: {quota_info.get('message')}")
print(f" - Calculated retry: {quota_info.get('calculated_retry')}s")
elif response.status_code == 503:
print("Got 503 (Service Unavailable) - GPU warming up")
data = response.json()
print(f"Details: {data.get('detail')}")
elif response.status_code == 200:
print("Success! Image processed")
data = response.json()
if 'processed_images' in data:
for img in data['processed_images']:
print(f" - URL: {img.get('url')}")
print(f" - Status: {img.get('status')}")
if img.get('status') == 'error':
print(f" - Error: {img.get('error')}")
else:
print(f"Unexpected status code: {response.status_code}")
try:
print(f"Response: {response.json()}")
except:
print(f"Response text: {response.text[:200]}")
def test_simulate_quota_error():
"""Simulate multiple requests to trigger quota error"""
print("\n=== Simulating Quota Exhaustion ===")
print("This will make multiple requests to exhaust GPU quota...")
for i in range(3):
print(f"\nAttempt {i + 1}:")
test_process_image()
# Check quota status after each attempt
time.sleep(2)
test_quota_info()
if i < 2:
print("\nWaiting 5 seconds before next attempt...")
time.sleep(5)
# ----------------------------------------------------------------------
# Main
# ----------------------------------------------------------------------
if __name__ == "__main__":
print("GPU Quota Handling Test")
print("=======================")
# Check if service is running
try:
response = requests.get(f"{BASE_URL}/health")
if response.status_code == 200:
health = response.json()
print(f"Service is healthy")
print(f"Device: {health.get('device')}")
print(f"Models loaded: {health.get('models_loaded')}")
print(f"GPU available: {health.get('gpu_available')}")
else:
print("Service health check failed")
sys.exit(1)
except Exception as e:
print(f"Cannot connect to service at {BASE_URL}: {e}")
print("Make sure the service is running!")
sys.exit(1)
# Run tests
if len(sys.argv) > 1 and sys.argv[1] == "exhaust":
test_simulate_quota_error()
else:
test_quota_info()
test_process_image()
print("\nTest completed!")
|