[Sonica-B]

Fixes #3759.

Description

The Dl Size... bar created in _Downloader.tqdm() counts MiB (unit=' MiB'), but two of the three paths that feed it add bytes:

• increase_tqdm(), called for files that are already cached (DownloadManager._log_skip_download), added url_info.size in bytes.
• _sync_file_copy(), used for non-HTTP URLs copied through tf.io.gfile, also added url_info.size in bytes.
• _sync_download() (HTTP/FTP) converts to MiB, but its modulo bookkeeping never flushed the last MiB when a download ended exactly on a MiB boundary. A 3 MiB file read in 64 KiB chunks stops at 2/3.

This PR routes the first two paths through one helper, _update_dl_size_pbar(size_in_bytes), which converts with // units.MiB the same way _sync_download computes its total. _sync_download now reports size // units.MiB as it grows, so the progress always matches the Content-length // MiB total. Files smaller than 1 MiB count as 0 MiB on every path, which is what the HTTP path already did. UrlInfo.size and checksums are unchanged.

If you would rather review the HTTP last-MiB change separately, I can split it into its own PR.

flowchart LR
  subgraph Before
    A1["cached file<br>increase_tqdm()"] -->|"+157286400 (bytes)"| B1[("Dl Size bar<br>unit: MiB")]
    A2["gfile copy<br>_sync_file_copy()"] -->|"+2097152 (bytes)"| B1
    A3["HTTP<br>_sync_download()"] -->|"+2 of 3 MiB"| B1
  end
  subgraph After
    C1["cached file<br>increase_tqdm()"] --> H["_update_dl_size_pbar()<br>bytes // MiB"]
    C2["gfile copy<br>_sync_file_copy()"] --> H
    H -->|"+150, +2"| B2[("Dl Size bar<br>unit: MiB")]
    C3["HTTP<br>_sync_download()<br>size // MiB"] -->|"+3"| B2
  end

Loading

Before / after

One cached 150 MiB file, then a 2 MiB local file copy, then a 3 MiB HTTP download (mocked, 64 KiB chunks). The expected bar is 155/155 MiB.

Before (master at 7b86414c):

after 150 MiB cached file : n=157286400 total=157286400
after 2 MiB local copy   : n=159383552 total=159383552
after 3 MiB HTTP download: n=159383554 total=159383555  (expected 155 / 155 MiB)
Dl Size...: 100%|#########9| 159383554/159383555 [00:36<00:00, 4410061.32 MiB/s]

After:

after 150 MiB cached file : n=150 total=150
after 2 MiB local copy   : n=152 total=152
after 3 MiB HTTP download: n=155 total=155  (expected 155 / 155 MiB)
Dl Size...: 100%|##########| 155/155 [00:44<00:00,  3.39 MiB/s]

Repro script

"""Repro for tensorflow/datasets#3759: `Dl Size...` bar mixes bytes and MiB."""
import contextlib, os, tempfile
from unittest import mock
from etils import epath
from tensorflow_datasets.core import units, utils
from tensorflow_datasets.core.download import checksums, downloader

class _Resp:
  def __init__(self, n): self.headers = {'Content-length': n}
  url = 'http://example.com/fresh.bin'; cookies = {}; status_code = 200

def _open_url(url, verify):  # 3 MiB HTTP response in 64 KiB chunks
  n = 3 * units.MiB
  chunks = iter([b'x' * (64 * 1024)] * (n // (64 * 1024)))
  @contextlib.contextmanager
  def cm(): yield _Resp(n), chunks
  return cm()

tmp = tempfile.mkdtemp()
src = os.path.join(tmp, 'local_2mib.bin')
with open(src, 'wb') as f: f.write(b'y' * (2 * units.MiB))
dest = os.path.join(tmp, 'dl'); os.makedirs(dest)

d = downloader.get_downloader(1)
with mock.patch.object(downloader, '_get_filename', lambda r: 'fresh.bin'), \
     mock.patch.object(downloader, '_open_url', _open_url), d.tqdm():
  bar = d._pbar_dl_size._pbar
  d.increase_tqdm(checksums.UrlInfo(size=utils.Size(150 * units.MiB), checksum='0' * 64, filename='cached.tar'))
  print(f'after 150 MiB cached file : n={bar.n} total={bar.total}')
  d.download(src, epath.Path(dest)).get()
  print(f'after 2 MiB local copy   : n={bar.n} total={bar.total}')
  d.download('http://example.com/fresh.bin', epath.Path(dest)).get()
  print(f'after 3 MiB HTTP download: n={bar.n} total={bar.total}  (expected 155 / 155 MiB)')

Tests

Three regression tests in downloader_test.py, one per path: cached file, gfile copy, and HTTP with a final MiB on a chunk boundary. On master they fail with:

AssertionError: Tuples differ: (3145728, 3145728) != (3, 3)
AssertionError: Tuples differ: (2097152, 2097152) != (2, 2)
AssertionError: Tuples differ: (3, 2) != (3, 3)

With this change, pytest tensorflow_datasets/core/download gives 81 passed, 5 skipped, 7 failed (master: 78 passed, 5 skipped, 7 failed). The 7 failures are the same extractor_test.py cases on both runs and are unrelated to this PR. On a Git-for-Windows checkout with core.autocrlf=true, testing/test_data/foo.csv is converted to CRLF and no longer matches the LF copy inside the test archives. I can send a separate .gitattributes fix if that is welcome.

Validated on Windows 11 with CPU TensorFlow (Python 3.11, tensorflow-cpu 2.21.0, tqdm 4.70.1). Not run on Linux; CI will confirm there on Python 3.10-3.12.

Checklist

• Add passing test(s)
• Lint code. pylint --rcfile=.pylintrc reports no new messages beyond protected-access on _pbar_dl_size in the new test helper, the same access setUp already makes.
• Dataset-specific items (checksums, BibTeX, test data, download_and_prepare): not applicable, this is a library fix.