Fairly often tests are failing due to connection issues when downloading models. Proposal to add retry logic to the download function so that if there's a socket hang up etc. it retries up to three times.
|
def download(url: str, path: Path, md5_hash: str) -> None: |
|
""" |
|
Downloads a file from a URL to a path. |
|
""" |
|
|
|
if path.exists() and (md5_hash == "" or md5_hash == hash_file(path)): |
|
return |
|
|
|
path.parent.mkdir(parents=True, exist_ok=True) |
|
|
|
resp = get(url, stream=True) |
|
total = int(resp.headers.get("content-length", 0)) |
|
|
|
with path.open("wb") as file, tqdm( |
|
desc="Downloading Model", |
|
total=total, |
|
unit="iB", |
|
unit_scale=True, |
|
unit_divisor=1024, |
|
) as progress_bar: |
|
for data in resp.iter_content(chunk_size=1024): |
|
size = file.write(data) |
|
progress_bar.update(size) |
|
|
|
if md5_hash != "": |
|
assert md5_hash == hash_file(path) |
