velixio
b7caa494eb
Five workstreams that finish the remote-GPU line, plus the test hole that let a broken signature reach a commit. **Downloads go through the normal path** (Phase 5). Rather than a second remote-only route, the existing Models install flow became target-aware, so a model landing on a worker uses the same code, the same progress events and the same UI as a local one. Progress rows key on (target, repo_id) — the aggregator keyed on bare repo_id, so the same model downloading here and on a worker at once collapsed into one row that told the user nothing true about either. **Audiobooks render chapter by chapter on the worker** (Phase 8), with per-chapter local fallback and ONE aggregated notice. The failure that shape exists to prevent: a remote GPU that sleeps at chapter 40 of 200 must not turn a working book into 160 rows of PROGRESS_LEASE_EXPIRED. Dictation is deliberately NOT ported — it runs ASR per utterance inside a live WebSocket loop, and paying queue admission plus a round trip there would spend the one thing that route is for. **Dubbing stays local, and says so** (Phase 7). The coarse worker operation is not finished, so the picker still reports dubbing as local rather than showing a green remote chip over work this machine is doing. What could not wait is the in-loop OOM retry: it sniffed the error string and flushed the *local* CUDA cache, which under remote execution is the wrong machine's GPU entirely. That is fixed now, before the path that would have exercised it exists. **Two instances can no longer share the control plane.** A second VoiceStudio silently bound the same worker port and coexisted, so remote workers landed on whichever process won the race — a session that registers with one instance and appears dead to the other. This produced hours of misdiagnosis during hardware testing and would hit any user with the app open twice. The second instance now keeps running locally and explains the conflict instead of quietly competing. **And the hole that allowed all this to be missable.** gpu_gateway called Scheduler.submit(pinned_worker_id=...) one commit before that parameter existed. Every remote generation raised TypeError; 5236 tests passed anyway, because nothing exercised the gateway against the real scheduler. tests/test_gpu_gateway_scheduler_contract.py now runs that path for real and binds every gateway→dependency call signature. Verified by renaming the parameter away and watching both tests fail with the original error. Gallery previews also fall back to a local render when a downloaded clip cannot be decoded, rather than yielding silence. Backend 5274 passed, frontend 1812 passed. Not yet verified on hardware: Phases 4, 5, 6, 7, 8. Only the TTS path and its artifact transport have been proven on a real GPU.
310 lines
12 KiB
Python
310 lines
12 KiB
Python
"""HuggingFace download progress — one monkey-patch, every `hf_hub_download`
|
|
reports bytes downloaded through a central callback.
|
|
|
|
`huggingface_hub` uses tqdm for progress bars; we subclass it, intercept
|
|
`update()` calls, and forward (filename, downloaded_bytes, total_bytes) to
|
|
whatever callback is registered. No changes to calling sites across
|
|
transformers / mlx_whisper / diffusers / accelerate — they all route through
|
|
`hf_hub_download`, which uses the patched tqdm.
|
|
|
|
Usage:
|
|
from utils.hf_progress import install, register_listener, unregister_listener
|
|
|
|
install() # once at app startup
|
|
listener_id = register_listener(lambda ev: print(ev))
|
|
# …models download, listener fires…
|
|
unregister_listener(listener_id)
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import contextvars
|
|
import itertools
|
|
import logging
|
|
import threading
|
|
from typing import Callable, Optional
|
|
|
|
logger = logging.getLogger("omnivoice.hf_progress")
|
|
|
|
# Context-scoped active repo_id. Set in the install/delete handler so every
|
|
# tqdm event fired while a snapshot_download runs can be stamped with the
|
|
# originating repo, letting the frontend route per-file events to the right
|
|
# row instead of heuristically matching filename substrings.
|
|
current_repo_id: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar(
|
|
"omnivoice_hf_progress_repo_id", default=None,
|
|
)
|
|
current_target: contextvars.ContextVar[str] = contextvars.ContextVar(
|
|
"omnivoice_hf_progress_target", default="local",
|
|
)
|
|
|
|
# Event shape forwarded to listeners. Typed loosely on purpose — SSE encodes
|
|
# it as JSON so consumers read the dict directly.
|
|
# {
|
|
# "filename": str, # desc on the tqdm bar, usually the HF file path
|
|
# "downloaded": int, # bytes pulled so far
|
|
# "total": int | None, # total bytes or None if unknown
|
|
# "pct": float, # 0.0-1.0 (or 0.0 if total unknown)
|
|
# "phase": "start"|"progress"|"done",
|
|
# }
|
|
ProgressEvent = dict
|
|
Listener = Callable[[ProgressEvent], None]
|
|
|
|
_listeners: dict[int, Listener] = {}
|
|
_listener_lock = threading.Lock()
|
|
_listener_counter = itertools.count(1)
|
|
_installed = False
|
|
_install_lock = threading.Lock()
|
|
|
|
# Set by install() to the TrackedTqdm subclass so call sites can drive it
|
|
# explicitly via snapshot_download(tqdm_class=...) instead of relying solely on
|
|
# the global monkey-patch. Xet feeds bytes into whatever tqdm_class is passed,
|
|
# so this is also the xet-aware progress hook (FDL-02).
|
|
_tracked_tqdm_class: Optional[type] = None
|
|
|
|
|
|
def tracked_tqdm_class() -> Optional[type]:
|
|
"""Return the progress-emitting tqdm subclass (or None if install() hasn't
|
|
run / huggingface_hub's tqdm couldn't be patched). Pass it as
|
|
``snapshot_download(tqdm_class=...)`` to drive progress deterministically."""
|
|
return _tracked_tqdm_class
|
|
|
|
|
|
# Optional sink fed every per-file (repo_id, filename, downloaded, total) byte
|
|
# update — used by utils.download_aggregator to build the overall aggregate bar
|
|
# (FDL-06). Kept as a setter to avoid a circular import (this module must not
|
|
# import the aggregator). Signature: fn(repo_id, filename, downloaded, total).
|
|
_byte_sink: Optional[Callable] = None
|
|
|
|
|
|
def set_byte_sink(fn: Optional[Callable]) -> None:
|
|
global _byte_sink
|
|
_byte_sink = fn
|
|
|
|
|
|
def register_listener(cb: Listener) -> int:
|
|
"""Register a callback that receives progress events. Returns an id that
|
|
can be passed to `unregister_listener` when the listener is done."""
|
|
with _listener_lock:
|
|
lid = next(_listener_counter)
|
|
_listeners[lid] = cb
|
|
return lid
|
|
|
|
|
|
def unregister_listener(lid: int) -> None:
|
|
with _listener_lock:
|
|
_listeners.pop(lid, None)
|
|
|
|
|
|
def _emit(event: ProgressEvent) -> None:
|
|
"""Fan out to all registered listeners. Never raise — a bad listener
|
|
shouldn't break a download."""
|
|
# Stamp the active repo_id so frontends can route events to the right
|
|
# row. Only set when this emit is happening inside an install handler.
|
|
rid = current_repo_id.get()
|
|
if rid is not None and "repo_id" not in event:
|
|
event = {**event, "repo_id": rid}
|
|
if "target" not in event:
|
|
event = {**event, "target": current_target.get()}
|
|
with _listener_lock:
|
|
listeners = list(_listeners.values())
|
|
for cb in listeners:
|
|
try:
|
|
cb(event)
|
|
except Exception as e: # noqa: BLE001
|
|
logger.debug("hf_progress listener raised: %s", e)
|
|
|
|
|
|
def emit(event: ProgressEvent) -> None:
|
|
"""Public emit — lets non-tqdm operations (delete, verify, etc.) push
|
|
lifecycle events onto the same SSE stream."""
|
|
_emit(event)
|
|
|
|
|
|
class SafeFileWrapper:
|
|
def __init__(self, fp):
|
|
self.fp = fp
|
|
self._is_safe_wrapper = True
|
|
def write(self, s):
|
|
try:
|
|
self.fp.write(s)
|
|
except (OSError, UnicodeError):
|
|
# OSError: EPIPE from a dead parent shell (the wrapper's original
|
|
# job). UnicodeError (#1155): a library print of user text hitting
|
|
# a non-UTF-8 stream — cp1252 stdout on Windows — must not abort
|
|
# the operation that printed. Logs are best-effort; work is not.
|
|
pass
|
|
def flush(self):
|
|
try:
|
|
getattr(self.fp, 'flush', lambda: None)()
|
|
except OSError:
|
|
pass
|
|
def __getattr__(self, name):
|
|
return getattr(self.fp, name)
|
|
|
|
def install() -> None:
|
|
"""Monkey-patch `huggingface_hub`'s tqdm so every download reports to our
|
|
listeners. Safe to call multiple times — second call is a no-op."""
|
|
global _installed
|
|
with _install_lock:
|
|
if _installed:
|
|
return
|
|
# `huggingface_hub.utils.__init__` does `from .tqdm import tqdm`,
|
|
# which shadows the `tqdm` SUBMODULE with the CLASS of the same name
|
|
# when accessed via attribute lookup. Pull the real module out of
|
|
# sys.modules after an explicit import so we patch the right thing.
|
|
try:
|
|
import sys
|
|
import huggingface_hub.utils.tqdm # noqa: F401
|
|
hf_tqdm_module = sys.modules.get("huggingface_hub.utils.tqdm")
|
|
if hf_tqdm_module is None:
|
|
raise ImportError("huggingface_hub.utils.tqdm not in sys.modules after import")
|
|
except Exception as e: # noqa: BLE001
|
|
logger.warning(
|
|
"hf_progress.install: huggingface_hub.utils.tqdm missing (%s); "
|
|
"progress tracking disabled.", e,
|
|
)
|
|
return
|
|
|
|
original = getattr(hf_tqdm_module, "tqdm", None)
|
|
if original is None or not isinstance(original, type):
|
|
logger.warning("hf_progress.install: no `tqdm` class on the module; aborting")
|
|
return
|
|
|
|
class TrackedTqdm(original): # type: ignore[misc,valid-type]
|
|
"""tqdm subclass that emits a progress event on every update."""
|
|
|
|
_last_emit_time: float = 0.0
|
|
|
|
@staticmethod
|
|
def status_printer(file):
|
|
if file is not None and not getattr(file, "_is_safe_wrapper", False):
|
|
file = SafeFileWrapper(file)
|
|
try:
|
|
return original.status_printer(file)
|
|
except Exception:
|
|
return lambda s: None
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
if 'file' in kwargs and kwargs['file'] is not None and not getattr(kwargs['file'], "_is_safe_wrapper", False):
|
|
kwargs['file'] = SafeFileWrapper(kwargs['file'])
|
|
try:
|
|
super().__init__(*args, **kwargs)
|
|
except OSError:
|
|
pass
|
|
|
|
if hasattr(self, 'fp') and getattr(self, 'fp', None) is not None and not getattr(self.fp, "_is_safe_wrapper", False):
|
|
self.fp = SafeFileWrapper(self.fp)
|
|
|
|
import time as _t
|
|
self._last_emit_time = _t.monotonic()
|
|
try:
|
|
desc = getattr(self, "desc", None)
|
|
total = int(getattr(self, "total", 0) or 0)
|
|
_emit({
|
|
"filename": str(desc or "download"),
|
|
"downloaded": 0,
|
|
"total": total,
|
|
"pct": 0.0,
|
|
"phase": "start",
|
|
})
|
|
except Exception:
|
|
pass
|
|
|
|
def _emit_progress(self):
|
|
"""Emit current state as a progress event."""
|
|
try:
|
|
desc = getattr(self, "desc", None)
|
|
total = int(getattr(self, "total", 0) or 0)
|
|
done = int(getattr(self, "n", 0) or 0)
|
|
pct = (done / total) if total > 0 else 0.0
|
|
# Pull rate from tqdm's own calculations if available
|
|
rate = None
|
|
try:
|
|
rate = self.format_dict.get("rate")
|
|
except Exception:
|
|
pass
|
|
event = {
|
|
"filename": str(desc or "download"),
|
|
"downloaded": done,
|
|
"total": total,
|
|
"pct": pct,
|
|
"phase": "done" if (total > 0 and done >= total) else "progress",
|
|
}
|
|
if rate and rate > 0:
|
|
event["rate"] = rate # bytes/sec from tqdm
|
|
_emit(event)
|
|
# Feed the overall aggregator (FDL-06), if wired.
|
|
self._feed_sink(done, total, complete=False)
|
|
except Exception:
|
|
pass
|
|
|
|
def _feed_sink(self, done, total, *, complete: bool):
|
|
"""Forward a byte/count update to the overall aggregator sink.
|
|
|
|
Passes the tqdm `unit` so the aggregator can tell a byte bar
|
|
(unit 'B') from the "Fetching N files" count bar, and a stable
|
|
per-bar key (id(self)) because every per-file byte bar shares
|
|
the default desc 'download' under Xet — keying by desc would
|
|
collapse them into one.
|
|
"""
|
|
sink = _byte_sink
|
|
if sink is None:
|
|
return
|
|
rid = current_repo_id.get()
|
|
if not rid:
|
|
return
|
|
try:
|
|
unit = getattr(self, "unit", None)
|
|
sink(rid, id(self), unit, int(done or 0), int(total or 0), complete)
|
|
except Exception:
|
|
pass
|
|
|
|
def update(self, n=1):
|
|
try:
|
|
super().update(n)
|
|
except OSError:
|
|
pass
|
|
import time as _t
|
|
now = _t.monotonic()
|
|
# Throttle: emit at most every 0.3s to avoid flooding SSE
|
|
if (now - self._last_emit_time) >= 0.3:
|
|
self._last_emit_time = now
|
|
self._emit_progress()
|
|
|
|
def display(self, msg=None, pos=None):
|
|
"""tqdm calls display() on its refresh cycle; piggyback for
|
|
periodic emits even when update() intervals are large."""
|
|
import time as _t
|
|
now = _t.monotonic()
|
|
if (now - self._last_emit_time) >= 0.5:
|
|
self._last_emit_time = now
|
|
self._emit_progress()
|
|
try:
|
|
return super().display(msg, pos)
|
|
except OSError:
|
|
pass
|
|
|
|
def close(self):
|
|
# Credit the file's full size to the aggregator on close. Under
|
|
# Xet a per-file byte bar often never increments `n` (Xet fetches
|
|
# chunks out-of-band), so completion is the only reliable signal
|
|
# that the file's bytes landed. Harmless for classic LFS bars
|
|
# (n already == total).
|
|
try:
|
|
total = int(getattr(self, "total", 0) or 0)
|
|
if total > 0:
|
|
self._feed_sink(total, total, complete=True)
|
|
except Exception:
|
|
pass
|
|
try:
|
|
super().close()
|
|
except OSError:
|
|
pass
|
|
|
|
# Stash the original for inspection / uninstall, then swap.
|
|
hf_tqdm_module._omnivoice_original_tqdm = original # type: ignore[attr-defined]
|
|
hf_tqdm_module.tqdm = TrackedTqdm # type: ignore[assignment]
|
|
global _tracked_tqdm_class
|
|
_tracked_tqdm_class = TrackedTqdm
|
|
_installed = True
|
|
logger.info("hf_progress: installed tqdm patch on huggingface_hub.utils.tqdm")
|