feat: add outfit batch runner
This commit is contained in:
@@ -0,0 +1 @@
|
||||
|
||||
@@ -0,0 +1,29 @@
|
||||
from PySide6.QtCore import QObject, Signal, Slot
|
||||
|
||||
from core.outfit_batch import OutfitBatchOptions, OutfitBatchRunner
|
||||
|
||||
|
||||
class AiOutfitWorker(QObject):
|
||||
"""Qt signal wrapper for the core outfit batch runner."""
|
||||
|
||||
log = Signal(str)
|
||||
progress = Signal(int, int, object) # completed, total, OutfitResult
|
||||
finished = Signal(object) # OutfitBatchSummary
|
||||
|
||||
def __init__(self, tasks, generate_func, options=None, parent=None):
|
||||
super().__init__(parent)
|
||||
self._runner = OutfitBatchRunner(
|
||||
tasks=tasks,
|
||||
generate_func=generate_func,
|
||||
options=options or OutfitBatchOptions(),
|
||||
progress_callback=self.progress.emit,
|
||||
log_callback=self.log.emit,
|
||||
)
|
||||
|
||||
@Slot()
|
||||
def run(self):
|
||||
self.finished.emit(self._runner.run())
|
||||
|
||||
@Slot()
|
||||
def stop(self):
|
||||
self._runner.stop()
|
||||
@@ -0,0 +1,235 @@
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Callable, List
|
||||
|
||||
from core.models import OutfitResult, OutfitTask
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
GenerateFunc = Callable[[OutfitTask], OutfitResult]
|
||||
ProgressCallback = Callable[[int, int, OutfitResult], None]
|
||||
LogCallback = Callable[[str], None]
|
||||
|
||||
|
||||
@dataclass
|
||||
class OutfitBatchOptions:
|
||||
"""Batch generation controls for AI outfit tasks."""
|
||||
|
||||
concurrency: int = 1
|
||||
request_interval: float = 0.0
|
||||
task_cooldown: float = 0.0
|
||||
retry_count: int = 0
|
||||
heartbeat_after: float = 30.0
|
||||
heartbeat_interval: float = 30.0
|
||||
|
||||
|
||||
@dataclass
|
||||
class OutfitBatchSummary:
|
||||
"""Result summary for an outfit batch run."""
|
||||
|
||||
results: List[OutfitResult] = field(default_factory=list)
|
||||
total: int = 0
|
||||
success_count: int = 0
|
||||
failure_count: int = 0
|
||||
stopped: bool = False
|
||||
|
||||
|
||||
class RateLimiter:
|
||||
"""Thread-safe minimum interval between request starts."""
|
||||
|
||||
def __init__(self, interval_seconds, now_func=None, sleep_func=None):
|
||||
self.interval_seconds = max(0.0, float(interval_seconds or 0.0))
|
||||
self._now = now_func or time.monotonic
|
||||
self._sleep = sleep_func or time.sleep
|
||||
self._lock = threading.Lock()
|
||||
self._next_time = 0.0
|
||||
|
||||
def wait(self):
|
||||
if self.interval_seconds <= 0:
|
||||
return
|
||||
with self._lock:
|
||||
now = self._now()
|
||||
if now < self._next_time:
|
||||
self._sleep(self._next_time - now)
|
||||
now = self._now()
|
||||
self._next_time = now + self.interval_seconds
|
||||
|
||||
|
||||
class OutfitBatchRunner:
|
||||
"""Run outfit tasks with concurrency, rate limiting, retries and stop."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
tasks,
|
||||
generate_func,
|
||||
options=None,
|
||||
progress_callback=None,
|
||||
log_callback=None,
|
||||
sleep_func=None,
|
||||
time_func=None,
|
||||
):
|
||||
self.tasks = list(tasks)
|
||||
self.generate_func = generate_func
|
||||
self.options = options or OutfitBatchOptions()
|
||||
self.progress_callback = progress_callback
|
||||
self.log_callback = log_callback
|
||||
self._sleep = sleep_func or time.sleep
|
||||
self._time = time_func or time.monotonic
|
||||
self._stop_event = threading.Event()
|
||||
self._limiter = RateLimiter(
|
||||
self.options.request_interval,
|
||||
now_func=self._time,
|
||||
sleep_func=self._sleep,
|
||||
)
|
||||
|
||||
def stop(self):
|
||||
"""Request a gentle stop: no new tasks, in-flight tasks finish."""
|
||||
self._stop_event.set()
|
||||
|
||||
def run(self):
|
||||
total = len(self.tasks)
|
||||
results = []
|
||||
completed = 0
|
||||
max_workers = max(1, int(self.options.concurrency or 1))
|
||||
|
||||
self._log("开始 AI 穿搭批量生成：{} 项，并发 {}".format(total, max_workers))
|
||||
|
||||
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||
task_iter = iter(self.tasks)
|
||||
futures = {}
|
||||
submitted_all = False
|
||||
|
||||
while futures or not submitted_all:
|
||||
while (
|
||||
not self._stop_event.is_set()
|
||||
and len(futures) < max_workers
|
||||
and not submitted_all
|
||||
):
|
||||
try:
|
||||
task = next(task_iter)
|
||||
except StopIteration:
|
||||
submitted_all = True
|
||||
break
|
||||
self._limiter.wait()
|
||||
future = executor.submit(self._run_one_with_retry, task)
|
||||
futures[future] = _RunningTask(
|
||||
task=task,
|
||||
started_at=self._time(),
|
||||
last_heartbeat_at=0.0,
|
||||
)
|
||||
self._log("提交第 {} 行：{}".format(task.row_index, task.product_id))
|
||||
|
||||
if not futures:
|
||||
break
|
||||
|
||||
poll_interval = min(0.2, max(0.01, float(self.options.heartbeat_interval or 0.2)))
|
||||
done, _pending = wait(futures.keys(), timeout=poll_interval, return_when=FIRST_COMPLETED)
|
||||
if not done:
|
||||
self._emit_heartbeats(futures)
|
||||
continue
|
||||
|
||||
for future in done:
|
||||
running = futures.pop(future)
|
||||
try:
|
||||
result = future.result()
|
||||
except Exception as exc:
|
||||
logger.exception("Outfit task failed unexpectedly: row %s", running.task.row_index)
|
||||
result = OutfitResult(
|
||||
task=running.task,
|
||||
success=False,
|
||||
error=str(exc),
|
||||
attempts=max(1, self.options.retry_count + 1),
|
||||
)
|
||||
|
||||
results.append(result)
|
||||
completed += 1
|
||||
if self.progress_callback is not None:
|
||||
self.progress_callback(completed, total, result)
|
||||
|
||||
if self._stop_event.is_set():
|
||||
submitted_all = True
|
||||
|
||||
summary = _summarize(results, total, stopped=self._stop_event.is_set())
|
||||
self._log(
|
||||
"AI 穿搭批量结束：完成 {}，失败 {}，停止 {}".format(
|
||||
summary.success_count,
|
||||
summary.failure_count,
|
||||
"是" if summary.stopped else "否",
|
||||
)
|
||||
)
|
||||
return summary
|
||||
|
||||
def _run_one_with_retry(self, task):
|
||||
attempts_allowed = max(1, int(self.options.retry_count or 0) + 1)
|
||||
last_result = None
|
||||
for attempt in range(1, attempts_allowed + 1):
|
||||
if self._stop_event.is_set() and attempt > 1:
|
||||
break
|
||||
result = self.generate_func(task)
|
||||
result.attempts = attempt
|
||||
if result.success:
|
||||
if self.options.task_cooldown > 0:
|
||||
self._sleep(self.options.task_cooldown)
|
||||
return result
|
||||
last_result = result
|
||||
if attempt < attempts_allowed:
|
||||
delay = retry_delay_seconds(result.error, attempt)
|
||||
self._log("第 {} 行失败，{} 秒后重试：{}".format(task.row_index, delay, result.error))
|
||||
self._sleep(delay)
|
||||
|
||||
if self.options.task_cooldown > 0:
|
||||
self._sleep(self.options.task_cooldown)
|
||||
return last_result or OutfitResult(task=task, success=False, error="未知错误", attempts=0)
|
||||
|
||||
def _emit_heartbeats(self, futures):
|
||||
now = self._time()
|
||||
for running in futures.values():
|
||||
elapsed = now - running.started_at
|
||||
if elapsed < self.options.heartbeat_after:
|
||||
continue
|
||||
if running.last_heartbeat_at and now - running.last_heartbeat_at < self.options.heartbeat_interval:
|
||||
continue
|
||||
running.last_heartbeat_at = now
|
||||
self._log(
|
||||
"第 {} 行等待中转响应，已等待 {} 秒".format(
|
||||
running.task.row_index,
|
||||
int(elapsed),
|
||||
)
|
||||
)
|
||||
|
||||
def _log(self, message):
|
||||
logger.info(message)
|
||||
if self.log_callback is not None:
|
||||
self.log_callback(message)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _RunningTask:
|
||||
task: OutfitTask
|
||||
started_at: float
|
||||
last_heartbeat_at: float = 0.0
|
||||
|
||||
|
||||
def retry_delay_seconds(error, attempt):
|
||||
"""Short stepped retry delay; rate-limit errors back off more."""
|
||||
text = str(error or "").lower()
|
||||
attempt = max(1, int(attempt or 1))
|
||||
if "429" in text or "rate" in text or "limit" in text or "限流" in text:
|
||||
return min(30, 3 * attempt)
|
||||
return min(10, attempt)
|
||||
|
||||
|
||||
def _summarize(results, total, stopped=False):
|
||||
success_count = sum(1 for result in results if result.success)
|
||||
failure_count = len(results) - success_count
|
||||
return OutfitBatchSummary(
|
||||
results=list(results),
|
||||
total=total,
|
||||
success_count=success_count,
|
||||
failure_count=failure_count,
|
||||
stopped=stopped,
|
||||
)
|
||||
@@ -1056,8 +1056,8 @@
|
||||
|
||||
前置阅读：`docs/11-ai-outfit.md`（§8）
|
||||
|
||||
- [ ] `Worker(QObject)`：`ThreadPoolExecutor(并发数)` + `RateLimiter(请求间隔)` + 单任务冷却 + 阶梯重试 + 温和停止 + >30s 心跳；纯逻辑尽量可测
|
||||
- [ ] 子线程只经 signal 回主线程，不直接碰控件
|
||||
- [x] `Worker(QObject)`：`ThreadPoolExecutor(并发数)` + `RateLimiter(请求间隔)` + 单任务冷却 + 阶梯重试 + 温和停止 + >30s 心跳；纯逻辑尽量可测
|
||||
- [x] 子线程只经 signal 回主线程，不直接碰控件
|
||||
|
||||
### 19.3 UI 页签 — docs/11 §14 阶段 3
|
||||
|
||||
|
||||
@@ -0,0 +1,220 @@
|
||||
"""Tests for AI outfit batch orchestration."""
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).parent.parent / "src"))
|
||||
|
||||
|
||||
def _task(row, product_id=None):
|
||||
from core.models import OutfitTask
|
||||
|
||||
return OutfitTask(
|
||||
row_index=row,
|
||||
title="title{}".format(row),
|
||||
product_id=product_id or "TY{:03d}".format(row),
|
||||
garment_path="g{}.png".format(row),
|
||||
)
|
||||
|
||||
|
||||
def _result(task, success=True, error=""):
|
||||
from core.models import OutfitResult
|
||||
|
||||
return OutfitResult(task=task, success=success, error=error)
|
||||
|
||||
|
||||
class TestRateLimiter(unittest.TestCase):
|
||||
def test_wait_enforces_interval(self):
|
||||
from core.outfit_batch import RateLimiter
|
||||
|
||||
now = [0.0]
|
||||
sleeps = []
|
||||
|
||||
def fake_now():
|
||||
return now[0]
|
||||
|
||||
def fake_sleep(seconds):
|
||||
sleeps.append(seconds)
|
||||
now[0] += seconds
|
||||
|
||||
limiter = RateLimiter(2.0, now_func=fake_now, sleep_func=fake_sleep)
|
||||
limiter.wait()
|
||||
limiter.wait()
|
||||
|
||||
self.assertEqual(sleeps, [2.0])
|
||||
|
||||
|
||||
class TestOutfitBatchRunner(unittest.TestCase):
|
||||
def test_run_success_summary_and_progress(self):
|
||||
from core.outfit_batch import OutfitBatchOptions, OutfitBatchRunner
|
||||
|
||||
tasks = [_task(2), _task(3)]
|
||||
progress = []
|
||||
|
||||
runner = OutfitBatchRunner(
|
||||
tasks,
|
||||
generate_func=lambda task: _result(task, True),
|
||||
options=OutfitBatchOptions(concurrency=1),
|
||||
progress_callback=lambda done, total, result: progress.append((done, total, result.task.row_index)),
|
||||
)
|
||||
|
||||
summary = runner.run()
|
||||
|
||||
self.assertEqual(summary.total, 2)
|
||||
self.assertEqual(summary.success_count, 2)
|
||||
self.assertEqual(summary.failure_count, 0)
|
||||
self.assertEqual(progress, [(1, 2, 2), (2, 2, 3)])
|
||||
|
||||
def test_retry_until_success(self):
|
||||
from core.outfit_batch import OutfitBatchOptions, OutfitBatchRunner
|
||||
|
||||
calls = []
|
||||
|
||||
def generate(task):
|
||||
calls.append(task.row_index)
|
||||
if len(calls) == 1:
|
||||
return _result(task, False, "temporary")
|
||||
return _result(task, True)
|
||||
|
||||
runner = OutfitBatchRunner(
|
||||
[_task(2)],
|
||||
generate_func=generate,
|
||||
options=OutfitBatchOptions(concurrency=1, retry_count=1),
|
||||
sleep_func=lambda _seconds: None,
|
||||
)
|
||||
|
||||
summary = runner.run()
|
||||
|
||||
self.assertEqual(len(calls), 2)
|
||||
self.assertEqual(summary.success_count, 1)
|
||||
self.assertEqual(summary.results[0].attempts, 2)
|
||||
|
||||
def test_concurrency_runs_more_than_one_task_at_once(self):
|
||||
from core.outfit_batch import OutfitBatchOptions, OutfitBatchRunner
|
||||
|
||||
active = [0]
|
||||
max_active = [0]
|
||||
lock = threading.Lock()
|
||||
|
||||
def generate(task):
|
||||
with lock:
|
||||
active[0] += 1
|
||||
max_active[0] = max(max_active[0], active[0])
|
||||
time.sleep(0.05)
|
||||
with lock:
|
||||
active[0] -= 1
|
||||
return _result(task, True)
|
||||
|
||||
runner = OutfitBatchRunner(
|
||||
[_task(2), _task(3), _task(4), _task(5)],
|
||||
generate_func=generate,
|
||||
options=OutfitBatchOptions(concurrency=2),
|
||||
)
|
||||
|
||||
summary = runner.run()
|
||||
|
||||
self.assertEqual(summary.success_count, 4)
|
||||
self.assertGreaterEqual(max_active[0], 2)
|
||||
|
||||
def test_task_cooldown_is_applied(self):
|
||||
from core.outfit_batch import OutfitBatchOptions, OutfitBatchRunner
|
||||
|
||||
sleeps = []
|
||||
runner = OutfitBatchRunner(
|
||||
[_task(2)],
|
||||
generate_func=lambda task: _result(task, True),
|
||||
options=OutfitBatchOptions(concurrency=1, task_cooldown=1.5),
|
||||
sleep_func=sleeps.append,
|
||||
)
|
||||
|
||||
runner.run()
|
||||
|
||||
self.assertEqual(sleeps, [1.5])
|
||||
|
||||
def test_failed_after_retries(self):
|
||||
from core.outfit_batch import OutfitBatchOptions, OutfitBatchRunner
|
||||
|
||||
runner = OutfitBatchRunner(
|
||||
[_task(2)],
|
||||
generate_func=lambda task: _result(task, False, "boom"),
|
||||
options=OutfitBatchOptions(concurrency=1, retry_count=2),
|
||||
sleep_func=lambda _seconds: None,
|
||||
)
|
||||
|
||||
summary = runner.run()
|
||||
|
||||
self.assertEqual(summary.failure_count, 1)
|
||||
self.assertEqual(summary.results[0].attempts, 3)
|
||||
|
||||
def test_gentle_stop_allows_inflight_to_finish_without_new_tasks(self):
|
||||
from core.outfit_batch import OutfitBatchOptions, OutfitBatchRunner
|
||||
|
||||
started = []
|
||||
|
||||
def generate(task):
|
||||
started.append(task.row_index)
|
||||
runner.stop()
|
||||
return _result(task, True)
|
||||
|
||||
runner = OutfitBatchRunner(
|
||||
[_task(2), _task(3), _task(4)],
|
||||
generate_func=generate,
|
||||
options=OutfitBatchOptions(concurrency=1),
|
||||
)
|
||||
|
||||
summary = runner.run()
|
||||
|
||||
self.assertEqual(started, [2])
|
||||
self.assertTrue(summary.stopped)
|
||||
self.assertEqual(summary.total, 3)
|
||||
self.assertEqual(len(summary.results), 1)
|
||||
|
||||
def test_heartbeat_log_for_long_running_task(self):
|
||||
from core.outfit_batch import OutfitBatchOptions, OutfitBatchRunner
|
||||
|
||||
logs = []
|
||||
task_started = threading.Event()
|
||||
release = threading.Event()
|
||||
|
||||
def generate(task):
|
||||
task_started.set()
|
||||
release.wait(2)
|
||||
return _result(task, True)
|
||||
|
||||
runner = OutfitBatchRunner(
|
||||
[_task(2)],
|
||||
generate_func=generate,
|
||||
options=OutfitBatchOptions(
|
||||
concurrency=1,
|
||||
heartbeat_after=0.01,
|
||||
heartbeat_interval=0.01,
|
||||
),
|
||||
log_callback=logs.append,
|
||||
)
|
||||
|
||||
thread = threading.Thread(target=runner.run)
|
||||
thread.start()
|
||||
task_started.wait(1)
|
||||
time.sleep(0.08)
|
||||
release.set()
|
||||
thread.join(1)
|
||||
|
||||
self.assertTrue(any("等待中转响应" in line for line in logs))
|
||||
|
||||
def test_retry_delay_rate_limit_errors_back_off_more(self):
|
||||
from core.outfit_batch import retry_delay_seconds
|
||||
|
||||
self.assertGreater(retry_delay_seconds("429 rate limit", 2), retry_delay_seconds("boom", 2))
|
||||
|
||||
|
||||
class TestAiOutfitWorkerImport(unittest.TestCase):
|
||||
def test_worker_symbols_import(self):
|
||||
from app.workers.ai_outfit_worker import AiOutfitWorker
|
||||
|
||||
self.assertIsNotNone(AiOutfitWorker)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user