208 lines
9.4 KiB
Python
208 lines
9.4 KiB
Python
"""自适应线程池。
|
|
|
|
用于逐帧 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()
|