Upload backend/hue_portal/chatbot/download_progress.py with huggingface_hub

backend/hue_portal/chatbot/download_progress.py

@@ -0,0 +1,294 @@
+"""
+Download progress tracker for Hugging Face models.
+Tracks real-time download progress in bytes.
+"""
+import threading
+import time
+from typing import Dict, Optional
+from dataclasses import dataclass, field
+
+
+@dataclass
+class DownloadProgress:
+    """Track download progress for a single file."""
+    filename: str
+    total_bytes: int = 0
+    downloaded_bytes: int = 0
+    started_at: Optional[float] = None
+    completed_at: Optional[float] = None
+    speed_bytes_per_sec: float = 0.0
+    
+    @property
+    def percentage(self) -> float:
+        """Calculate download percentage."""
+        if self.total_bytes == 0:
+            return 0.0
+        return min(100.0, (self.downloaded_bytes / self.total_bytes) * 100.0)
+    
+    @property
+    def is_complete(self) -> bool:
+        """Check if download is complete."""
+        return self.total_bytes > 0 and self.downloaded_bytes >= self.total_bytes
+    
+    @property
+    def elapsed_time(self) -> float:
+        """Get elapsed time in seconds."""
+        if self.started_at is None:
+            return 0.0
+        end_time = self.completed_at or time.time()
+        return end_time - self.started_at
+
+
+@dataclass
+class ModelDownloadProgress:
+    """Track overall download progress for a model."""
+    model_path: str
+    files: Dict[str, DownloadProgress] = field(default_factory=dict)
+    started_at: Optional[float] = None
+    completed_at: Optional[float] = None
+    
+    def update_file(self, filename: str, downloaded: int, total: int):
+        """Update progress for a specific file."""
+        if filename not in self.files:
+            self.files[filename] = DownloadProgress(
+                filename=filename,
+                started_at=time.time()
+            )
+            if self.started_at is None:
+                self.started_at = time.time()
+        
+        file_progress = self.files[filename]
+        file_progress.downloaded_bytes = downloaded
+        file_progress.total_bytes = total
+        
+        # Calculate speed
+        if file_progress.started_at:
+            elapsed = time.time() - file_progress.started_at
+            if elapsed > 0:
+                file_progress.speed_bytes_per_sec = downloaded / elapsed
+        
+        # Mark as complete
+        if total > 0 and downloaded >= total:
+            file_progress.completed_at = time.time()
+    
+    def complete_file(self, filename: str):
+        """Mark a file as complete."""
+        if filename in self.files:
+            self.files[filename].completed_at = time.time()
+    
+    @property
+    def total_bytes(self) -> int:
+        """Get total bytes across all files."""
+        return sum(f.total_bytes for f in self.files.values())
+    
+    @property
+    def downloaded_bytes(self) -> int:
+        """Get downloaded bytes across all files."""
+        return sum(f.downloaded_bytes for f in self.files.values())
+    
+    @property
+    def percentage(self) -> float:
+        """Calculate overall download percentage."""
+        total = self.total_bytes
+        if total == 0:
+            # If no total yet, count completed files
+            if len(self.files) == 0:
+                return 0.0
+            completed = sum(1 for f in self.files.values() if f.is_complete)
+            return (completed / len(self.files)) * 100.0
+        return min(100.0, (self.downloaded_bytes / total) * 100.0)
+    
+    @property
+    def is_complete(self) -> bool:
+        """Check if all files are downloaded."""
+        if len(self.files) == 0:
+            return False
+        return all(f.is_complete for f in self.files.values())
+    
+    @property
+    def speed_bytes_per_sec(self) -> float:
+        """Get overall download speed."""
+        total_speed = sum(f.speed_bytes_per_sec for f in self.files.values() if f.started_at)
+        return total_speed
+    
+    @property
+    def elapsed_time(self) -> float:
+        """Get elapsed time in seconds."""
+        if self.started_at is None:
+            return 0.0
+        end_time = self.completed_at or time.time()
+        return end_time - self.started_at
+    
+    def to_dict(self) -> Dict:
+        """Convert to dictionary for JSON serialization."""
+        return {
+            "model_path": self.model_path,
+            "total_bytes": self.total_bytes,
+            "downloaded_bytes": self.downloaded_bytes,
+            "percentage": round(self.percentage, 2),
+            "speed_bytes_per_sec": round(self.speed_bytes_per_sec, 2),
+            "speed_mb_per_sec": round(self.speed_bytes_per_sec / (1024 * 1024), 2),
+            "elapsed_time": round(self.elapsed_time, 2),
+            "is_complete": self.is_complete,
+            "files_count": len(self.files),
+            "files_completed": sum(1 for f in self.files.values() if f.is_complete),
+            "files": {
+                name: {
+                    "filename": f.filename,
+                    "total_bytes": f.total_bytes,
+                    "downloaded_bytes": f.downloaded_bytes,
+                    "percentage": round(f.percentage, 2),
+                    "speed_mb_per_sec": round(f.speed_bytes_per_sec / (1024 * 1024), 2),
+                    "is_complete": f.is_complete
+                }
+                for name, f in self.files.items()
+            }
+        }
+
+
+class ProgressTracker:
+    """Thread-safe progress tracker for multiple models."""
+    
+    def __init__(self):
+        self._progress: Dict[str, ModelDownloadProgress] = {}
+        self._lock = threading.Lock()
+    
+    def get_or_create(self, model_path: str) -> ModelDownloadProgress:
+        """Get or create progress tracker for a model."""
+        with self._lock:
+            if model_path not in self._progress:
+                self._progress[model_path] = ModelDownloadProgress(model_path=model_path)
+            return self._progress[model_path]
+    
+    def get(self, model_path: str) -> Optional[ModelDownloadProgress]:
+        """Get progress tracker for a model."""
+        with self._lock:
+            return self._progress.get(model_path)
+    
+    def update(self, model_path: str, filename: str, downloaded: int, total: int):
+        """Update download progress for a file."""
+        progress = self.get_or_create(model_path)
+        progress.update_file(filename, downloaded, total)
+    
+    def complete_file(self, model_path: str, filename: str):
+        """Mark a file as complete."""
+        progress = self.get(model_path)
+        if progress:
+            progress.complete_file(filename)
+    
+    def complete_model(self, model_path: str):
+        """Mark entire model download as complete."""
+        progress = self.get(model_path)
+        if progress:
+            progress.completed_at = time.time()
+    
+    def get_all(self) -> Dict[str, Dict]:
+        """Get all progress as dictionary."""
+        with self._lock:
+            return {
+                path: prog.to_dict()
+                for path, prog in self._progress.items()
+            }
+    
+    def get_model_progress(self, model_path: str) -> Optional[Dict]:
+        """Get progress for a specific model."""
+        progress = self.get(model_path)
+        if progress:
+            return progress.to_dict()
+        return None
+
+
+# Global progress tracker instance
+_global_tracker = ProgressTracker()
+
+
+def get_progress_tracker() -> ProgressTracker:
+    """Get global progress tracker instance."""
+    return _global_tracker
+
+
+def create_progress_callback(model_path: str):
+    """
+    Create a progress callback for huggingface_hub downloads.
+    
+    Usage:
+        from huggingface_hub import snapshot_download
+        callback = create_progress_callback("Qwen/Qwen2.5-32B-Instruct")
+        snapshot_download(repo_id=model_path, resume_download=True, 
+                         tqdm_class=callback)
+    """
+    tracker = get_progress_tracker()
+    
+    class ProgressCallback:
+        """Progress callback for tqdm."""
+        
+        def __init__(self, *args, **kwargs):
+            # Store tqdm arguments but don't initialize yet
+            self.tqdm_args = args
+            self.tqdm_kwargs = kwargs
+            self.current_file = None
+        
+        def __call__(self, *args, **kwargs):
+            # This will be called by huggingface_hub
+            # We'll intercept the progress updates
+            pass
+        
+        def update(self, n: int = 1):
+            """Update progress."""
+            if self.current_file:
+                # Get current progress from tqdm
+                if hasattr(self, 'n'):
+                    downloaded = self.n
+                else:
+                    downloaded = n
+                if hasattr(self, 'total'):
+                    total = self.total
+                else:
+                    total = 0
+                tracker.update(model_path, self.current_file, downloaded, total)
+        
+        def set_description(self, desc: str):
+            """Set description (filename)."""
+            # Extract filename from description
+            if desc:
+                self.current_file = desc.split()[-1] if ' ' in desc else desc
+        
+        def close(self):
+            """Close progress bar."""
+            if self.current_file:
+                tracker.complete_file(model_path, self.current_file)
+    
+    return ProgressCallback
+
+
+def create_hf_progress_callback(model_path: str):
+    """
+    Create a progress callback compatible with huggingface_hub.
+    Returns a function that can be used with tqdm.
+    """
+    tracker = get_progress_tracker()
+    current_file = [None]  # Use list to allow modification in nested function
+    
+    def progress_callback(tqdm_bar):
+        """Progress callback function."""
+        if tqdm_bar.desc:
+            # Extract filename from description
+            filename = tqdm_bar.desc.split()[-1] if ' ' in tqdm_bar.desc else tqdm_bar.desc
+            if filename != current_file[0]:
+                current_file[0] = filename
+                if current_file[0] not in tracker.get_or_create(model_path).files:
+                    tracker.get_or_create(model_path).files[current_file[0]] = DownloadProgress(
+                        filename=current_file[0],
+                        started_at=time.time()
+                    )
+        
+        if current_file[0]:
+            downloaded = getattr(tqdm_bar, 'n', 0)
+            total = getattr(tqdm_bar, 'total', 0)
+            tracker.update(model_path, current_file[0], downloaded, total)
+    
+    return progress_callback
+
+
+
+