v0.1.17; update 7-9
This commit is contained in:
wty-yy
@@ -2,16 +2,17 @@
|
||||
'''
|
||||
@File : progress_monitor.py
|
||||
@Time : 2025/12/27 00:10:07
|
||||
@Author : wty-yy
|
||||
@Author : wty-yy (with Gemini 3)
|
||||
@Version : 1.0
|
||||
@Blog : https://wty-yy.github.io/
|
||||
@Desc : Centralized progress monitoring for multiprocessing tasks using tqdm
|
||||
@Desc : Centralized progress monitoring with Dynamic Slot Management
|
||||
'''
|
||||
from tqdm import tqdm
|
||||
import multiprocessing
|
||||
from typing import Dict
|
||||
from threading import Thread
|
||||
from dataclasses import dataclass
|
||||
import heapq
|
||||
|
||||
class ProgressTypes:
|
||||
INIT = 'init' # Init progress (set total, desc)
|
||||
@@ -50,76 +51,111 @@ def report_progress(progress_data: ProgressData, msg_type, value=None, desc=None
|
||||
class ProgressMonitor:
|
||||
def __init__(self, total_rows):
|
||||
self.total_rows = total_rows
|
||||
self.bars: Dict[int, tqdm] = {}
|
||||
self.active_bars: Dict[int, Dict] = {}
|
||||
self.free_slots = []
|
||||
self.max_slot_used = 0
|
||||
|
||||
def _get_free_slot(self):
|
||||
""" Get a free screen display position """
|
||||
if self.free_slots:
|
||||
return heapq.heappop(self.free_slots)
|
||||
else:
|
||||
self.max_slot_used += 1
|
||||
return self.max_slot_used
|
||||
|
||||
def _release_slot(self, slot):
|
||||
""" Release a screen display position """
|
||||
heapq.heappush(self.free_slots, slot)
|
||||
|
||||
def listener_loop(self, queue):
|
||||
"""
|
||||
Run in a separate thread in the main process to consume the Queue and update tqdm.
|
||||
Listener thread in the main process
|
||||
"""
|
||||
# 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
|
||||
)
|
||||
# Create the main progress bar (Position 0)
|
||||
main_bar = tqdm(
|
||||
total=self.total_rows,
|
||||
position=0,
|
||||
desc="🚀 Total Progress",
|
||||
bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [Done: {n_fmt}] [{elapsed}<{remaining}]",
|
||||
leave=True,
|
||||
smoothing=0.1 # Optional: set smoothing factor (0-1) to prevent drastic ETA fluctuations due to slow tasks
|
||||
)
|
||||
|
||||
active_tasks = self.total_rows
|
||||
# Number of tasks still pending
|
||||
pending_tasks = self.total_rows
|
||||
|
||||
while active_tasks > 0:
|
||||
while pending_tasks > 0:
|
||||
record = queue.get()
|
||||
if record is None: # Poison pill signal
|
||||
break
|
||||
if record is None: 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:
|
||||
# --- Dynamically create/get sub progress bars ---
|
||||
if task_id not in self.active_bars:
|
||||
# Only create bar when message received to avoid Pending screen flicker
|
||||
slot = self._get_free_slot()
|
||||
# leave=False is key: after task completion, clear the line for new task reuse
|
||||
bar = tqdm(
|
||||
total=100,
|
||||
position=slot,
|
||||
desc=f"Task {task_id} Starting...",
|
||||
bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}]",
|
||||
leave=False
|
||||
)
|
||||
self.active_bars[task_id] = {'bar': bar, 'slot': slot}
|
||||
|
||||
bar_info = self.active_bars[task_id]
|
||||
bar = bar_info['bar']
|
||||
|
||||
# --- Handle messages ---
|
||||
try:
|
||||
if msg_type == ProgressTypes.INIT:
|
||||
total = data.get('total', 100)
|
||||
desc = data.get('desc', '')
|
||||
bar.reset(total=total)
|
||||
bar.set_description(desc)
|
||||
bar.refresh()
|
||||
|
||||
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:
|
||||
elif msg_type == ProgressTypes.UPDATE:
|
||||
val = data.get('value', 1)
|
||||
bar.update(val)
|
||||
|
||||
elif msg_type == ProgressTypes.DESC:
|
||||
desc = data.get('desc')
|
||||
if desc: bar.set_description(desc)
|
||||
|
||||
elif msg_type == ProgressTypes.RESET:
|
||||
total = data.get('total', 100)
|
||||
desc = data.get('desc', '')
|
||||
bar.reset(total=total)
|
||||
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
|
||||
bar.refresh()
|
||||
|
||||
elif msg_type == ProgressTypes.FINISH or msg_type == ProgressTypes.ERROR:
|
||||
# Update sub progress bar status
|
||||
desc = data.get('desc')
|
||||
if desc: bar.set_description(desc)
|
||||
bar.refresh()
|
||||
|
||||
# Close sub progress bar and release slot
|
||||
bar.close() # 因为 leave=False，这行会从屏幕消失
|
||||
slot = bar_info['slot']
|
||||
self._release_slot(slot)
|
||||
del self.active_bars[task_id]
|
||||
|
||||
# Update the main progress bar
|
||||
main_bar.update(1)
|
||||
main_bar.refresh()
|
||||
|
||||
pending_tasks -= 1
|
||||
|
||||
# After all tasks are finished, close all bars
|
||||
for bar in self.bars.values():
|
||||
bar.close()
|
||||
except Exception as e:
|
||||
pass # Prevent printing errors from disrupting layout
|
||||
|
||||
# Close all remaining bars (theoretically should be empty except main_bar)
|
||||
main_bar.close()
|
||||
for info in self.active_bars.values():
|
||||
info['bar'].close()
|
||||
|
||||
def start_progress_monitor_thread(total_rows):
|
||||
"""
|
||||
|
||||
Reference in New Issue
Block a user