"""自适应线程池。 用于逐帧 VLM OCR、逐条 LLM 判定等独立 I/O 子任务： - 从 min_workers 起步，每个时间窗口按平均单任务耗时增减目标并发； - 快响应增加 1 个在途任务额度，慢响应减少 1 个，受 min/max_workers 限制； - 限流降低有效上限，发生错误的窗口禁止扩容，干净窗口逐步恢复上限。 执行器按需创建线程并复用；map 只提交目标额度内的任务，不把整批输入压入 执行器队列。缩容立即限制后续提交，已发出的请求允许完成，不强制中断。 结果与进度由 map 所在线程统一收集，worker 只负责处理输入；返回结果保持 输入顺序，worker 异常作为结果交给调用方决定是否重试。 """ from __future__ import annotations import threading import time from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait from typing import Callable def decide( current: int, avg: float, min_workers: int, max_workers: int, fast_threshold: float, slow_threshold: float, ) -> int: """按平均耗时返回目标并发：快则 +1、慢则 -1，达到上下界后保持。""" if avg < fast_threshold and current < max_workers: return current + 1 if avg > slow_threshold and current > min_workers: return current - 1 return current class AdaptiveThreadPool: """有界自适应执行器；允许顺序重复 map，不允许同一实例并行调用 map。""" def __init__( self, worker: Callable, min_workers: int = 1, max_workers: int = 16, window_seconds: float = 10.0, fast_threshold: float = 0.3, slow_threshold: float = 1.0, clock=time.monotonic, on_progress: Callable[[int, int, float, float, int], None] | None = None, ) -> None: """保存 worker、窗口策略及进度回调；clock 可在测试中注入。""" self._worker = worker self.min_workers = max(1, min_workers) self.max_workers = max(self.min_workers, max_workers) self.window_seconds = window_seconds self.fast_threshold = fast_threshold self.slow_threshold = slow_threshold self._clock = clock # 并发目标是提交额度，不能用执行器已创建的线程数判断缩容。 self._target_workers = 0 self._lock = threading.Lock() self._map_lock = threading.Lock() # 保留 cancel 的调用约定：抑制回调，worker 自行检测暂停并返回异常。 # OCR 因此仍能为每个输入得到结果，同时不会产生上万条暂停进度日志。 self._cancel_event = threading.Event() # 有效上限跨 map 保留：LLM 失败条目重试时继续遵守已收紧的配额。 self._effective_max_workers = self.max_workers self._window_failures = 0 self._window_start = clock() # 记录实际提交时的最大在途数量，供监控与测试检查。 self.max_concurrency = 0 self._on_progress = on_progress self._completed = 0 self._total = 0 self._started_at = 0.0 self._elapsed_sum = 0.0 self._window_times: list[float] = [] self._window_avg_time: float | None = None def _run(self, item) -> tuple[object, float]: """执行一次 worker，保留异常对象并记录真实单任务耗时。""" start = self._clock() try: result = self._worker(item) except Exception as exc: result = exc return result, self._clock() - start def _tick(self, elapsed: float) -> None: """收集窗口耗时并调整额度；与 worker 报告限流共用锁，避免决策竞态。""" with self._lock: self._window_times.append(elapsed) now = self._clock() if now - self._window_start < self.window_seconds: return avg = sum(self._window_times) / len(self._window_times) self._window_start = now self._window_times.clear() self._window_avg_time = avg had_failures = self._window_failures > 0 # 有错误的窗口禁止恢复上限或增加并发；干净窗口每次只恢复 1。 if not had_failures and self._effective_max_workers < self.max_workers: self._effective_max_workers += 1 self._window_failures = 0 target = decide( self._target_workers, avg, self.min_workers, self._effective_max_workers, self.fast_threshold, self.slow_threshold, ) if had_failures: target = min(target, self._target_workers) self._target_workers = max( self.min_workers, min(target, self._effective_max_workers) ) def _current_avg_time(self, elapsed_total: float) -> float: """返回最近窗口均值；窗口未满时返回实际单任务耗时均值。 elapsed_total 保留旧调用签名；墙钟时间除以任务数会受并发倍数影响， 因此均值改用 worker 耗时总和计算。 """ if self._window_avg_time is not None: return self._window_avg_time return self._elapsed_sum / max(self._completed, 1) def _resize(self, target: int) -> None: """幂等调整提交额度；不创建退出哨兵，不等待积压队列消费完再缩容。""" with self._lock: self._target_workers = max( self.min_workers, min(target, self._effective_max_workers) ) def cancel(self) -> None: """抑制本批后续进度回调；暂停与输入结果处理仍交给 worker。 下一批 map 重置此标记，保持 OCR 暂停后继续及 LLM 重试的调用约定。 """ self._cancel_event.set() def report_failure(self) -> None: """限流/服务端错误收紧有效上限，仅允许保持或减少当前提交额度。""" with self._lock: self._window_failures += 1 self._effective_max_workers = max( self.min_workers, self._effective_max_workers - 1 ) # 上限 20 -> 19 不意味着当前 1 个任务应扩到 19 个。 self._target_workers = min(self._target_workers, self._effective_max_workers) def map(self, items) -> list: """有界提交并按输入顺序返回结果；上下文退出时回收执行器线程。""" if not self._map_lock.acquire(blocking=False): raise RuntimeError("同一自适应线程池不能同时执行多个 map") try: self._completed = 0 self._total = len(items) self._started_at = self._clock() self._elapsed_sum = 0.0 self._cancel_event.clear() with self._lock: # 新批次重新计时，空闲时间不构成快响应窗口；错误上限仍保留。 self._window_start = self._started_at self._window_times.clear() self._window_avg_time = None self._target_workers = self.min_workers results = [None] * self._total pending = {} iterator = iter(enumerate(items)) exhausted = False with ThreadPoolExecutor(max_workers=self.max_workers) as executor: while True: # 与 report_failure 共用锁：每次提交都依据最新额度。 # pending 包含尚未收集的完成任务，限制只会更保守，不会超额。 with self._lock: while not exhausted and len(pending) < self._target_workers: try: seq, item = next(iterator) except StopIteration: exhausted = True break future = executor.submit(self._run, item) pending[future] = seq self.max_concurrency = max(self.max_concurrency, len(pending)) if not pending: break done, _ = wait(pending, return_when=FIRST_COMPLETED) for future in sorted(done, key=pending.__getitem__): seq = pending.pop(future) result, elapsed = future.result() results[seq] = result self._completed += 1 self._elapsed_sum += elapsed # 单一收集线程串行回调和统计，不再发生 queue.task_done # 因回调异常未执行而使整批永久挂起的问题。 if self._on_progress is not None and not self._cancel_event.is_set(): elapsed_total = max(self._clock() - self._started_at, 1e-9) with self._lock: workers = self._target_workers self._on_progress( self._completed, self._total, self._completed / elapsed_total, self._current_avg_time(elapsed_total), workers, ) self._tick(elapsed) return results finally: self._map_lock.release()