Skip to content

Fix progress bar not always closed in file_download.py #2308

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Jun 3, 2024
Merged

Fix progress bar not always closed in file_download.py #2308

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
1ecf5d3
Fix progress bar not always closed in file_download.py
Wauplin Jun 3, 2024
63dfdaf
Fix feedback from charles
Wauplin Jun 3, 2024
5b54b42
Merge branch 'main' into fix-progress-bar-not-closed-when-downloading
Wauplin Jun 3, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
133 changes: 69 additions & 64 deletions src/huggingface_hub/file_download.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import contextlib
import copy
import errno
import fnmatch
Expand Down Expand Up @@ -487,9 +488,8 @@ def http_get(
)

# Stream file to buffer
progress = _tqdm_bar
if progress is None:
progress = tqdm(
progress_cm: tqdm = (
tqdm( # type: ignore[assignment]
unit="B",
unit_scale=True,
total=total,
Expand All @@ -500,71 +500,76 @@ def http_get(
# see https://github.com/huggingface/huggingface_hub/pull/2000
name="huggingface_hub.http_get",
)
if _tqdm_bar is None
else contextlib.nullcontext(_tqdm_bar)
# ^ `contextlib.nullcontext` mimics a context manager that does nothing
# Makes it easier to use the same code path for both cases but in the later
# case, the progress bar is not closed when exiting the context manager.
)

if hf_transfer and total is not None and total > 5 * DOWNLOAD_CHUNK_SIZE:
supports_callback = "callback" in inspect.signature(hf_transfer.download).parameters
if not supports_callback:
warnings.warn(
"You are using an outdated version of `hf_transfer`. "
"Consider upgrading to latest version to enable progress bars "
"using `pip install -U hf_transfer`."
)
with progress_cm as progress:
if hf_transfer and total is not None and total > 5 * DOWNLOAD_CHUNK_SIZE:
supports_callback = "callback" in inspect.signature(hf_transfer.download).parameters
if not supports_callback:
warnings.warn(
"You are using an outdated version of `hf_transfer`. "
"Consider upgrading to latest version to enable progress bars "
"using `pip install -U hf_transfer`."
)
try:
hf_transfer.download(
url=url,
filename=temp_file.name,
max_files=HF_TRANSFER_CONCURRENCY,
chunk_size=DOWNLOAD_CHUNK_SIZE,
headers=headers,
parallel_failures=3,
max_retries=5,
**({"callback": progress.update} if supports_callback else {}),
)
except Exception as e:
raise RuntimeError(
"An error occurred while downloading using `hf_transfer`. Consider"
" disabling HF_HUB_ENABLE_HF_TRANSFER for better error handling."
) from e
if not supports_callback:
progress.update(total)
if expected_size is not None and expected_size != os.path.getsize(temp_file.name):
raise EnvironmentError(
consistency_error_message.format(
actual_size=os.path.getsize(temp_file.name),
)
)
return
new_resume_size = resume_size
try:
hf_transfer.download(
for chunk in r.iter_content(chunk_size=DOWNLOAD_CHUNK_SIZE):
if chunk: # filter out keep-alive new chunks
progress.update(len(chunk))
temp_file.write(chunk)
new_resume_size += len(chunk)
# Some data has been downloaded from the server so we reset the number of retries.
_nb_retries = 5
except (requests.ConnectionError, requests.ReadTimeout) as e:
# If ConnectionError (SSLError) or ReadTimeout happen while streaming data from the server, it is most likely
# a transient error (network outage?). We log a warning message and try to resume the download a few times
# before giving up. Tre retry mechanism is basic but should be enough in most cases.
if _nb_retries <= 0:
logger.warning("Error while downloading from %s: %s\nMax retries exceeded.", url, str(e))
raise
logger.warning("Error while downloading from %s: %s\nTrying to resume download...", url, str(e))
time.sleep(1)
reset_sessions() # In case of SSLError it's best to reset the shared requests.Session objects
return http_get(
url=url,
filename=temp_file.name,
max_files=HF_TRANSFER_CONCURRENCY,
chunk_size=DOWNLOAD_CHUNK_SIZE,
headers=headers,
parallel_failures=3,
max_retries=5,
**({"callback": progress.update} if supports_callback else {}),
)
except Exception as e:
raise RuntimeError(
"An error occurred while downloading using `hf_transfer`. Consider"
" disabling HF_HUB_ENABLE_HF_TRANSFER for better error handling."
) from e
if not supports_callback:
progress.update(total)
if expected_size is not None and expected_size != os.path.getsize(temp_file.name):
raise EnvironmentError(
consistency_error_message.format(
actual_size=os.path.getsize(temp_file.name),
)
temp_file=temp_file,
proxies=proxies,
resume_size=new_resume_size,
headers=initial_headers,
expected_size=expected_size,
_nb_retries=_nb_retries - 1,
_tqdm_bar=_tqdm_bar,
)
return
new_resume_size = resume_size
try:
for chunk in r.iter_content(chunk_size=DOWNLOAD_CHUNK_SIZE):
if chunk: # filter out keep-alive new chunks
progress.update(len(chunk))
temp_file.write(chunk)
new_resume_size += len(chunk)
# Some data has been downloaded from the server so we reset the number of retries.
_nb_retries = 5
except (requests.ConnectionError, requests.ReadTimeout) as e:
# If ConnectionError (SSLError) or ReadTimeout happen while streaming data from the server, it is most likely
# a transient error (network outage?). We log a warning message and try to resume the download a few times
# before giving up. Tre retry mechanism is basic but should be enough in most cases.
if _nb_retries <= 0:
logger.warning("Error while downloading from %s: %s\nMax retries exceeded.", url, str(e))
raise
logger.warning("Error while downloading from %s: %s\nTrying to resume download...", url, str(e))
time.sleep(1)
reset_sessions() # In case of SSLError it's best to reset the shared requests.Session objects
return http_get(
url=url,
temp_file=temp_file,
proxies=proxies,
resume_size=new_resume_size,
headers=initial_headers,
expected_size=expected_size,
_nb_retries=_nb_retries - 1,
_tqdm_bar=_tqdm_bar,
)

progress.close()

if expected_size is not None and expected_size != temp_file.tell():
raise EnvironmentError(
Expand Down