This commit is contained in:
bennykok 2024-09-18 08:20:35 -07:00
parent aa47f3523f
commit 61acffd355

View File

@ -24,7 +24,7 @@ from typing import Dict, List, Union, Any, Optional
from PIL import Image from PIL import Image
import copy import copy
import struct import struct
from aiohttp import web, ClientSession, ClientError, ClientTimeout from aiohttp import web, ClientSession, ClientError, ClientTimeout, ClientResponseError
import atexit import atexit
# Global session # Global session
@ -1209,6 +1209,37 @@ async def file_sender(file_object, chunk_size):
chunk_size = 1024 * 1024 # 1MB chunks, adjust as needed chunk_size = 1024 * 1024 # 1MB chunks, adjust as needed
async def upload_with_retry(session, url, headers, data, max_retries=3, initial_delay=1):
start_time = time.time() # Start timing here
for attempt in range(max_retries):
try:
async with session.put(url, headers=headers, data=data) as response:
upload_duration = time.time() - start_time
logger.info(f"Upload attempt {attempt + 1} completed in {upload_duration:.2f} seconds")
logger.info(f"Upload response status: {response.status}")
response.raise_for_status() # This will raise an exception for 4xx and 5xx status codes
response_text = await response.text()
logger.info(f"Response body: {response_text[:1000]}...")
logger.info("Upload successful")
return response # Successful upload, exit the retry loop
except (ClientError, ClientResponseError) as e:
logger.error(f"Upload attempt {attempt + 1} failed: {str(e)}")
if attempt < max_retries - 1: # If it's not the last attempt
delay = initial_delay * (2 ** attempt) # Exponential backoff
logger.info(f"Retrying in {delay} seconds...")
await asyncio.sleep(delay)
else:
logger.error("Max retries reached. Upload failed.")
raise # Re-raise the last exception if all retries are exhausted
except Exception as e:
logger.error(f"Unexpected error during upload: {str(e)}")
logger.error(traceback.format_exc())
raise # Re-raise unexpected exceptions immediately
async def upload_file(prompt_id, filename, subfolder=None, content_type="image/png", type="output", item=None): async def upload_file(prompt_id, filename, subfolder=None, content_type="image/png", type="output", item=None):
""" """
Uploads file to S3 bucket using S3 client object Uploads file to S3 bucket using S3 client object
@ -1277,29 +1308,11 @@ async def upload_file(prompt_id, filename, subfolder=None, content_type="image/p
async with aiohttp.ClientSession() as session: async with aiohttp.ClientSession() as session:
try: try:
async with session.put(ok.get("url"), headers=headers, data=data) as response: response = await upload_with_retry(session, ok.get("url"), headers, data)
upload_duration = time.time() - start_time # Process successful response...
logger.info(f"Upload completed in {upload_duration:.2f} seconds")
logger.info(f"Upload response status: {response.status}")
logger.info(f"Upload response reason: {response.reason}")
response_headers = response.headers
logger.info(f"Response headers: {dict(response_headers)}")
response_text = await response.text()
logger.info(f"Response body: {response_text[:1000]}...") # Log first 1000 characters of response body
if response.status not in [200, 201, 204]:
logger.error(f"Upload failed with status {response.status}")
logger.error(f"Full response body: {response_text}")
else:
logger.info("Upload successful")
except aiohttp.ClientError as e:
logger.error(f"Client error during upload: {str(e)}")
except Exception as e: except Exception as e:
logger.error(f"Unexpected error during upload: {str(e)}") # Handle final failure...
logger.error(traceback.format_exc()) logger.error(f"Upload ultimately failed: {str(e)}")
end_time = time.time() # End timing after the request is complete end_time = time.time() # End timing after the request is complete
logger.info("Upload time: {:.2f} seconds".format(end_time - start_time)) logger.info("Upload time: {:.2f} seconds".format(end_time - start_time))