diff --git a/src/app/workers/__init__.py b/src/app/workers/__init__.py new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/src/app/workers/__init__.py @@ -0,0 +1 @@ + diff --git a/src/app/workers/ai_outfit_worker.py b/src/app/workers/ai_outfit_worker.py new file mode 100644 index 0000000..7209d06 --- /dev/null +++ b/src/app/workers/ai_outfit_worker.py @@ -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() diff --git a/src/core/outfit_batch.py b/src/core/outfit_batch.py new file mode 100644 index 0000000..15e3249 --- /dev/null +++ b/src/core/outfit_batch.py @@ -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, + ) diff --git a/tasks.md b/tasks.md index 909d9a6..52d5bf5 100644 --- a/tasks.md +++ b/tasks.md @@ -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 diff --git a/tests/test_outfit_batch.py b/tests/test_outfit_batch.py new file mode 100644 index 0000000..a4bbdee --- /dev/null +++ b/tests/test_outfit_batch.py @@ -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()