v0.1.16

robogauge/utils/progress_monitor.py

@@ -0,0 +1,137 @@
+# -*- coding: utf-8 -*-
+'''
+@File    : progress_monitor.py
+@Time    : 2025/12/27 00:10:07
+@Author  : wty-yy
+@Version : 1.0
+@Blog    : https://wty-yy.github.io/
+@Desc    : Centralized progress monitoring for multiprocessing tasks using tqdm
+'''
+from tqdm import tqdm
+import multiprocessing
+from typing import Dict
+from threading import Thread
+from dataclasses import dataclass
+
+class ProgressTypes:
+    INIT = 'init'       # Init progress (set total, desc)
+    UPDATE = 'update'   # Update progress value (set value)
+    DESC = 'desc'       # Update description text only (set desc)
+    RESET = 'reset'     # Reset progress bar (set total, desc)
+    FINISH = 'finish'   # Mark completion (set desc)
+    ERROR = 'error'     # Mark error (set desc)
+
+@dataclass
+class ProgressData:
+    progress_queue: multiprocessing.Queue
+    task_id: int = 0
+    msg_prefix: str = ''
+
+def report_progress(progress_data: ProgressData, msg_type, value=None, desc=None, total=None):
+    """
+    Assistant function to report progress to the main process.
+    Args:
+        queue: multiprocessing.Queue
+        task_id: int, line number of the task
+        msg_type: ProgressTypes
+        value: int, value to update (for UPDATE type)
+        desc: str, description text (for DESC, INIT, FINISH types)
+        total: int, total value (for INIT, RESET types)
+    """
+    if progress_data is None:
+        return
+    queue = progress_data.progress_queue
+    task_id = progress_data.task_id
+    msg_prefix = progress_data.msg_prefix
+    if desc is not None:
+        desc = msg_prefix + desc
+    queue.put((task_id, msg_type, {'value': value, 'desc': desc, 'total': total}))
+
+class ProgressMonitor:
+    def __init__(self, total_rows):
+        self.total_rows = total_rows
+        self.bars: Dict[int, tqdm] = {}
+        
+    def listener_loop(self, queue):
+        """
+        Run in a separate thread in the main process to consume the Queue and update tqdm.
+        """
+        # print(f"Monitor started for {self.total_rows} tasks...")
+        for i in range(self.total_rows):
+            self.bars[i] = tqdm(
+                total=100, 
+                position=i, 
+                desc=f"Task {i} Pending...", 
+                bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}]",
+                leave=True
+            )
+        
+        active_tasks = self.total_rows
+        
+        while active_tasks > 0:
+            record = queue.get()
+            if record is None:  # Poison pill signal
+                break
+            
+            task_id, msg_type, data = record
+            if task_id not in self.bars:
+                continue
+                
+            bar = self.bars[task_id]
+            
+            if msg_type == ProgressTypes.INIT:
+                total = data['total']
+                desc = data['desc']
+                bar.reset(total=total)
+                bar.set_description(desc)
+                bar.refresh()
+                
+            elif msg_type == ProgressTypes.UPDATE:
+                val = data['value']
+                bar.update(val)
+                
+            elif msg_type == ProgressTypes.DESC:
+                desc = data['desc']
+                if desc:
+                    bar.set_description(desc)
+                    
+            elif msg_type == ProgressTypes.RESET:
+                # Scenario: Level search finished, starting MultiPipeline, reset progress bar
+                total = data['total']
+                desc = data['desc']
+                bar.reset(total=total)
+                bar.set_description(desc)
+                bar.refresh()
+                
+            elif msg_type == ProgressTypes.FINISH:
+                desc = data['desc']
+                if desc:
+                    bar.set_description(desc)
+                bar.refresh()
+                # Note: We do not close the bar here to keep it displayed until all tasks are finished and closed together.
+                active_tasks -= 1
+                
+            elif msg_type == ProgressTypes.ERROR:
+                desc = data['desc']
+                bar.set_description(desc)
+                bar.refresh()
+                active_tasks -= 1
+
+        # After all tasks are finished, close all bars
+        for bar in self.bars.values():
+            bar.close()
+
+def start_progress_monitor_thread(total_rows):
+    """
+    Create and start a ProgressMonitor thread. Queue is returned for reporting progress.
+    Args:
+        total_rows: int, number of tasks to monitor
+    Returns:
+        queue: multiprocessing.Queue
+        monitor_thread: threading.Thread
+    """
+    queue = multiprocessing.Manager().Queue()
+    monitor = ProgressMonitor(total_rows)
+    monitor_thread = Thread(target=monitor.listener_loop, args=(queue,))
+    monitor_thread.start()
+    return queue, monitor_thread