From 0d97776fe883fce765f81e82b956316dda2d0a5c Mon Sep 17 00:00:00 2001 From: Suraj Mishra Date: Wed, 11 Nov 2026 05:11:42 +0000 Subject: [PATCH] Keep the weights download progress output short in logs urlretrieve calls the progress hook for every 8 KB block, about 8,000 times for the weights, and each call rewrote the progress line. In a terminal it now redraws only when the whole percentage changes; when stderr is not a terminal it writes one line per 10%. Fixes #5 Co-Authored-By: Claude Opus 5.5 Assisted-by: claude-code (claude-opus-5-5) --- CHANGELOG.md | 6 +++++ src/subcellseg/weights.py | 47 ++++++++++++++++++++++++++++++--------- tests/test_subcellseg.py | 21 +++++++++++++++++ 3 files changed, 64 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e744cdb..8a182d3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,12 @@ All notable changes to SubCellSeg are listed here. The format follows (was 10%). Fine-tuning uses a learning rate of 5e-5 and stops after 50 epochs without improvement; training from scratch keeps 1e-4 and 100. +### Fixed +- The progress line for the first-time weights download no longer floods + logs: in a terminal it redraws once per percent, and when output goes to a + file, CI or a notebook it writes one line per 10% (was about 8,000 redraws, + ~150 KB). ([#5](https://github.com/AllenCell/SubCellSeg/issues/5)) + ### Added - Benchmarks section in the README and `tools/benchmark/` scripts. - Contributing guide, code of conduct, security policy, citation file, diff --git a/src/subcellseg/weights.py b/src/subcellseg/weights.py index 7cf7a24..8e89d9d 100644 --- a/src/subcellseg/weights.py +++ b/src/subcellseg/weights.py @@ -32,14 +32,40 @@ def sha256(path: Path) -> str: return digest.hexdigest() -def _report(blocks: int, block_size: int, total: int) -> None: - done = blocks * block_size / 2**20 - if total > 0: - pct = min(100, int(done * 2**20 * 100 / total)) - sys.stderr.write(f"\r {done:.0f} / {total / 2**20:.0f} MB ({pct}%)") - else: - sys.stderr.write(f"\r {done:.0f} MB") - sys.stderr.flush() +class _Progress: + """``urlretrieve`` reporthook that writes only when the shown value changes. + + urlretrieve calls the hook for every 8 KB block (about 8,000 times for the + weights). In a terminal the line is redrawn in place when the whole percentage + (or, with no known size, the whole MB) changes. When stderr is not a terminal + (CI, log files, notebooks) one line is written per 10%, so logs stay short. + """ + + def __init__(self) -> None: + self.tty = sys.stderr.isatty() + self.last: int | None = None + + def __call__(self, blocks: int, block_size: int, total: int) -> None: + done = blocks * block_size + if total > 0: + done = min(done, total) + pct = done * 100 // total + step = pct if self.tty else pct // 10 * 10 + text = f"{done / 2**20:.0f} / {total / 2**20:.0f} MB ({step}%)" + elif self.tty: + step = done // 2**20 + text = f"{step} MB" + else: + return + if step == self.last: + return + self.last = step + sys.stderr.write(f"\r {text}" if self.tty else f" {text}\n") + sys.stderr.flush() + + def finish(self) -> None: + if self.tty and self.last is not None: + sys.stderr.write("\n") def fetch_pretrained() -> Path: @@ -59,16 +85,17 @@ def fetch_pretrained() -> Path: path.parent.mkdir(parents=True, exist_ok=True) partial = path.with_name(path.name + ".part") print(f"Downloading SubCellSeg weights to {path}", file=sys.stderr) + progress = _Progress() try: try: - urllib.request.urlretrieve(url, partial, reporthook=_report) + urllib.request.urlretrieve(url, partial, reporthook=progress) except OSError as err: # no network, blocked download, missing file raise RuntimeError( f"Could not download the pretrained weights from {url} ({err}).\n" f"Download that file in a web browser and save it as {path}, " "or pass a weights file with --checkpoint." ) from None - sys.stderr.write("\n") + progress.finish() if WEIGHTS_SHA256 and not os.environ.get("SUBCELLSEG_WEIGHTS_URL"): found = sha256(partial) if found != WEIGHTS_SHA256: diff --git a/tests/test_subcellseg.py b/tests/test_subcellseg.py index c40c595..65723cc 100644 --- a/tests/test_subcellseg.py +++ b/tests/test_subcellseg.py @@ -2,6 +2,7 @@ they need neither a GPU nor the pretrained weights.""" import csv +import sys import numpy as np import pytest @@ -333,3 +334,23 @@ def test_weights_download_is_verified_and_cached(tmp_path, monkeypatch): path.unlink() with pytest.raises(RuntimeError, match="Could not download"): weights.fetch_pretrained() + + +@pytest.mark.parametrize("tty", [False, True]) +def test_download_progress_stays_short(tty, monkeypatch, capsys): + from subcellseg import weights + + monkeypatch.setattr(sys.stderr, "isatty", lambda: tty) + progress = weights._Progress() + total, block = 64 * 2**20, 8192 # the real weights: about 8,000 hook calls + for blocks in range(total // block + 2): + progress(blocks, block, total) + progress.finish() + err = capsys.readouterr().err + assert "(100%)" in err + if tty: + assert err.count("\r") == 101 # one redraw per whole percent + else: + assert "\r" not in err + lines = err.splitlines() + assert [line.rsplit("(", 1)[1] for line in lines] == [f"{p}%)" for p in range(0, 101, 10)]