also download to a tmp file not the live file

to avoid compounding errors from network failure
This commit is contained in:
Alex "mcmonkey" Goodwin 2024-09-19 20:06:53 +09:00
parent c539ee55f4
commit cfb5d6a532

View File

@ -158,7 +158,10 @@ async def track_download_progress(response: aiohttp.ClientResponse,
await progress_callback(model_name, status)
last_update_time = time.time()
with open(file_path, 'wb') as f:
if os.path.exists(file_path + '.tmp'):
os.remove(file_path + '.tmp')
with open(file_path + '.tmp', 'wb') as f:
chunk_iterator = response.content.iter_chunked(8192)
while True:
try:
@ -171,6 +174,8 @@ async def track_download_progress(response: aiohttp.ClientResponse,
if time.time() - last_update_time >= interval:
await update_progress()
os.rename(file_path + '.tmp', file_path)
await update_progress()
logging.info(f"Successfully downloaded {model_name}. Total downloaded: {downloaded}")