Downloader: Make timeout configurable

Add an API parameter to set the timeout in seconds. Keep it to None
by default for uninterrupted downloads.

Signed-off-by: kingbri <bdashore3@proton.me>
This commit is contained in:
kingbri 2024-07-23 21:42:38 -04:00
parent 8c02fe9771
commit 71de3060bb
2 changed files with 4 additions and 2 deletions

View file

@ -101,6 +101,7 @@ async def hf_repo_download(
chunk_limit: Optional[float],
include: Optional[List[str]],
exclude: Optional[List[str]],
timeout: Optional[int],
repo_type: Optional[str] = "model",
):
"""Gets a repo's information from HuggingFace and downloads it locally."""
@ -145,8 +146,8 @@ async def hf_repo_download(
logger.info(f"Saving {repo_id} to {str(download_path)}")
try:
timeout = aiohttp.ClientTimeout(total=None) # Turn off timeout
async with aiohttp.ClientSession(timeout=timeout) as session:
client_timeout = aiohttp.ClientTimeout(total=timeout) # Turn off timeout
async with aiohttp.ClientSession(timeout=client_timeout) as session:
tasks = []
logger.info(f"Starting download for {repo_id}")