mirror of
https://github.com/jamiepine/voicebox.git
synced 2026-09-15 04:40:40 -07:00
Refactor model download progress tracking and enhance SSE handling
- Rearranged imports for consistency in useModelDownloadToast hook. - Improved logging in useModelDownloadToast for better debugging during download events. - Updated progress calculation to handle cases where progress exceeds 100%. - Enhanced toast notifications to reflect download completion and error states. - Introduced throttling in ProgressManager to optimize SSE updates and prevent overwhelming clients. - Added new test scripts for monitoring SSE events during model downloads, ensuring accurate progress reporting.
This commit is contained in:
Jamie Pine
@@ -32,7 +32,6 @@ class HFProgressTracker:
|
||||
"""A tqdm subclass that reports progress to our tracker."""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
print(f"[DEBUG TrackedTqdm] __init__ called with desc: {kwargs.get('desc', '')}")
|
||||
# Extract filename from desc before passing to parent
|
||||
desc = kwargs.get("desc", "")
|
||||
if not desc and args:
|
||||
@@ -81,7 +80,6 @@ class HFProgressTracker:
|
||||
}
|
||||
|
||||
def update(self, n=1):
|
||||
print(f"[DEBUG TrackedTqdm] update called with n={n}, filter_non_downloads={tracker.filter_non_downloads}")
|
||||
result = super().update(n)
|
||||
|
||||
# Report progress
|
||||
@@ -121,10 +119,7 @@ class HFProgressTracker:
|
||||
|
||||
def _is_download_progress(self, filename: str) -> bool:
|
||||
"""Check if this is a real download progress bar vs internal processing."""
|
||||
print(f"[DEBUG _is_download_progress] Checking filename: '{filename}'")
|
||||
|
||||
if not filename or filename == "unknown":
|
||||
print(f"[DEBUG _is_download_progress] Rejected: empty/unknown filename")
|
||||
return False
|
||||
|
||||
# Real downloads have file extensions
|
||||
@@ -141,9 +136,7 @@ class HFProgressTracker:
|
||||
skip_patterns = ['segment', 'processing', 'generating', 'loading']
|
||||
has_skip_pattern = any(pattern in filename_lower for pattern in skip_patterns)
|
||||
|
||||
result = has_extension and not has_skip_pattern
|
||||
print(f"[DEBUG _is_download_progress] has_extension={has_extension}, has_skip_pattern={has_skip_pattern}, result={result}")
|
||||
return result
|
||||
return has_extension and not has_skip_pattern
|
||||
|
||||
def close(self):
|
||||
with tracker._lock:
|
||||
@@ -156,13 +149,11 @@ class HFProgressTracker:
|
||||
@contextmanager
|
||||
def patch_download(self):
|
||||
"""Context manager to patch tqdm for progress tracking."""
|
||||
print("[DEBUG HFProgressTracker] patch_download called")
|
||||
try:
|
||||
import tqdm as tqdm_module
|
||||
|
||||
# Store original tqdm class
|
||||
self._original_tqdm_class = tqdm_module.tqdm
|
||||
print(f"[DEBUG HFProgressTracker] Original tqdm class: {self._original_tqdm_class}")
|
||||
|
||||
# Reset totals
|
||||
with self._lock:
|
||||
@@ -175,22 +166,18 @@ class HFProgressTracker:
|
||||
|
||||
# Create our tracked tqdm class
|
||||
tracked_tqdm = self._create_tracked_tqdm_class()
|
||||
print(f"[DEBUG HFProgressTracker] Created TrackedTqdm class: {tracked_tqdm}")
|
||||
|
||||
# Patch tqdm.tqdm
|
||||
tqdm_module.tqdm = tracked_tqdm
|
||||
print(f"[DEBUG HFProgressTracker] Patched tqdm.tqdm")
|
||||
|
||||
# Also patch tqdm.auto.tqdm if it exists (used by huggingface_hub)
|
||||
self._original_tqdm_auto = None
|
||||
if hasattr(tqdm_module, "auto") and hasattr(tqdm_module.auto, "tqdm"):
|
||||
self._original_tqdm_auto = tqdm_module.auto.tqdm
|
||||
tqdm_module.auto.tqdm = tracked_tqdm
|
||||
print(f"[DEBUG HFProgressTracker] Patched tqdm.auto.tqdm")
|
||||
|
||||
# Patch in sys.modules to catch already-imported references
|
||||
self._patched_modules = {}
|
||||
patched_count = 0
|
||||
for module_name in list(sys.modules.keys()):
|
||||
if "huggingface" in module_name or module_name.startswith("tqdm"):
|
||||
try:
|
||||
@@ -203,11 +190,8 @@ class HFProgressTracker:
|
||||
):
|
||||
self._patched_modules[module_name] = attr
|
||||
setattr(module, "tqdm", tracked_tqdm)
|
||||
patched_count += 1
|
||||
print(f"[DEBUG HFProgressTracker] Patched {module_name}.tqdm")
|
||||
except (AttributeError, TypeError):
|
||||
pass
|
||||
print(f"[DEBUG HFProgressTracker] Patched {patched_count} modules in sys.modules")
|
||||
|
||||
yield
|
||||
|
||||
|
||||
+34
-10
@@ -16,11 +16,17 @@ class ProgressManager:
|
||||
Thread-safe: can be called from background threads (e.g., via asyncio.to_thread).
|
||||
"""
|
||||
|
||||
# Throttle settings to prevent overwhelming SSE clients
|
||||
THROTTLE_INTERVAL_SECONDS = 0.5 # Minimum time between updates
|
||||
THROTTLE_PROGRESS_DELTA = 1.0 # Minimum progress change (%) to force update
|
||||
|
||||
def __init__(self):
|
||||
self._progress: Dict[str, Dict] = {}
|
||||
self._listeners: Dict[str, list] = {}
|
||||
self._lock = threading.Lock() # Thread-safe lock for progress dict
|
||||
self._main_loop: Optional[asyncio.AbstractEventLoop] = None
|
||||
self._last_notify_time: Dict[str, float] = {} # Last notification time per model
|
||||
self._last_notify_progress: Dict[str, float] = {} # Last notified progress per model
|
||||
|
||||
def _set_main_loop(self, loop: asyncio.AbstractEventLoop):
|
||||
"""Set the main event loop for thread-safe operations."""
|
||||
@@ -67,6 +73,10 @@ class ProgressManager:
|
||||
Update progress for a model download.
|
||||
|
||||
Thread-safe: can be called from background threads.
|
||||
|
||||
Progress updates are throttled to prevent overwhelming SSE clients.
|
||||
Updates are sent at most every THROTTLE_INTERVAL_SECONDS, or when
|
||||
progress changes by at least THROTTLE_PROGRESS_DELTA percent.
|
||||
|
||||
Args:
|
||||
model_name: Name of the model (e.g., "qwen-tts-1.7B", "whisper-base")
|
||||
@@ -76,6 +86,7 @@ class ProgressManager:
|
||||
status: Status string (downloading, extracting, complete, error)
|
||||
"""
|
||||
import logging
|
||||
import time
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
progress_pct = (current / total * 100) if total > 0 else 0
|
||||
@@ -90,25 +101,38 @@ class ProgressManager:
|
||||
"timestamp": datetime.now().isoformat(),
|
||||
}
|
||||
|
||||
print(f"[DEBUG] update_progress called: {model_name}, {progress_pct:.1f}%")
|
||||
|
||||
# Thread-safe update of progress dict
|
||||
# Thread-safe update of progress dict (always update internal state)
|
||||
with self._lock:
|
||||
self._progress[model_name] = progress_data
|
||||
|
||||
# Check if we should notify listeners (throttling)
|
||||
current_time = time.time()
|
||||
last_time = self._last_notify_time.get(model_name, 0)
|
||||
last_progress = self._last_notify_progress.get(model_name, -100)
|
||||
|
||||
time_delta = current_time - last_time
|
||||
progress_delta = abs(progress_pct - last_progress)
|
||||
|
||||
# Always notify for complete/error status, or if throttle conditions are met
|
||||
should_notify = (
|
||||
status in ("complete", "error") or
|
||||
time_delta >= self.THROTTLE_INTERVAL_SECONDS or
|
||||
progress_delta >= self.THROTTLE_PROGRESS_DELTA
|
||||
)
|
||||
|
||||
if not should_notify:
|
||||
return # Skip this update (throttled)
|
||||
|
||||
# Update throttle tracking
|
||||
self._last_notify_time[model_name] = current_time
|
||||
self._last_notify_progress[model_name] = progress_pct
|
||||
|
||||
# Notify all listeners (thread-safe)
|
||||
listener_count = len(self._listeners.get(model_name, []))
|
||||
print(f"[DEBUG] Listener count for {model_name}: {listener_count}")
|
||||
print(f"[DEBUG] All listeners: {list(self._listeners.keys())}")
|
||||
print(f"[DEBUG] Main loop set: {self._main_loop is not None}")
|
||||
if self._main_loop:
|
||||
print(f"[DEBUG] Main loop running: {self._main_loop.is_running()}")
|
||||
|
||||
if listener_count > 0:
|
||||
logger.debug(f"Notifying {listener_count} listeners for {model_name}: {progress_pct:.1f}% ({filename})")
|
||||
print(f"[DEBUG] About to notify listeners...")
|
||||
self._notify_listeners_threadsafe(model_name, progress_data)
|
||||
print(f"[DEBUG] Notified listeners")
|
||||
else:
|
||||
logger.debug(f"No listeners for {model_name}, progress update stored: {progress_pct:.1f}%")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user