From 17106b1e404680b4e1e13ef1e495e5b66b53d214 Mon Sep 17 00:00:00 2001 From: Jamie Pine Date: Fri, 30 Jan 2026 16:47:54 -0800 Subject: [PATCH] Add progress tracking and caching checks for model downloads - Introduced methods to check if models are cached locally in MLX and PyTorch backends. - Enhanced progress tracking during model loading to filter out non-download progress when models are cached. - Updated HFProgressTracker to conditionally report progress based on download status. - Added test scripts for monitoring SSE events during model downloads and verifying progress tracking functionality. - Improved overall error handling and logging for better debugging during model download processes. --- backend/backends/mlx_backend.py | 135 ++++++--- backend/backends/pytorch_backend.py | 134 ++++++--- backend/tests/test_check_progress_state.py | 48 +++ backend/tests/test_generation_progress.py | 321 +++++++++++++++++++++ backend/tests/test_progress.py | 313 ++++++++++++++++++++ backend/tests/test_real_download.py | 178 ++++++++++++ backend/utils/hf_progress.py | 66 ++++- 7 files changed, 1094 insertions(+), 101 deletions(-) create mode 100644 backend/tests/test_check_progress_state.py create mode 100644 backend/tests/test_generation_progress.py create mode 100644 backend/tests/test_progress.py create mode 100644 backend/tests/test_real_download.py diff --git a/backend/backends/mlx_backend.py b/backend/backends/mlx_backend.py index a019f418..8d84776c 100644 --- a/backend/backends/mlx_backend.py +++ b/backend/backends/mlx_backend.py @@ -52,6 +52,24 @@ class MLXTTSBackend: return hf_model_id + def _is_model_cached(self, model_size: str) -> bool: + """ + Check if the model is already cached locally. + + Args: + model_size: Model size to check + + Returns: + True if model is cached, False otherwise + """ + try: + from huggingface_hub import constants as hf_constants + model_path = self._get_model_path(model_size) + repo_cache = Path(hf_constants.HF_HUB_CACHE) / ("models--" + model_path.replace("/", "--")) + return repo_cache.exists() + except Exception: + return False + async def load_model_async(self, model_size: Optional[str] = None): """ Lazy load the MLX TTS model. @@ -86,39 +104,47 @@ class MLXTTSBackend: # Set up progress tracking progress_manager = get_progress_manager() + task_manager = get_task_manager() model_name = f"qwen-tts-{model_size}" - # Start tracking download task - task_manager = get_task_manager() - task_manager.start_download(model_name) + # Check if model is already cached + is_cached = self._is_model_cached(model_size) + + # Set up progress callback + # If cached: filter out non-download progress + # If not cached: report all progress (we're actually downloading) + progress_callback = create_hf_progress_callback(model_name, progress_manager) + tracker = HFProgressTracker(progress_callback, filter_non_downloads=is_cached) print(f"Loading MLX TTS model {model_size}...") - # Initialize progress state - progress_manager.update_progress( - model_name=model_name, - current=0, - total=1, - filename="", - status="downloading", - ) + # Only track download progress if model is NOT cached + if not is_cached: + # Start tracking download task + task_manager.start_download(model_name) + + # Initialize progress state + progress_manager.update_progress( + model_name=model_name, + current=0, + total=1, + filename="", + status="downloading", + ) - # Set up progress callback - progress_callback = create_hf_progress_callback(model_name, progress_manager) - tracker = HFProgressTracker(progress_callback) - - # Use progress tracker during download + # Use progress tracker (tqdm is patched, but filters out non-download progress) with tracker.patch_download(): # Load MLX model (downloads automatically) self.model = load(model_path) + # Only mark download as complete if we were tracking it + if not is_cached: + progress_manager.mark_complete(model_name) + task_manager.complete_download(model_name) + self._current_model_size = model_size self.model_size = model_size - # Mark as complete - progress_manager.mark_complete(model_name) - task_manager.complete_download(model_name) - print(f"MLX TTS model {model_size} loaded successfully") except ImportError as e: @@ -332,6 +358,24 @@ class MLXSTTBackend: """Check if model is loaded.""" return self.model is not None + def _is_model_cached(self, model_size: str) -> bool: + """ + Check if the Whisper model is already cached locally. + + Args: + model_size: Model size to check + + Returns: + True if model is cached, False otherwise + """ + try: + from huggingface_hub import constants as hf_constants + model_name = f"openai/whisper-{model_size}" + repo_cache = Path(hf_constants.HF_HUB_CACHE) / ("models--" + model_name.replace("/", "--")) + return repo_cache.exists() + except Exception: + return False + async def load_model_async(self, model_size: Optional[str] = None): """ Lazy load the MLX Whisper model. @@ -354,55 +398,60 @@ class MLXSTTBackend: def _load_model_sync(self, model_size: str): """Synchronous model loading.""" try: - # IMPORTANT: Set up progress tracking BEFORE importing mlx_audio - # This ensures tqdm is patched before any HuggingFace Hub imports progress_manager = get_progress_manager() + task_manager = get_task_manager() progress_model_name = f"whisper-{model_size}" + # Check if model is already cached + is_cached = self._is_model_cached(model_size) + # Set up progress callback and tracker + # If cached: filter out non-download progress + # If not cached: report all progress (we're actually downloading) progress_callback = create_hf_progress_callback(progress_model_name, progress_manager) - tracker = HFProgressTracker(progress_callback) + tracker = HFProgressTracker(progress_callback, filter_non_downloads=is_cached) # Patch tqdm BEFORE importing mlx_audio - # This is critical because mlx_audio imports huggingface_hub which imports tqdm print("[DEBUG] Starting tqdm patch BEFORE mlx_audio import") tracker_context = tracker.patch_download() tracker_context.__enter__() - print("[DEBUG] tqdm patched, now importing mlx_audio") + print("[DEBUG] tqdm patched") - # NOW import mlx_audio - it will use our patched tqdm + # Import mlx_audio from mlx_audio.stt import load # MLX Whisper uses the standard OpenAI models model_name = f"openai/whisper-{model_size}" - - # Start tracking download task - task_manager = get_task_manager() - task_manager.start_download(progress_model_name) print(f"Loading MLX Whisper model {model_size}...") - # Initialize progress state - progress_manager.update_progress( - model_name=progress_model_name, - current=0, - total=1, - filename="", - status="downloading", - ) + # Only track download progress if model is NOT cached + if not is_cached: + # Start tracking download task + task_manager.start_download(progress_model_name) - # Load the model (tqdm is already patched from above) + # Initialize progress state + progress_manager.update_progress( + model_name=progress_model_name, + current=0, + total=1, + filename="", + status="downloading", + ) + + # Load the model (tqdm is patched, but filters out non-download progress) try: self.model = load(model_name) finally: # Exit the patch context tracker_context.__exit__(None, None, None) - self.model_size = model_size + # Only mark download as complete if we were tracking it + if not is_cached: + progress_manager.mark_complete(progress_model_name) + task_manager.complete_download(progress_model_name) - # Mark as complete - progress_manager.mark_complete(progress_model_name) - task_manager.complete_download(progress_model_name) + self.model_size = model_size print(f"MLX Whisper model {model_size} loaded successfully") diff --git a/backend/backends/pytorch_backend.py b/backend/backends/pytorch_backend.py index cd1257cb..7f8c2379 100644 --- a/backend/backends/pytorch_backend.py +++ b/backend/backends/pytorch_backend.py @@ -58,6 +58,24 @@ class PyTorchTTSBackend: return hf_model_map[model_size] + def _is_model_cached(self, model_size: str) -> bool: + """ + Check if the model is already cached locally. + + Args: + model_size: Model size to check + + Returns: + True if model is cached, False otherwise + """ + try: + from huggingface_hub import constants as hf_constants + model_path = self._get_model_path(model_size) + repo_cache = Path(hf_constants.HF_HUB_CACHE) / ("models--" + model_path.replace("/", "--")) + return repo_cache.exists() + except Exception: + return False + async def load_model_async(self, model_size: Optional[str] = None): """ Lazy load the TTS model with automatic downloading from HuggingFace Hub. @@ -85,20 +103,24 @@ class PyTorchTTSBackend: def _load_model_sync(self, model_size: str): """Synchronous model loading.""" try: - # IMPORTANT: Set up progress tracking BEFORE importing qwen_tts - # This ensures tqdm is patched before any HuggingFace Hub imports progress_manager = get_progress_manager() + task_manager = get_task_manager() model_name = f"qwen-tts-{model_size}" + # Check if model is already cached + is_cached = self._is_model_cached(model_size) + # Set up progress callback and tracker + # If cached: filter out non-download progress (like "Segment 1/1" during generation) + # If not cached: report all progress (we're actually downloading) progress_callback = create_hf_progress_callback(model_name, progress_manager) - tracker = HFProgressTracker(progress_callback) + tracker = HFProgressTracker(progress_callback, filter_non_downloads=is_cached) # Patch tqdm BEFORE importing qwen_tts tracker_context = tracker.patch_download() tracker_context.__enter__() - # NOW import qwen_tts - it will use our patched tqdm + # Import qwen_tts from qwen_tts import Qwen3TTSModel # Get model path (local or HuggingFace Hub ID) @@ -106,20 +128,21 @@ class PyTorchTTSBackend: print(f"Loading TTS model {model_size} on {self.device}...") - # Start tracking download task - task_manager = get_task_manager() - task_manager.start_download(model_name) + # Only track download progress if model is NOT cached + if not is_cached: + # Start tracking download task + task_manager.start_download(model_name) - # Initialize progress state to show download has started - progress_manager.update_progress( - model_name=model_name, - current=0, - total=1, # Set to 1 initially, will be updated by callback - filename="", - status="downloading", - ) + # Initialize progress state to show download has started + progress_manager.update_progress( + model_name=model_name, + current=0, + total=1, # Set to 1 initially, will be updated by callback + filename="", + status="downloading", + ) - # Load the model (tqdm is already patched from above) + # Load the model (tqdm is patched, but filters out non-download progress) try: self.model = Qwen3TTSModel.from_pretrained( model_path, @@ -130,9 +153,10 @@ class PyTorchTTSBackend: # Exit the patch context tracker_context.__exit__(None, None, None) - # Mark as complete - progress_manager.mark_complete(model_name) - task_manager.complete_download(model_name) + # Only mark download as complete if we were tracking it + if not is_cached: + progress_manager.mark_complete(model_name) + task_manager.complete_download(model_name) self._current_model_size = model_size self.model_size = model_size @@ -321,6 +345,24 @@ class PyTorchSTTBackend: """Check if model is loaded.""" return self.model is not None + def _is_model_cached(self, model_size: str) -> bool: + """ + Check if the Whisper model is already cached locally. + + Args: + model_size: Model size to check + + Returns: + True if model is cached, False otherwise + """ + try: + from huggingface_hub import constants as hf_constants + model_name = f"openai/whisper-{model_size}" + repo_cache = Path(hf_constants.HF_HUB_CACHE) / ("models--" + model_name.replace("/", "--")) + return repo_cache.exists() + except Exception: + return False + async def load_model_async(self, model_size: Optional[str] = None): """ Lazy load the Whisper model. @@ -349,14 +391,18 @@ class PyTorchSTTBackend: """Synchronous model loading.""" print(f"[DEBUG] _load_model_sync called for Whisper {model_size}") try: - # IMPORTANT: Set up progress tracking BEFORE importing transformers - # This ensures tqdm is patched before any HuggingFace Hub imports progress_manager = get_progress_manager() + task_manager = get_task_manager() progress_model_name = f"whisper-{model_size}" + # Check if model is already cached + is_cached = self._is_model_cached(model_size) + # Set up progress callback and tracker + # If cached: filter out non-download progress + # If not cached: report all progress (we're actually downloading) progress_callback = create_hf_progress_callback(progress_model_name, progress_manager) - tracker = HFProgressTracker(progress_callback) + tracker = HFProgressTracker(progress_callback, filter_non_downloads=is_cached) # Patch tqdm BEFORE importing transformers print("[DEBUG] Starting tqdm patch BEFORE transformers import") @@ -364,31 +410,32 @@ class PyTorchSTTBackend: tracker_context.__enter__() print("[DEBUG] tqdm patched, now importing transformers") - # NOW import transformers - it will use our patched tqdm + # Import transformers from transformers import WhisperProcessor, WhisperForConditionalGeneration model_name = f"openai/whisper-{model_size}" print(f"[DEBUG] Model name: {model_name}") - # Start tracking download task - task_manager = get_task_manager() - task_manager.start_download(progress_model_name) - print(f"[DEBUG] Task manager started download") - print(f"Loading Whisper model {model_size} on {self.device}...") - # Initialize progress state to show download has started - print(f"[DEBUG] Calling update_progress...") - progress_manager.update_progress( - model_name=progress_model_name, - current=0, - total=1, # Set to 1 initially, will be updated by callback - filename="", - status="downloading", - ) - print(f"[DEBUG] update_progress called, listeners: {len(progress_manager._listeners.get(progress_model_name, []))}") + # Only track download progress if model is NOT cached + if not is_cached: + # Start tracking download task + task_manager.start_download(progress_model_name) + print(f"[DEBUG] Task manager started download") - # Load models (tqdm is already patched from above) + # Initialize progress state to show download has started + print(f"[DEBUG] Calling update_progress...") + progress_manager.update_progress( + model_name=progress_model_name, + current=0, + total=1, # Set to 1 initially, will be updated by callback + filename="", + status="downloading", + ) + print(f"[DEBUG] update_progress called, listeners: {len(progress_manager._listeners.get(progress_model_name, []))}") + + # Load models (tqdm is patched, but filters out non-download progress) try: self.processor = WhisperProcessor.from_pretrained(model_name) self.model = WhisperForConditionalGeneration.from_pretrained(model_name) @@ -396,13 +443,14 @@ class PyTorchSTTBackend: # Exit the patch context tracker_context.__exit__(None, None, None) + # Only mark download as complete if we were tracking it + if not is_cached: + progress_manager.mark_complete(progress_model_name) + task_manager.complete_download(progress_model_name) + self.model.to(self.device) self.model_size = model_size - # Mark as complete - progress_manager.mark_complete(progress_model_name) - task_manager.complete_download(progress_model_name) - print(f"Whisper model {model_size} loaded successfully") except Exception as e: diff --git a/backend/tests/test_check_progress_state.py b/backend/tests/test_check_progress_state.py new file mode 100644 index 00000000..6737547f --- /dev/null +++ b/backend/tests/test_check_progress_state.py @@ -0,0 +1,48 @@ +""" +Check the internal state of ProgressManager. +""" +import sys +import time + +# Add backend to path +sys.path.insert(0, '.') + +from utils.progress import get_progress_manager +from utils.tasks import get_task_manager + +def main(): + pm = get_progress_manager() + tm = get_task_manager() + + print("=" * 60) + print("ProgressManager State Inspector") + print("=" * 60) + + # Check all progress + print("\nAll progress entries:") + print(f" _progress dict: {pm._progress}") + + # Check listeners + print("\nActive listeners:") + print(f" _listeners dict: {pm._listeners}") + + # Check main loop + print(f"\nMain loop set: {pm._main_loop is not None}") + if pm._main_loop: + print(f" Loop running: {pm._main_loop.is_running()}") + + # Check task manager + print("\nTaskManager state:") + print(f" Active downloads: {tm.get_active_downloads()}") + print(f" Active generations: {tm.get_active_generations()}") + + # Check for specific model + print("\nChecking whisper-base specifically:") + progress = pm.get_progress("whisper-base") + if progress: + print(f" Progress: {progress}") + else: + print(" No progress data found") + +if __name__ == "__main__": + main() diff --git a/backend/tests/test_generation_progress.py b/backend/tests/test_generation_progress.py new file mode 100644 index 00000000..5cbe3fdf --- /dev/null +++ b/backend/tests/test_generation_progress.py @@ -0,0 +1,321 @@ +""" +Test TTS generation with SSE progress monitoring. +This test captures the exact SSE events triggered during generation +to identify UX issues where users see download progress even when +the model is already cached. +""" + +import asyncio +import json +import httpx +from typing import List, Dict, Optional +from datetime import datetime + + +async def monitor_sse_stream(model_name: str, timeout: int = 120): + """Monitor SSE stream for a model during generation.""" + events: List[Dict] = [] + url = f"http://localhost:8000/models/progress/{model_name}" + + print(f"[{_timestamp()}] Connecting to SSE endpoint: {url}") + + try: + async with httpx.AsyncClient(timeout=timeout) as client: + async with client.stream("GET", url) as response: + print(f"[{_timestamp()}] SSE connected, status: {response.status_code}") + + if response.status_code != 200: + print(f"[{_timestamp()}] Error: SSE endpoint returned {response.status_code}") + return events + + async for line in response.aiter_lines(): + if not line: + continue + + timestamp = _timestamp() + + if line.startswith("data: "): + try: + data = json.loads(line[6:]) + print(f"[{timestamp}] → SSE Event: {data['status']:12} {data.get('progress', 0):6.1f}% {data.get('filename', '')}") + events.append({ + **data, + "_timestamp": timestamp + }) + + # Stop if complete or error + if data.get("status") in ("complete", "error"): + print(f"[{timestamp}] → Model {data['status']}!") + break + + except json.JSONDecodeError as e: + print(f"[{timestamp}] Error parsing JSON: {e}") + print(f" Line was: {line}") + + elif line.startswith(": heartbeat"): + print(f"[{timestamp}] ♥ heartbeat") + + except asyncio.TimeoutError: + print(f"[{_timestamp()}] SSE monitoring timed out") + except Exception as e: + print(f"[{_timestamp()}] SSE error: {e}") + + return events + + +async def trigger_generation(profile_id: str, text: str, model_size: str = "1.7B"): + """Trigger TTS generation via the API.""" + url = "http://localhost:8000/generate" + + print(f"\n[{_timestamp()}] Triggering generation...") + print(f" Profile: {profile_id}") + print(f" Text: {text[:50]}...") + print(f" Model: {model_size}") + + try: + async with httpx.AsyncClient(timeout=120) as client: + response = await client.post(url, json={ + "profile_id": profile_id, + "text": text, + "language": "en", + "model_size": model_size, + }) + + print(f"[{_timestamp()}] Response: {response.status_code}") + + if response.status_code == 200: + result = response.json() + print(f"[{_timestamp()}] ✓ Generation successful!") + print(f" Generation ID: {result.get('id')}") + print(f" Duration: {result.get('duration', 0):.2f}s") + return True, result + elif response.status_code == 202: + # Model is being downloaded + result = response.json() + print(f"[{_timestamp()}] → Model download in progress") + print(f" Detail: {result}") + return False, result + else: + print(f"[{_timestamp()}] ✗ Error: {response.text}") + return False, None + + except Exception as e: + print(f"[{_timestamp()}] ✗ Exception: {e}") + return False, None + + +async def get_first_profile(): + """Get the first available voice profile.""" + url = "http://localhost:8000/profiles" + + try: + async with httpx.AsyncClient(timeout=10) as client: + response = await client.get(url) + if response.status_code == 200: + profiles = response.json() + if profiles: + return profiles[0]["id"] + except Exception as e: + print(f"Error getting profiles: {e}") + + return None + + +async def check_server(): + """Check if the server is running.""" + try: + async with httpx.AsyncClient(timeout=5) as client: + response = await client.get("http://localhost:8000/health") + return response.status_code == 200 + except Exception as e: + print(f"Server not running: {e}") + return False + + +def _timestamp(): + """Get current timestamp for logging.""" + return datetime.now().strftime("%H:%M:%S.%f")[:-3] + + +async def test_generation_with_cached_model(): + """ + Test Case 1: Generation when model is already cached. + + This should NOT show any download progress events. + If it does, that's the UX bug we're trying to fix. + """ + print("\n" + "=" * 80) + print("TEST CASE 1: Generation with Cached Model") + print("=" * 80) + print("Expected: No download progress events (or minimal/instant completion)") + print("Actual UX Issue: Users see 'started' and 'finished' events even for cached models") + print("=" * 80) + + model_size = "1.7B" + model_name = f"qwen-tts-{model_size}" + + # Get a profile + profile_id = await get_first_profile() + if not profile_id: + print("✗ No voice profiles found. Please create a profile first.") + return False + + print(f"\nUsing profile: {profile_id}") + + # Start SSE monitor BEFORE triggering generation + monitor_task = asyncio.create_task(monitor_sse_stream(model_name, timeout=30)) + + # Wait for SSE to connect + await asyncio.sleep(1) + + # Trigger generation + test_text = "Hello, this is a test of the voice generation system." + success, result = await trigger_generation(profile_id, test_text, model_size) + + if not success and result and result.get("downloading"): + print("\n⚠ Model is being downloaded. Waiting for download to complete...") + # Wait for SSE monitor to capture download events + events = await monitor_task + return events + + # Wait a bit more to catch any progress events + await asyncio.sleep(3) + + # Cancel SSE monitor + monitor_task.cancel() + try: + events = await monitor_task + except asyncio.CancelledError: + events = [] + + return events + + +async def test_generation_with_fresh_download(): + """ + Test Case 2: Generation when model needs to be downloaded. + + This SHOULD show download progress events. + """ + print("\n" + "=" * 80) + print("TEST CASE 2: Generation with Model Download") + print("=" * 80) + print("Expected: Download progress events from 0% to 100%") + print("=" * 80) + + # Use a different model size to force download + model_size = "0.6B" # Smaller model for faster testing + model_name = f"qwen-tts-{model_size}" + + # Get a profile + profile_id = await get_first_profile() + if not profile_id: + print("✗ No voice profiles found. Please create a profile first.") + return False + + print(f"\nUsing profile: {profile_id}") + print("Note: This will download the model if not cached") + + # Start SSE monitor BEFORE triggering generation + monitor_task = asyncio.create_task(monitor_sse_stream(model_name, timeout=300)) + + # Wait for SSE to connect + await asyncio.sleep(1) + + # Trigger generation + test_text = "This should trigger a model download if the model is not cached." + success, result = await trigger_generation(profile_id, test_text, model_size) + + if not success and result and result.get("downloading"): + print("\n→ Model download initiated. Monitoring progress...") + # Wait for download to complete + events = await monitor_task + + # Try generation again + print(f"\n[{_timestamp()}] Retrying generation after download...") + await asyncio.sleep(2) + success, result = await trigger_generation(profile_id, test_text, model_size) + + if success: + print("✓ Generation successful after download") + + return events + + # If model was already cached + await asyncio.sleep(3) + monitor_task.cancel() + try: + events = await monitor_task + except asyncio.CancelledError: + events = [] + + return events + + +async def main(): + print("=" * 80) + print("TTS Generation Progress Test") + print("=" * 80) + print("Purpose: Capture exact SSE events during generation to identify UX issues") + print("=" * 80) + + # Check if server is running + print(f"\n[{_timestamp()}] Checking if server is running...") + if not await check_server(): + print("✗ Server is not running on http://localhost:8000") + print("\nPlease start the server first:") + print(" cd backend && python main.py") + return False + + print("✓ Server is running") + + # Test Case 1: Cached model + print("\n" + "🧪 " * 20) + events_cached = await test_generation_with_cached_model() + + # Results for Test Case 1 + print("\n" + "=" * 80) + print("TEST CASE 1 RESULTS: Generation with Cached Model") + print("=" * 80) + + if not events_cached: + print("✓ GOOD: No SSE progress events received") + print(" This is the expected behavior for a cached model.") + else: + print(f"⚠ ISSUE FOUND: Received {len(events_cached)} SSE events:") + print("\nEvent Timeline:") + for i, event in enumerate(events_cached, 1): + timestamp = event.pop("_timestamp", "??:??:??.???") + print(f" {i}. [{timestamp}] {event}") + + print("\n⚠ This explains the UX issue!") + print(" Users see progress events even when the model is already cached,") + print(" making them think the model is downloading again.") + + # Test Case 2: Fresh download (optional, commented out by default) + # Uncomment if you want to test download progress + # print("\n" + "🧪 " * 20) + # events_download = await test_generation_with_fresh_download() + # + # print("\n" + "=" * 80) + # print("TEST CASE 2 RESULTS: Generation with Model Download") + # print("=" * 80) + # + # if not events_download: + # print("ℹ Model was already cached, no download occurred") + # else: + # print(f"✓ Received {len(events_download)} download progress events") + # print("\nDownload Timeline:") + # for i, event in enumerate(events_download, 1): + # timestamp = event.pop("_timestamp", "??:??:??.???") + # print(f" {i}. [{timestamp}] {event}") + + print("\n" + "=" * 80) + print("Test Complete!") + print("=" * 80) + + return True + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/backend/tests/test_progress.py b/backend/tests/test_progress.py new file mode 100644 index 00000000..a66ba079 --- /dev/null +++ b/backend/tests/test_progress.py @@ -0,0 +1,313 @@ +""" +Test script to debug model download progress tracking. +""" + +import asyncio +import json +import time +from typing import List, Dict +import logging + +# Set up logging to see what's happening +logging.basicConfig( + level=logging.DEBUG, + format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' +) + +from utils.progress import ProgressManager, get_progress_manager +from utils.hf_progress import HFProgressTracker, create_hf_progress_callback + + +def test_progress_manager_basic(): + """Test 1: Basic ProgressManager functionality.""" + print("\n" + "=" * 60) + print("Test 1: ProgressManager Basic Operations") + print("=" * 60) + + pm = ProgressManager() + + # Test update_progress + pm.update_progress( + model_name="test-model", + current=50, + total=100, + filename="test.bin", + status="downloading" + ) + + # Test get_progress + progress = pm.get_progress("test-model") + print(f"✓ Progress stored: {progress}") + assert progress is not None + assert progress["progress"] == 50.0 + assert progress["filename"] == "test.bin" + assert progress["status"] == "downloading" + + # Test mark_complete + pm.mark_complete("test-model") + progress = pm.get_progress("test-model") + print(f"✓ Marked complete: {progress}") + assert progress["status"] == "complete" + assert progress["progress"] == 100.0 + + print("✓ Test 1 PASSED\n") + return True + + +async def test_progress_manager_sse(): + """Test 2: ProgressManager SSE streaming.""" + print("\n" + "=" * 60) + print("Test 2: ProgressManager SSE Streaming") + print("=" * 60) + + pm = ProgressManager() + collected_events: List[Dict] = [] + + # Simulate SSE client + async def sse_client(): + """Simulates a frontend SSE connection.""" + print(" SSE client: Subscribing to test-model-sse...") + async for event in pm.subscribe("test-model-sse"): + # Parse SSE event + if event.startswith("data: "): + data = json.loads(event[6:]) + print(f" SSE client: Received event: {data['status']} - {data.get('progress', 0):.1f}%") + collected_events.append(data) + + # Stop when complete + if data.get("status") in ("complete", "error"): + break + elif event.startswith(": heartbeat"): + print(" SSE client: Received heartbeat") + + # Simulate download progress updates (from backend thread) + async def simulate_download(): + """Simulates backend sending progress updates.""" + print(" Backend: Starting simulated download...") + await asyncio.sleep(0.2) # Let SSE client subscribe first + + # Send progress updates + for i in range(0, 101, 20): + print(f" Backend: Updating progress to {i}%") + pm.update_progress( + model_name="test-model-sse", + current=i, + total=100, + filename=f"file_{i}.bin", + status="downloading" if i < 100 else "downloading" + ) + await asyncio.sleep(0.1) + + # Mark complete + print(" Backend: Marking download complete") + pm.mark_complete("test-model-sse") + + # Run SSE client and download simulation concurrently + await asyncio.gather( + sse_client(), + simulate_download() + ) + + # Verify we got events + print(f"\n Collected {len(collected_events)} events") + assert len(collected_events) > 0, "Should have received at least one event" + assert collected_events[-1]["status"] == "complete", "Last event should be 'complete'" + + print("✓ Test 2 PASSED\n") + return True + + +def test_hf_progress_tracker(): + """Test 3: HFProgressTracker tqdm patching.""" + print("\n" + "=" * 60) + print("Test 3: HFProgressTracker tqdm Patching") + print("=" * 60) + + captured_progress: List[tuple] = [] + + def progress_callback(downloaded: int, total: int, filename: str): + """Capture progress updates.""" + captured_progress.append((downloaded, total, filename)) + print(f" Progress callback: {downloaded}/{total} bytes ({filename})") + + tracker = HFProgressTracker(progress_callback) + + # Simulate a download with tqdm + with tracker.patch_download(): + try: + from tqdm import tqdm + + # Simulate downloading a file + print(" Simulating download with tqdm...") + total_size = 1000 + with tqdm(total=total_size, desc="model.bin", unit="B", unit_scale=True) as pbar: + for chunk in range(0, total_size, 100): + pbar.update(100) + time.sleep(0.01) + + print(f" Captured {len(captured_progress)} progress updates") + assert len(captured_progress) > 0, "Should have captured progress updates" + + # Verify progress increases + last_downloaded = 0 + for downloaded, total, filename in captured_progress: + assert downloaded >= last_downloaded, "Downloaded bytes should increase" + assert total == total_size, "Total should be consistent" + last_downloaded = downloaded + + print("✓ Test 3 PASSED\n") + return True + + except ImportError: + print("✗ tqdm not available, skipping test\n") + return None + + +async def test_full_integration(): + """Test 4: Full integration test.""" + print("\n" + "=" * 60) + print("Test 4: Full Integration (ProgressManager + HFProgressTracker)") + print("=" * 60) + + pm = get_progress_manager() + collected_events: List[Dict] = [] + + # SSE client + async def sse_client(): + print(" SSE client: Subscribing...") + async for event in pm.subscribe("integration-test"): + if event.startswith("data: "): + data = json.loads(event[6:]) + print(f" SSE client: {data['status']} - {data.get('progress', 0):.1f}% - {data.get('filename', '')}") + collected_events.append(data) + if data.get("status") in ("complete", "error"): + break + + # Simulate backend download with HFProgressTracker + async def simulate_real_download(): + await asyncio.sleep(0.2) # Let SSE subscribe + + print(" Backend: Starting download with HFProgressTracker...") + + # Set up tracking (like the real backend does) + progress_callback = create_hf_progress_callback("integration-test", pm) + tracker = HFProgressTracker(progress_callback) + + # Initialize progress + pm.update_progress( + model_name="integration-test", + current=0, + total=1, + filename="", + status="downloading" + ) + + # Simulate download with tqdm patching + with tracker.patch_download(): + try: + from tqdm import tqdm + + # Simulate multi-file download (like HuggingFace does) + files = [ + ("model.safetensors", 5000), + ("config.json", 1000), + ("tokenizer.json", 500), + ] + + for filename, size in files: + print(f" Backend: Downloading {filename}...") + with tqdm(total=size, desc=filename, unit="B") as pbar: + for chunk in range(0, size, 500): + chunk_size = min(500, size - chunk) + pbar.update(chunk_size) + await asyncio.sleep(0.05) + + # Mark complete + print(" Backend: Download complete") + pm.mark_complete("integration-test") + + except ImportError: + print(" ✗ tqdm not available") + pm.mark_error("integration-test", "tqdm not available") + + # Run both + await asyncio.gather( + sse_client(), + simulate_real_download() + ) + + # Verify + print(f"\n Collected {len(collected_events)} events") + if len(collected_events) > 0: + print(f" First event: {collected_events[0]}") + print(f" Last event: {collected_events[-1]}") + assert collected_events[-1]["status"] == "complete", "Should end with 'complete'" + print("✓ Test 4 PASSED\n") + return True + else: + print("✗ Test 4 FAILED - No events received\n") + return False + + +async def main(): + """Run all tests.""" + print("\n" + "=" * 60) + print("Voicebox Progress Tracking Test Suite") + print("=" * 60) + + results = [] + + # Test 1: Basic operations + try: + results.append(("Basic Operations", test_progress_manager_basic())) + except Exception as e: + print(f"✗ Test 1 FAILED: {e}\n") + results.append(("Basic Operations", False)) + + # Test 2: SSE streaming + try: + results.append(("SSE Streaming", await test_progress_manager_sse())) + except Exception as e: + print(f"✗ Test 2 FAILED: {e}\n") + results.append(("SSE Streaming", False)) + + # Test 3: tqdm patching + try: + results.append(("tqdm Patching", test_hf_progress_tracker())) + except Exception as e: + print(f"✗ Test 3 FAILED: {e}\n") + results.append(("tqdm Patching", False)) + + # Test 4: Full integration + try: + results.append(("Full Integration", await test_full_integration())) + except Exception as e: + print(f"✗ Test 4 FAILED: {e}\n") + results.append(("Full Integration", False)) + + # Summary + print("\n" + "=" * 60) + print("Test Results Summary") + print("=" * 60) + + for name, result in results: + status = "✓ PASS" if result else ("⊘ SKIP" if result is None else "✗ FAIL") + print(f" {status:8} {name}") + + passed = sum(1 for _, r in results if r is True) + failed = sum(1 for _, r in results if r is False) + skipped = sum(1 for _, r in results if r is None) + + print() + print(f" Total: {len(results)} tests") + print(f" Passed: {passed}") + print(f" Failed: {failed}") + print(f" Skipped: {skipped}") + print("=" * 60 + "\n") + + return failed == 0 + + +if __name__ == "__main__": + success = asyncio.run(main()) + exit(0 if success else 1) diff --git a/backend/tests/test_real_download.py b/backend/tests/test_real_download.py new file mode 100644 index 00000000..9bebca85 --- /dev/null +++ b/backend/tests/test_real_download.py @@ -0,0 +1,178 @@ +""" +Test real model download with SSE progress monitoring. +""" + +import asyncio +import json +import httpx +import time +from typing import List, Dict + +async def monitor_sse_stream(model_name: str, timeout: int = 300): + """Monitor SSE stream for a model download.""" + events: List[Dict] = [] + url = f"http://localhost:8000/models/progress/{model_name}" + + print(f"Connecting to SSE endpoint: {url}") + + async with httpx.AsyncClient(timeout=timeout) as client: + async with client.stream("GET", url) as response: + print(f"SSE connected, status: {response.status_code}") + + if response.status_code != 200: + print(f"Error: SSE endpoint returned {response.status_code}") + return events + + async for line in response.aiter_lines(): + if not line: + continue + + print(f" Raw SSE: {line[:100]}...") # Print first 100 chars + + if line.startswith("data: "): + try: + data = json.loads(line[6:]) + print(f" → {data['status']:12} {data.get('progress', 0):6.1f}% {data.get('filename', '')}") + events.append(data) + + # Stop if complete or error + if data.get("status") in ("complete", "error"): + print(f" Download {data['status']}!") + break + + except json.JSONDecodeError as e: + print(f" Error parsing JSON: {e}") + print(f" Line was: {line}") + + elif line.startswith(": heartbeat"): + print(" ♥ heartbeat") + + return events + + +async def trigger_download(model_name: str): + """Trigger a model download via the API.""" + url = "http://localhost:8000/models/download" + + print(f"\nTriggering download for: {model_name}") + + async with httpx.AsyncClient(timeout=300) as client: + response = await client.post(url, json={"model_name": model_name}) + print(f"Response: {response.status_code} - {response.json()}") + return response.status_code == 200 + + +async def check_server(): + """Check if the server is running.""" + try: + async with httpx.AsyncClient(timeout=5) as client: + response = await client.get("http://localhost:8000/health") + return response.status_code == 200 + except Exception as e: + print(f"Server not running: {e}") + return False + + +async def main(): + print("=" * 60) + print("Real Model Download Progress Test") + print("=" * 60) + + # Check if server is running + print("\nChecking if server is running...") + if not await check_server(): + print("✗ Server is not running on http://localhost:8000") + print("\nPlease start the server first:") + print(" cd backend && python main.py") + return False + + print("✓ Server is running") + + # Choose a small model for testing + model_name = "whisper-base" # ~150MB, faster to download + print(f"\nUsing model: {model_name}") + + # Option to delete model first if it exists + print("\nDo you want to delete the model first to force a fresh download? (y/n)") + # For automated testing, skip deletion prompt + # delete_first = input().strip().lower() == 'y' + delete_first = False + + if delete_first: + print(f"Deleting {model_name}...") + async with httpx.AsyncClient(timeout=30) as client: + response = await client.delete(f"http://localhost:8000/models/{model_name}") + print(f"Delete response: {response.status_code}") + + print("\n" + "=" * 60) + print("Starting Test") + print("=" * 60) + + # Start monitoring SSE stream BEFORE triggering download + async def run_test(): + # Start SSE monitor in background + monitor_task = asyncio.create_task(monitor_sse_stream(model_name)) + + # Wait a bit to ensure SSE is connected + await asyncio.sleep(1) + + # Trigger download + success = await trigger_download(model_name) + + if not success: + print("✗ Failed to trigger download") + monitor_task.cancel() + return False + + # Wait for SSE monitor to complete + events = await monitor_task + + return events + + events = await run_test() + + # Results + print("\n" + "=" * 60) + print("Test Results") + print("=" * 60) + + if not events: + print("✗ FAILED - No SSE events received!") + print("\nPossible causes:") + print(" 1. SSE endpoint not working") + print(" 2. Progress updates not being sent") + print(" 3. Model already downloaded (no progress to report)") + print("\nTry deleting the model first to force a fresh download:") + print(f" curl -X DELETE http://localhost:8000/models/{model_name}") + return False + + print(f"✓ Received {len(events)} SSE events") + print(f"\nFirst event: {events[0]}") + print(f"Last event: {events[-1]}") + + # Check if we got meaningful progress + has_progress = any(e.get('progress', 0) > 0 for e in events) + has_complete = any(e.get('status') == 'complete' for e in events) + + if has_progress: + print("✓ Progress updates received") + else: + print("✗ No progress updates (might be already downloaded)") + + if has_complete: + print("✓ Download completed successfully") + else: + print("✗ Download did not complete") + + success = has_progress and has_complete + + if success: + print("\n✓ TEST PASSED - Progress tracking works!") + else: + print("\n⊘ TEST INCONCLUSIVE - Try with a fresh download") + + return success + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/backend/utils/hf_progress.py b/backend/utils/hf_progress.py index 3b88fe4b..014645d5 100644 --- a/backend/utils/hf_progress.py +++ b/backend/utils/hf_progress.py @@ -11,8 +11,9 @@ import sys class HFProgressTracker: """Tracks HuggingFace Hub download progress by intercepting tqdm.""" - def __init__(self, progress_callback: Optional[Callable] = None): + def __init__(self, progress_callback: Optional[Callable] = None, filter_non_downloads: bool = False): self.progress_callback = progress_callback + self.filter_non_downloads = filter_non_downloads # Only filter if True self._original_tqdm_class = None self._lock = threading.Lock() self._total_downloaded = 0 @@ -80,7 +81,7 @@ class HFProgressTracker: } def update(self, n=1): - print(f"[DEBUG TrackedTqdm] update called with n={n}") + print(f"[DEBUG TrackedTqdm] update called with n={n}, filter_non_downloads={tracker.filter_non_downloads}") result = super().update(n) # Report progress @@ -91,24 +92,59 @@ class HFProgressTracker: total = getattr(self, "total", 0) if total and total > 0: - # Update per-file tracking - tracker._file_sizes[filename] = total - tracker._file_downloaded[filename] = current + # Determine if we should report this progress + should_report = True - # Calculate totals across all files - tracker._total_size = sum(tracker._file_sizes.values()) - tracker._total_downloaded = sum(tracker._file_downloaded.values()) + if tracker.filter_non_downloads: + # Filter out non-download progress bars (e.g., "Segment 1/1" during generation) + # Only report progress for actual file downloads from HuggingFace + should_report = self._is_download_progress(filename) - # Call progress callback - if tracker.progress_callback: - tracker.progress_callback( - tracker._total_downloaded, - tracker._total_size, - filename - ) + if should_report: + # Update per-file tracking + tracker._file_sizes[filename] = total + tracker._file_downloaded[filename] = current + + # Calculate totals across all files + tracker._total_size = sum(tracker._file_sizes.values()) + tracker._total_downloaded = sum(tracker._file_downloaded.values()) + + # Call progress callback + if tracker.progress_callback: + tracker.progress_callback( + tracker._total_downloaded, + tracker._total_size, + filename + ) return result + 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 + download_extensions = [ + '.safetensors', '.bin', '.pt', '.pth', # Model weights + '.json', '.txt', '.py', # Config files + '.msgpack', '.h5', # Other formats + ] + + filename_lower = filename.lower() + has_extension = any(filename_lower.endswith(ext) for ext in download_extensions) + + # Skip internal progress indicators + 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 + def close(self): with tracker._lock: if id(self) in tracker._active_tqdms: