← back to Atmos Game
assets/gen/.venv/lib/python3.12/site-packages/huggingface_hub/utils/_xet_progress_reporting.py
354 lines
from collections import OrderedDict
from typing import Any
from . import is_google_colab, is_notebook
from .tqdm import _create_progress_bar, tqdm
def _format_speed_postfix(speed: float | None) -> str:
s = tqdm.format_sizeof(speed) if speed is not None else "???"
return f"{s}B/s ".rjust(10, " ")
# Transfer byte count is hard to predict (dedup/compression), so we omit a total and show bytes only.
XET_TRANSFER_BAR_FORMAT = "{desc}: {bar}| {n_fmt:>5}B{postfix:>12}"
XET_BYTES_BAR_FORMAT = "{l_bar}{bar}| {n_fmt:>5}B / {total_fmt:>5}B{postfix:>12}"
def _set_monotonic_total(bar, total: int | None) -> None:
if total is None or not hasattr(bar, "total"):
return
bar.total = max(bar.total or 0, total)
def _update_transfer_bar(bar, inc: int) -> None:
"""Update the transfer bar and grow its hidden total so the bar graphic advances.
Network bytes are hard to predict (dedup/compression), so the display omits a denominator.
tqdm still needs an internal total for the bar width — seeded from file size when known,
then expanded here if bytes received exceed that estimate.
"""
n_after = getattr(bar, "n", 0) + inc
current_total = getattr(bar, "total", 0) or 0
if n_after > 0 and current_total < n_after:
bar.total = max(current_total, int(n_after * 1.25) + 1)
bar.update(inc)
def _finish_transfer_bar(bar) -> None:
"""Snap the transfer bar to 100% when downloading stops.
Transfer totals are seeded from file size, but actual network bytes are often lower.
Set ``total = n`` so the bar fills completely instead of stopping partway.
"""
n = getattr(bar, "n", 0)
if n > 0 and hasattr(bar, "total") and bar.total != n:
bar.total = n
bar.refresh()
def _set_aggregate_rate_postfix(bar) -> None:
"""Show a shared bar's own throughput as its rate.
When many files feed one bar (snapshot download), each file only knows its own per-item rate — a
fraction of the total. Deriving the rate from the shared bar's aggregated byte count reports the
true combined speed instead of whichever single file reported last.
"""
bar.set_postfix_str(_format_speed_postfix(bar.format_dict.get("rate")), refresh=False)
class XetDownloadProgressReporter:
"""Dual progress bars for Xet downloads: network transfer and file reconstruction.
``total_transfer_bytes_completed`` tracks bytes received from the network (updated continuously).
``total_bytes_completed`` tracks bytes written to disk (updated after buffered chunks are flushed).
Showing both bars gives responsive feedback on slow connections where reconstruction lags behind transfer.
"""
def __init__(
self,
*,
reconstruction_desc: str,
transfer_desc: str = "Downloading bytes",
total: int | None = None,
log_level: int,
name: str | None = None,
tqdm_class: type | None = None,
external_reconstruction_bar: Any | None = None,
position: int = 0,
):
self._prev_bytes_completed = 0
self._prev_transfer_bytes_completed = 0
cls = tqdm_class or tqdm
routes_transfer_via_reconstruction = external_reconstruction_bar is not None and callable(
getattr(external_reconstruction_bar, "update_transfer", None)
)
uses_aggregated_tqdm_class = external_reconstruction_bar is None and callable(
getattr(cls, "update_transfer", None)
)
if external_reconstruction_bar is not None:
self.reconstruction_bar = external_reconstruction_bar
self._owns_reconstruction_bar = False
else:
self.reconstruction_bar = _create_progress_bar(
cls=cls, # ty: ignore[invalid-argument-type]
log_level=log_level,
name=name,
desc=reconstruction_desc,
total=total,
unit="B",
unit_scale=True,
position=position + 1,
bar_format=XET_BYTES_BAR_FORMAT,
leave=True,
)
self._owns_reconstruction_bar = True
if routes_transfer_via_reconstruction or uses_aggregated_tqdm_class:
self.transfer_bar = self.reconstruction_bar
self._owns_transfer_bar = False
elif external_reconstruction_bar is not None:
self.transfer_bar = None
self._owns_transfer_bar = False
else:
self.transfer_bar = _create_progress_bar(
cls=tqdm,
log_level=log_level,
name=f"{name}.transfer" if name else None,
desc=transfer_desc,
total=total,
unit="B",
unit_scale=True,
position=position,
bar_format=XET_TRANSFER_BAR_FORMAT,
leave=True,
)
self._owns_transfer_bar = True
@property
def _aggregated(self) -> bool:
return self.transfer_bar is not None and self.transfer_bar is self.reconstruction_bar
def update_progress(self, group_report, _item_reports: dict | None = None) -> None:
bytes_inc = max(0, group_report.total_bytes_completed - self._prev_bytes_completed)
transfer_inc = max(0, group_report.total_transfer_bytes_completed - self._prev_transfer_bytes_completed)
self._prev_bytes_completed = group_report.total_bytes_completed
self._prev_transfer_bytes_completed = group_report.total_transfer_bytes_completed
if bytes_inc > 0:
self.reconstruction_bar.update(bytes_inc)
self.reconstruction_bar.set_postfix_str(
_format_speed_postfix(group_report.total_bytes_completion_rate), refresh=False
)
if transfer_inc > 0 and self.transfer_bar is not None:
if self._aggregated:
self.reconstruction_bar.update_transfer(transfer_inc)
self.reconstruction_bar.set_transfer_postfix_str(
_format_speed_postfix(group_report.total_transfer_bytes_completion_rate), refresh=False
)
else:
_update_transfer_bar(self.transfer_bar, transfer_inc)
self.transfer_bar.set_postfix_str(
_format_speed_postfix(group_report.total_transfer_bytes_completion_rate), refresh=False
)
if group_report.total_bytes:
_set_monotonic_total(self.reconstruction_bar, group_report.total_bytes)
def close(self) -> None:
"""Close bars owned by this reporter.
Standalone downloads finish the transfer bar first (snap hidden total to ``n``), then close it.
Aggregated and external bars (e.g. snapshot ``_AggregatedTqdm`` or a reused ``_tqdm_bar``) are left
open — their parent caller owns their lifecycle.
"""
if self.transfer_bar is not None and self._owns_transfer_bar:
_finish_transfer_bar(self.transfer_bar)
if hasattr(self.transfer_bar, "close"):
self.transfer_bar.close()
if self._owns_reconstruction_bar and hasattr(self.reconstruction_bar, "close"):
self.reconstruction_bar.close()
def __enter__(self) -> "XetDownloadProgressReporter":
return self
def __exit__(self, *args) -> None:
self.close()
class XetUploadProgressReporter:
"""
Reports on progress for Xet uploads.
Shows summary progress bars when running in notebooks or GUIs, and detailed per-file progress in console environments.
"""
def __init__(self, n_lines: int = 10, description_width: int = 30, total_files: int | None = None):
self.n_lines = n_lines
self.description_width = description_width
self.total_files = total_files
self.per_file_progress = is_google_colab() or not is_notebook()
self.tqdm_settings = {
"unit": "B",
"unit_scale": True,
"leave": True,
"unit_divisor": 1000,
"nrows": n_lines + 3 if self.per_file_progress else 3,
"miniters": 1,
"bar_format": XET_BYTES_BAR_FORMAT,
}
# Overall progress bars
self.data_processing_bar = tqdm(
total=0, desc=self.format_desc("Processing Files (0 / 0)", False), position=0, **self.tqdm_settings
)
self.upload_bar = tqdm(
total=0, desc=self.format_desc("New Data Upload", False), position=1, **self.tqdm_settings
)
self.known_items: set[str] = set()
self.completed_items: set[str] = set()
# Track previous absolute values to compute increments
self._prev_bytes_completed: int = 0
self._prev_transfer_bytes_completed: int = 0
# Item bars (scrolling view)
self.item_state: OrderedDict[str, Any] = OrderedDict()
self.current_bars: list = [None] * self.n_lines
def format_desc(self, name: str, indent: bool) -> str:
"""
if name is longer than width characters, prints ... at the start and then the last width-3 characters of the name, otherwise
the whole name right justified into description_width characters. Also adds some padding.
"""
if not self.per_file_progress:
# Here we just use the defaults.
return name
padding = " " if indent else ""
width = self.description_width - len(padding)
if len(name) > width:
name = f"...{name[-(width - 3) :]}"
return f"{padding}{name.ljust(width)}"
def reset_for_next_commit(self):
"""Reset per-commit state so the reporter can be reused across multiple upload commits."""
self._prev_bytes_completed = 0
self._prev_transfer_bytes_completed = 0
self.known_items.clear()
self.completed_items.clear()
self.item_state.clear()
def update_progress(self, group_report, item_reports: dict):
# Update all the per-item values.
for item in item_reports.values():
item_name = item.item_name
self.known_items.add(item_name)
# Only care about items where the processing has already started.
if item.bytes_completed == 0:
continue
# Overwrite the existing value in there.
self.item_state[item_name] = item
bar_idx = 0
new_completed = []
# Now, go through and update all the bars
for name, item in self.item_state.items():
# Is this ready to be removed on the next update?
if item.bytes_completed == item.total_bytes:
self.completed_items.add(name)
new_completed.append(name)
# If we're only showing summary information, then don't update the individual bars
if not self.per_file_progress:
continue
# If we've run out of bars to use, then collapse the last ones together.
if bar_idx >= len(self.current_bars):
bar = self.current_bars[-1]
in_final_bar_mode = True
final_bar_aggregation_count = bar_idx + 1 - len(self.current_bars)
else:
bar = self.current_bars[bar_idx]
in_final_bar_mode = False
if bar is None:
self.current_bars[bar_idx] = tqdm(
desc=self.format_desc(name, True),
position=2 + bar_idx, # Set to the position past the initial bars.
total=item.total_bytes,
initial=item.bytes_completed,
**self.tqdm_settings,
)
elif in_final_bar_mode:
bar.n += item.bytes_completed
bar.total += item.total_bytes
bar.set_description(self.format_desc(f"[+ {final_bar_aggregation_count} files]", True), refresh=False)
else:
bar.set_description(self.format_desc(name, True), refresh=False)
bar.n = item.bytes_completed
bar.total = item.total_bytes
bar_idx += 1
# Remove all the completed ones from the ordered dictionary
for name in new_completed:
# Only remove ones from consideration to make room for more items coming in.
if len(self.item_state) <= self.n_lines:
break
del self.item_state[name]
if self.per_file_progress:
# Now manually refresh each of the bars
for bar in self.current_bars:
if bar:
bar.refresh()
# Update overall bars
bytes_inc = max(0, group_report.total_bytes_completed - self._prev_bytes_completed)
transfer_inc = max(0, group_report.total_transfer_bytes_completed - self._prev_transfer_bytes_completed)
self._prev_bytes_completed = group_report.total_bytes_completed
self._prev_transfer_bytes_completed = group_report.total_transfer_bytes_completed
self.data_processing_bar.total = group_report.total_bytes
total_files_count = self.total_files if self.total_files is not None else len(self.known_items)
self.data_processing_bar.set_description(
self.format_desc(f"Processing Files ({len(self.completed_items)} / {total_files_count})", False),
refresh=False,
)
self.data_processing_bar.set_postfix_str(
_format_speed_postfix(group_report.total_bytes_completion_rate), refresh=False
)
self.data_processing_bar.update(bytes_inc)
self.upload_bar.total = group_report.total_transfer_bytes
self.upload_bar.set_postfix_str(
_format_speed_postfix(group_report.total_transfer_bytes_completion_rate), refresh=False
)
self.upload_bar.update(transfer_inc)
def close(self):
self.data_processing_bar.close()
self.upload_bar.close()
if self.per_file_progress:
for bar in self.current_bars:
if bar:
bar.close()