diff --git a/.env.example b/.env.example index 32f950df..fa70c0f2 100644 --- a/.env.example +++ b/.env.example @@ -100,16 +100,25 @@ WINDUP_MQ_MAX_PUBLISH_ATTEMPTS=10 WINDUP_MQ_MAX_CONSUME_ATTEMPTS=5 WINDUP_MQ_CONSUME_LEASE_SECONDS=1800 WINDUP_MQ_EMAIL_HANDLER_RETRIES=3 -# worker 内 handler 并行度（单进程内 ThreadPoolExecutor；默认合计 ≤ 连接池容量） +# worker 内 handler 并行度（单进程内 ThreadPoolExecutor） +# 生成期间不占 Postgres 连接，可按机器内存与上游配额上调。 WINDUP_MQ_EMAIL_CONCURRENCY=8 -WINDUP_MQ_GENERATION_IMAGE_CONCURRENCY=6 -WINDUP_MQ_GENERATION_ACTION_CONCURRENCY=4 +WINDUP_MQ_GENERATION_IMAGE_CONCURRENCY=16 +WINDUP_MQ_GENERATION_ACTION_CONCURRENCY=8 +# poll 使用独立线程池,不与图/动作共享槽位 +WINDUP_MQ_GENERATION_POLL_CONCURRENCY=16 +# i2v 轮询走 ZSET 延迟队列,到期促进到 Stream;不用 Redis 过期事件 +WINDUP_MQ_DELAYED_ZSET=windup:zset:delayed +WINDUP_MQ_DELAYED_CLAIM_LIMIT=50 +WINDUP_MQ_DELAYED_TICK_SECONDS=1 +# 上传 / 参考图下载 / 多图 gen_image 共用的 IO 线程池 +WINDUP_IO_POOL_SIZE=32 # 生成任务排队 / 孤儿 RUNNING 超时（秒） WINDUP_GENERATION_PENDING_MAX_AGE=86400 WINDUP_GENERATION_RUNNING_STALE_SECONDS=1800 # SSE 跨进程桥接：worker publish → web subscribe → EventBus WINDUP_SSE_REDIS_CHANNEL=windup:pubsub:generation-task-events -# worker 并发高时建议调大连接池（默认 pool_size=5, max_overflow=10） +# web / worker 各有一份连接池。生成已短 session，不必再按 handler 并发配满。 POSTGRES_POOL_SIZE=5 POSTGRES_MAX_OVERFLOW=10 diff --git a/backend/packages/ai_engine/src/windup_ai_engine/impl/character_generator.py b/backend/packages/ai_engine/src/windup_ai_engine/impl/character_generator.py index 12fc8ab0..39208a96 100644 --- a/backend/packages/ai_engine/src/windup_ai_engine/impl/character_generator.py +++ b/backend/packages/ai_engine/src/windup_ai_engine/impl/character_generator.py @@ -128,6 +128,63 @@ def generate( ) return self._finish(frames, action, route, progress, canvas) + def can_defer_i2v(self) -> bool: + strategy = self._by_route.get(GenRoute.VIDEO_I2V) + video = getattr(strategy, "_video", None) + return hasattr(strategy, "start_job") and hasattr(video, "start_i2v") + + def start_video( + self, + card: CharacterCard, + action: ActionSpec, + master: bytes, + progress: ProgressPort, + canvas: tuple[int, int] | None = None, + ): + """花钱之前的预检 + 建 i2v 单,不跟单。""" + facts = check_master(master, canvas) + progress.step("precheck", _TICK_PRECHECK, _TOTAL, facts.note()) + route = ROUTE_MATRIX[action.action] + progress.step("route", _TICK_ROUTE, _TOTAL, f"{action.action.value} → {route.value}") + strategy = self._pick(route, action) + start_job = getattr(strategy, "start_job", None) + if start_job is None: + raise TypeError(f"{route.value} 不支持异步建单") + return start_job( + card, action, master, _BandProgress(progress, _DERIVE_FROM, _DERIVE_TO) + ) + + def poll_video(self, job_id: str, *, route_id: str | None = None) -> bytes | None: + """walk/idle/attack/custom 共用 VIDEO_I2V 桶,按 job 探一次。""" + strategy = self._by_route.get(GenRoute.VIDEO_I2V) + if strategy is None: + raise NotImplementedError("未注入视频路线,无法轮询 i2v") + poll_job = getattr(strategy, "poll_job", None) + if poll_job is None: + raise TypeError("视频路线不支持 poll_job") + return poll_job(job_id, route_id=route_id) + + def finish_video( + self, + video: bytes, + card: CharacterCard, + action: ActionSpec, + master: bytes, + progress: ProgressPort, + canvas: tuple[int, int] | None = None, + ) -> GeneratedAction: + """成片到位后走抽帧 / 抠图 / 最后一公里。""" + route = ROUTE_MATRIX[action.action] + progress.step("route", _TICK_ROUTE, _TOTAL, f"{action.action.value} → {route.value}") + strategy = self._pick(route, action) + frames_from = getattr(strategy, "frames_from_video", None) + if frames_from is None: + raise TypeError(f"{route.value} 不支持从成片续跑") + frames = frames_from( + video, card, action, master, _BandProgress(progress, _DERIVE_FROM, _DERIVE_TO) + ) + return self._finish(frames, action, route, progress, canvas) + def generate_rendered( self, card: CharacterCard, diff --git a/backend/packages/ai_engine/src/windup_ai_engine/slicing/extract.py b/backend/packages/ai_engine/src/windup_ai_engine/slicing/extract.py index 10665f8b..7fb4547e 100644 --- a/backend/packages/ai_engine/src/windup_ai_engine/slicing/extract.py +++ b/backend/packages/ai_engine/src/windup_ai_engine/slicing/extract.py @@ -8,6 +8,7 @@ from __future__ import annotations +import io import logging import os import tempfile @@ -19,21 +20,18 @@ __all__ = ["extract_frames_bytes", "extract_all_frames_bytes"] +_PYAV = {"plugin": "pyav"} +_PYAV_BYTES = {"plugin": "pyav", "extension": ".mp4"} + def extract_frames_bytes(video: bytes, n: int) -> list[Image.Image]: - """从视频 bytes 均匀抽 ``n`` 帧（供后端 strategy 用，provider 返回的是 bytes）。""" - with tempfile.TemporaryDirectory() as tmp: - path = Path(tmp) / "source.mp4" - path.write_bytes(video) - return _extract_frames(str(path), n) + """从视频 bytes 均匀抽 ``n`` 帧。优先内存解码,不落盘。""" + return _extract_frames(video, n) def extract_all_frames_bytes(video: bytes, cap: int = 150) -> list[Image.Image]: """抽视频全部帧（至多 ``cap``，均匀降采样），供周期检测用。""" - with tempfile.TemporaryDirectory() as tmp: - path = Path(tmp) / "source.mp4" - path.write_bytes(video) - return _extract_frames(str(path), cap) + return _extract_frames(video, cap) def _uniform_indices(total: int, n: int) -> list[int]: @@ -42,7 +40,19 @@ def _uniform_indices(total: int, n: int) -> list[int]: return [round(i * (total - 1) / max(1, m - 1)) for i in range(m)] -def _frame_count(video_path: str) -> int: +def _pyav_source(source: str | bytes) -> tuple[str | io.BytesIO, dict]: + if isinstance(source, (bytes, bytearray)): + return io.BytesIO(source), dict(_PYAV_BYTES) + return source, dict(_PYAV) + + +def _rewind(handle: object) -> None: + seek = getattr(handle, "seek", None) + if seek is not None: + seek(0) + + +def _frame_count(source, kw: dict | None = None) -> int: """帧数。先问容器元数据,不可信时退回逐帧计数(计数不保留帧,内存不涨)。 元数据在 14 段真实 i2v 视频上与实际帧数全部一致(2026-08-10 实测),但不同容器/编码 @@ -51,32 +61,40 @@ def _frame_count(video_path: str) -> int: """ import imageio.v3 as iio + opts = dict(kw or _PYAV) try: - shape = iio.improps(video_path, plugin="pyav").shape + shape = iio.improps(source, **opts).shape if shape and isinstance(shape[0], int) and shape[0] > 0: return shape[0] except Exception: pass - return sum(1 for _ in iio.imiter(video_path, plugin="pyav")) + _rewind(source) + total = sum(1 for _ in iio.imiter(source, **opts)) + _rewind(source) + return total -def _extract_frames(video_path: str, n: int) -> list[Image.Image]: +def _extract_frames(source: str | bytes, n: int) -> list[Image.Image]: """从视频均匀抽 ``n`` 帧。优先 imageio(流式),回退系统 ffmpeg。 **流式而不是一次性读整段**(2026-08-10,机器审 P2):原先走 ``iio.imread`` 会把 ``(T, H, W, C)`` 整个 materialize 出来。实测 121 帧 720p 的真实 i2v 视频峰值 319 MiB,而我们只要其中 8~16 帧;并发 worker 叠加时这是实打实的内存墙。 现在峰值≈保留帧数 × 单帧,与视频长度无关。 + + bytes 入口走内存缓冲,避免把 mp4 写到 overlay/云盘再读回来。 """ + handle, kw = _pyav_source(source) try: import imageio.v3 as iio - total = _frame_count(video_path) + total = _frame_count(handle, kw) if total <= 0: raise RuntimeError("视频无可解码帧") + _rewind(handle) wanted = set(_uniform_indices(total, n)) out: list[Image.Image] = [] - for i, frame in enumerate(iio.imiter(video_path, plugin="pyav")): + for i, frame in enumerate(iio.imiter(handle, **kw)): if i in wanted: # convert 之后原始 ndarray 就可以被回收;不持有 frame 本身。 out.append(Image.fromarray(frame).convert("RGBA")) @@ -89,15 +107,26 @@ def _extract_frames(video_path: str, n: int) -> list[Image.Image]: # 两者都表现为悄悄换用 ffmpeg 分支、产出看着正常的帧。至少留一条日志。 logger.warning("imageio 抽帧失败,回退系统 ffmpeg", exc_info=True) + return _ffmpeg_extract(source, n) + + +def _ffmpeg_extract(source: str | bytes, n: int) -> list[Image.Image]: + """仅 pyav 失败时走。会把整段导出成 PNG,只应是冷路径。""" import glob import subprocess from imageio_ffmpeg import get_ffmpeg_exe with tempfile.TemporaryDirectory() as tmp: + if isinstance(source, str): + video_path = source + else: + video_path = str(Path(tmp) / "source.mp4") + Path(video_path).write_bytes(source) subprocess.run( [get_ffmpeg_exe(), "-y", "-i", video_path, "-vsync", "0", os.path.join(tmp, "f_%04d.png")], - capture_output=True, check=True, + capture_output=True, + check=True, ) files = sorted(glob.glob(os.path.join(tmp, "f_*.png"))) if not files: diff --git a/backend/packages/ai_engine/src/windup_ai_engine/strategy/concrete.py b/backend/packages/ai_engine/src/windup_ai_engine/strategy/concrete.py index 5c9b7e9b..2d77c1ce 100644 --- a/backend/packages/ai_engine/src/windup_ai_engine/strategy/concrete.py +++ b/backend/packages/ai_engine/src/windup_ai_engine/strategy/concrete.py @@ -130,17 +130,57 @@ def derive( master: bytes, progress: ProgressPort, ) -> list[bytes]: - # 帧数直接读契约字段:缺省值已收进 ActionSpec(DEFAULT_N_FRAMES),不再由本层 - # 用 `or 8` 兜底 —— 那等于把契约的缺省值写在实现里,换条 strategy 就换个默认值。 - n = action.n_frames # 进度文案里的枚举一律取 .value:Python 3.11+ 改了 str-mixin 枚举的 __format__, # f"{action.action}" 现在给的是 "ActionType.WALK" 而不是 "walk"(3.12.13 实测), # 而这串字会经 server 变成用户看到的 SSE 进度。 progress.step("derive", 0, 3, f"{action.action.value}: i2v 生成视频") - # 母版按动作预处理:jump 要在顶部补空间,否则角色腾空时头顶顶出视频画面被裁 + video = self._submit_i2v(card, action, master) + return self.frames_from_video(video, card, action, master, progress) + + def start_job( + self, + card: CharacterCard, + action: ActionSpec, + master: bytes, + progress: ProgressPort, + ): + """只建 i2v 单,把 job 交给延迟队列。""" + progress.step("derive", 0, 3, f"{action.action.value}: i2v 提交") framed = prepare_master(master, action.action.value) - video = self._video.i2v(framed, self._build_prompt(action, card.stance), seconds=5) + start = getattr(self._video, "start_i2v", None) + if start is None: + raise TypeError("当前 VideoProvider 不支持 start_i2v") + return start(framed, self._build_prompt(action, card.stance), seconds=5) + + def poll_job(self, job_id: str, *, route_id: str | None = None) -> bytes | None: + """探一次。``None`` = 仍在跑;bytes = 成片;失败抛错。""" + poll = getattr(self._video, "poll_i2v", None) + if poll is None: + raise TypeError("当前 VideoProvider 不支持 poll_i2v") + snap = poll(job_id, route_id=route_id) + if snap.ok and snap.body: + return snap.body + if getattr(snap, "error_type", None) is None: + return None + raise RuntimeError( + f"i2v 轮询失败(HTTP {snap.http_status} {snap.error_type}): " + f"{snap.edge_fingerprint or snap.job_status or ''}" + ) + def _submit_i2v(self, card: CharacterCard, action: ActionSpec, master: bytes) -> bytes: + framed = prepare_master(master, action.action.value) + return self._video.i2v(framed, self._build_prompt(action, card.stance), seconds=5) + + def frames_from_video( + self, + video: bytes, + card: CharacterCard, + action: ActionSpec, + master: bytes, + progress: ProgressPort, + ) -> list[bytes]: + del card + n = action.n_frames dense = extract_all_frames_bytes(video) # 跨动作一致性:用视频首帧(=母版姿态)的角色高当共同定标基准。各动作都从同一母版 # 起手,故此值一致 —— 否则各动作按自己最高帧定标,切状态时角色会忽大忽小。 diff --git a/backend/packages/app/src/windup_app/bootstrap/worker.py b/backend/packages/app/src/windup_app/bootstrap/worker.py index ce6c6a50..f09c3156 100644 --- a/backend/packages/app/src/windup_app/bootstrap/worker.py +++ b/backend/packages/app/src/windup_app/bootstrap/worker.py @@ -9,9 +9,13 @@ from windup_app.server.mq.catalog import all_stream_specs, email_stream_spec, generation_stream_spec from windup_app.server.orchestrator import task_repo -from windup_app.server.orchestrator.executor import run_action_task, run_image_task +from windup_app.server.orchestrator.executor import ( + resume_action_poll, + run_action_task, + run_image_task, +) from windup_app.server.orchestrator.recover import recover_orphaned_generation_tasks -from windup_app.worker.consumer import StreamConsumer, start_relay_loop +from windup_app.worker.consumer import StreamConsumer, start_delayed_loop, start_relay_loop from windup_app.worker.pending_timeout import release_stale_pending_tasks from windup_framework.db import Base, SessionLocal, engine from windup_framework.mq.model import MqMessage # noqa: F401 — register metadata @@ -48,6 +52,7 @@ def main() -> None: _recover_on_start(publisher) relay_pending_messages() release_stale_pending_tasks() + _warmup_local_inference() stop_event = threading.Event() @@ -70,10 +75,12 @@ def _handle_signal(_signum, _frame) -> None: run_image_task=run_image_task, run_action_task=run_action_task, stop_event=stop_event, + resume_action_poll=resume_action_poll, ), ] threads = [consumer.start() for consumer in consumers] relay_thread = start_relay_loop(stop_event) + delayed_thread = start_delayed_loop(stop_event) pending_thread = threading.Thread( target=_pending_timeout_loop, @@ -95,10 +102,31 @@ def _handle_signal(_signum, _frame) -> None: for thread in threads: thread.join(timeout=5) relay_thread.join(timeout=5) + delayed_thread.join(timeout=5) pending_thread.join(timeout=5) logger.info("windup worker 已停止") +def _warmup_local_inference() -> None: + """顺序预热抽帧 / 抠图依赖,把 .so 和 onnx 一次读进页缓存。 + + 首个动作任务再惰性加载时,会和并发 handler、同一块云盘上的 Postgres/Redis + 叠成 IOPS 脉冲。预热失败不挡启动,只记一条日志。 + """ + try: + import imageio.v3 # noqa: F401 + import av # noqa: F401 + except Exception: + logger.warning("抽帧后端预热失败", exc_info=True) + try: + from windup_framework.providers import OnnxU2NetMatteProvider + + OnnxU2NetMatteProvider().warmup() + logger.info("ONNX 抠图会话已预热") + except Exception: + logger.warning("ONNX 预热失败,首个抠图任务会再加载", exc_info=True) + + def _pending_timeout_loop(stop_event: threading.Event) -> None: while not stop_event.wait(timeout=300): try: diff --git a/backend/packages/app/src/windup_app/server/mq/catalog.py b/backend/packages/app/src/windup_app/server/mq/catalog.py index 22bc50d4..25aa0f5e 100644 --- a/backend/packages/app/src/windup_app/server/mq/catalog.py +++ b/backend/packages/app/src/windup_app/server/mq/catalog.py @@ -21,6 +21,7 @@ MSG_TYPE_VERIFICATION_CODE = "verification_code" MSG_TYPE_CHARACTER_IMAGE = "character_image" MSG_TYPE_CHARACTER_ACTION = "character_action" +MSG_TYPE_CHARACTER_ACTION_POLL = "character_action_poll" def _env_int(name: str, default: int) -> int: @@ -46,6 +47,8 @@ def email_stream_spec() -> StreamSpec: def generation_worker_pool_size() -> int: + # image/action 共用一个线程池。poll 走独立执行器,不能加进这个数字, + # 否则 image 占满线程后 poll 只能排队。 return generation_image_concurrency() + generation_action_concurrency() @@ -58,12 +61,17 @@ def generation_stream_spec() -> StreamSpec: def generation_image_concurrency() -> int: - # 默认对齐 .env.example 的 pool_size(5)+max_overflow(10)，长耗时 handler 会占连接。 - return _env_int("WINDUP_MQ_GENERATION_IMAGE_CONCURRENCY", 6) + # executor 已短 session:生成/上传不再占连接,默认不再被 15 连接池卡住。 + return _env_int("WINDUP_MQ_GENERATION_IMAGE_CONCURRENCY", 16) def generation_action_concurrency() -> int: - return _env_int("WINDUP_MQ_GENERATION_ACTION_CONCURRENCY", 4) + return _env_int("WINDUP_MQ_GENERATION_ACTION_CONCURRENCY", 8) + + +def generation_poll_concurrency() -> int: + # 单次 inspect + 偶尔下载,不 sleep,默认高于 action 建单并发。 + return _env_int("WINDUP_MQ_GENERATION_POLL_CONCURRENCY", 16) def all_stream_specs() -> tuple[StreamSpec, StreamSpec]: @@ -79,6 +87,7 @@ def all_stream_specs() -> tuple[StreamSpec, StreamSpec]: "GENERATION_PENDING_MAX_AGE_SECONDS", "GENERATION_RUNNING_STALE_SECONDS", "MSG_TYPE_CHARACTER_ACTION", + "MSG_TYPE_CHARACTER_ACTION_POLL", "MSG_TYPE_CHARACTER_IMAGE", "MSG_TYPE_VERIFICATION_CODE", "SSE_REDIS_CHANNEL", @@ -87,6 +96,7 @@ def all_stream_specs() -> tuple[StreamSpec, StreamSpec]: "email_stream_spec", "generation_action_concurrency", "generation_image_concurrency", + "generation_poll_concurrency", "generation_stream_spec", "generation_worker_pool_size", ] diff --git a/backend/packages/app/src/windup_app/server/mq/i2v_state.py b/backend/packages/app/src/windup_app/server/mq/i2v_state.py new file mode 100644 index 00000000..c77deed4 --- /dev/null +++ b/backend/packages/app/src/windup_app/server/mq/i2v_state.py @@ -0,0 +1,66 @@ +"""动作任务在等 i2v 成片时的 Redis 状态。 + +任务行保持 RUNNING;轮询间隔走 ZSET 延迟队列,不占 action worker。 +""" + +from __future__ import annotations + +import time +from typing import Any + +from windup_framework.db.redis import get_redis + +I2V_KEY_PREFIX = "windup:i2v:" +I2V_STATE_TTL_S = 2 * 3600 +I2V_FIRST_POLL_S = 5.0 +I2V_POLL_INTERVAL_S = 60.0 +I2V_MAX_WAIT_S = 30 * 60 + + +def _key(task_id: int) -> str: + return f"{I2V_KEY_PREFIX}{task_id}" + + +def save_i2v_state( + task_id: int, + *, + job_id: str, + poll_count: int, + next_wait: float, + started_at: float | None = None, + route_id: str = "", + model: str = "", +) -> None: + redis_client = get_redis() + mapping = { + "job_id": job_id, + "poll_count": str(poll_count), + "next_wait": str(next_wait), + "started_at": str(started_at if started_at is not None else time.time()), + "route_id": route_id, + "model": model, + } + redis_client.hset(_key(task_id), mapping=mapping) + redis_client.expire(_key(task_id), I2V_STATE_TTL_S) + + +def load_i2v_state(task_id: int) -> dict[str, Any] | None: + raw = get_redis().hgetall(_key(task_id)) + if not raw: + return None + return { + "job_id": str(raw.get("job_id") or ""), + "poll_count": int(raw.get("poll_count") or 0), + "next_wait": float(raw.get("next_wait") or I2V_FIRST_POLL_S), + "started_at": float(raw.get("started_at") or 0), + "route_id": str(raw.get("route_id") or ""), + "model": str(raw.get("model") or ""), + } + + +def delete_i2v_state(task_id: int) -> None: + get_redis().delete(_key(task_id)) + + +def has_i2v_state(task_id: int) -> bool: + return bool(get_redis().exists(_key(task_id))) diff --git a/backend/packages/app/src/windup_app/server/orchestrator/executor.py b/backend/packages/app/src/windup_app/server/orchestrator/executor.py index 7197e162..2d46946b 100644 --- a/backend/packages/app/src/windup_app/server/orchestrator/executor.py +++ b/backend/packages/app/src/windup_app/server/orchestrator/executor.py @@ -1,7 +1,8 @@ """动作生成后台编排(调 ai_engine)。 -编排链:``mark RUNNING → 取母版 → ai_engine 出帧 → 逐帧上传对象存储 → 写回结果/COMPLETED``。 -异常兜底为 FAILED,不抛。 +编排链:``短 session 标 RUNNING → 取母版 → ai_engine 出帧 → IO 池并行上传 +(与判官重叠) → 短 session 写回 COMPLETED``。异常兜底为 FAILED,不抛。 +生成期间不占 Postgres 连接,以便 worker 同时跑多路用户任务。 **分层**:本模块调 ai_engine,故 web/worker **不得 import 本模块**(否则牵出 ai_engine, 违反"入口层不经 ai_engine 直连"门禁)。由 bootstrap(composition root)import + 注入 @@ -13,12 +14,15 @@ from __future__ import annotations +import contextvars import dataclasses import logging +import os import threading -from collections.abc import Callable +from collections.abc import Callable, Sequence +from concurrent.futures import Future, ThreadPoolExecutor from dataclasses import dataclass -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, TypeVar from sqlalchemy.orm import Session @@ -31,8 +35,9 @@ from windup_framework.gateway.types import Scene from windup_framework.config.quality_gate import settings as gate_settings -from windup_app.server.orchestrator import billing, quality_gate, task_repo +from windup_app.server.orchestrator import billing, i2v_poll, quality_gate, task_repo from windup_app.server.orchestrator._failure import user_message +from windup_app.server.orchestrator.i2v_poll import ActionAwaitingVideo from windup_app.server.orchestrator._fetch import fetch_own_media from windup_app.server.orchestrator.model import ( ActionType, @@ -50,6 +55,87 @@ _ACTION_RESULT = "character_action" # task_repo._deserialize_result 按此标签反序列化 +class _PollSkip(Exception): + """续跑时任务已终态,直接 ACK。""" + + +T = TypeVar("T") +R = TypeVar("R") + + +def _env_int(name: str, default: int) -> int: + raw = os.getenv(name, "").strip() + return int(raw) if raw else default + + +# 进程级 IO 池:上传 / 参考图下载 / 多图 gen_image 共用。 +# 必须从 handler 线程往里 submit,禁止在池内任务再 _io_map(同池会死锁)。 +_IO_POOL_SIZE = max(1, _env_int("WINDUP_IO_POOL_SIZE", 32)) +_io_pool_lock = threading.Lock() +_io_pool: ThreadPoolExecutor | None = None + + +def _shared_io_pool() -> ThreadPoolExecutor: + global _io_pool + if _io_pool is None: + with _io_pool_lock: + if _io_pool is None: + _io_pool = ThreadPoolExecutor( + max_workers=_IO_POOL_SIZE, + thread_name_prefix="windup-io", + ) + return _io_pool + + +def _submit_io(fn: Callable[[T], R], items: Sequence[T]) -> list[Future[R]]: + """把独立 IO 丢进共享池。每个任务一份 Context 快照——同一 Context 不能被两线程同时 enter。""" + pool = _shared_io_pool() + futs: list[Future[R]] = [] + for item in items: + ctx = contextvars.copy_context() + futs.append(pool.submit(ctx.run, fn, item)) + return futs + + +def _io_map(fn: Callable[[T], R], items: Sequence[T]) -> list[R]: + if not items: + return [] + if len(items) == 1: + return [fn(items[0])] + return [fut.result() for fut in _submit_io(fn, items)] + + +def _upload_frames(upload: Callable[[bytes], str], pngs: Sequence[bytes]) -> list[str]: + """并行上传各帧,返回 URL 列表,下标与 ``pngs`` 对齐。 + + 失败语义与串行相同:已成功的 PUT 会留在桶里(孤儿),任务仍 FAILED。 + """ + return _io_map(upload, pngs) + + +def _using_session( + session: Session | None, + factory: Callable[[], Session], + fn: Callable[[Session], T], +) -> T: + """自开的 session 只包住 fn:commit 后立刻 close,生成/上传期间不占连接池。 + + 调用方传入的 session 不提交、不关闭(测试事务)。 + """ + if session is not None: + return fn(session) + owned = factory() + try: + out = fn(owned) + owned.commit() + return out + except Exception: + owned.rollback() + raise + finally: + owned.close() + + def _settle_credit(session: Session, task_id: int, *, success: bool) -> None: """任务终态时结清预付费：成功扣减，失败解冻。""" task = task_repo.get_task(session, task_id) @@ -255,54 +341,62 @@ def run_action_task( """跑一个动作任务;异常兜底为 FAILED,不抛。 先从 ``project`` 取全局约束(朝向/画风/尺寸/方向)再调 ai_engine。``session`` - 缺省时自开一个(后台场景);测试可传入自己的 session。 + 缺省时自开短 session(标 RUNNING / 写终态各一次),生成期间不占连接; + 测试可传入自己的 session,则全程复用、不代为 commit。 """ - own = session is None - session = session or self._make_session() reset = None try: - task_repo.update_status(session, task_id, TaskStatus.RUNNING) - if own: - session.commit() + def _mark_running(s: Session) -> ProjectConstraints: + task_repo.update_status(s, task_id, TaskStatus.RUNNING) + return (self._fetch_constraints or _load_constraints)(s, project_id) - cons = (self._fetch_constraints or _load_constraints)(session, project_id) + cons = _using_session(session, self._make_session, _mark_running) reset = bind_call_context( task_id=str(task_id), start_from_model=_resolve_video_model(input.video_model), ) - result = self._produce_action(input, cons) - task_repo.update_result(session, task_id, _ACTION_RESULT, result) - _settle_credit(session, task_id, success=True) - if own: - session.commit() + result = self._produce_action(input, cons, task_id=task_id) + + def _complete(s: Session) -> None: + task_repo.update_result(s, task_id, _ACTION_RESULT, result) + _settle_credit(s, task_id, success=True) + + _using_session(session, self._make_session, _complete) + except ActionAwaitingVideo: + logger.info("动作任务 %s 已提交 i2v,等待延迟轮询", task_id) except PromptRejected as exc: # 单独捕获而不是落进下面那个兜底:兜底只存 str(exc),``code`` 就丢了,server # 于是分不出"用户改一句话就能过的输入错"和"引擎故障",只能去解析异常文本。 logger.info("动作任务 %s 的描述被措辞门禁拒绝: %s", task_id, exc.code.value) - task_repo.fail_task( - session, task_id, error_message=user_message(exc), - ) - if own: - session.commit() + error_message = user_message(exc) + + def _reject(s: Session) -> None: + task_repo.fail_task(s, task_id, error_message=error_message) + + _using_session(session, self._make_session, _reject) except Exception as exc: # noqa: BLE001 —— 兜底任何生成/上传/网络异常 logger.exception("动作任务 %s 失败", task_id) - session.rollback() - task_repo.fail_task( - session, task_id, error_message=user_message(exc), - ) - _settle_credit(session, task_id, success=False) - if own: - session.commit() + if session is not None: + session.rollback() + error_message = user_message(exc) + + def _fail(s: Session) -> None: + task_repo.fail_task(s, task_id, error_message=error_message) + _settle_credit(s, task_id, success=False) + + _using_session(session, self._make_session, _fail) finally: if reset is not None: reset() - if own: - session.close() # -- 内部 -------------------------------------------------------------- def _produce_action( - self, input: CharacterActionInput, cons: ProjectConstraints + self, + input: CharacterActionInput, + cons: ProjectConstraints, + *, + task_id: int, ) -> dict: """母版 → ai_engine 按项目尺寸出帧 → 逐帧上传 → 组结果 dict。 @@ -391,51 +485,181 @@ def _produce_action( ).generate_rendered(card, action, rigged, progress, canvas=canvas) else: master = (self._fetch_master or self._download_master)(input) - generated = self._get_generator( + gen = self._get_generator( _resolve_video_model(input.video_model), cons.directions - ).generate(card, action, master, progress, canvas=canvas) - - upload = self._upload or self._upload_frame - checked = [_require_size(png, cons.sprite_w, cons.sprite_h) for png in generated.frames] - frames = [ - {"index": i, "image_url": upload(png), "duration_ms": dur} - for i, (png, dur) in enumerate(zip(checked, generated.durations)) - ] - # quality / prompt_version 只落库记账,不在此处据成色改判决:交付/重试是产品 - # 决策,该由读这本账的下游按阈值决定,任务状态仍只反映"生成流程是否跑完"。 - result = { - "type": "character_action", - "action_type": input.action_type.value, - "direction": input.direction.value, - "frames": frames, - "quality": dataclasses.asdict(generated.quality), - "prompt_version": generated.prompt_version, - } - # 落位几何随产物一起交出:消费方要把帧画到画布上、判角色有没有站在地上, - # 而这条线的比例是对齐那一步的实参。前端此前抄了一份 0.92 自己算 —— 两份 - # 常数只要有一次不同步,角色就不站在地上,而没有任何一道会红。 - if generated.geometry is not None: - g = generated.geometry - result["geometry"] = { - "canvas_width": g.canvas_w, - "canvas_height": g.canvas_h, - "anchor": {"x": g.anchor_x, "y": g.anchor_y}, - "foot_y": g.foot_y, - } - # master 为 None 时 review 按"没判"返回 None(三渲二路线没有可比的参照)。 - decision = quality_gate.review( - self._get_judge(), checked, master, _judged_action(input) - ) - if decision is not None: - result["judge"] = decision.as_payload() - if decision.blocked: - # 帧已经生成、已经上传,钱早就花完了。拦在这里的意义只剩"不把坏产物当成 - # 交付物交出去";这也正是拦截档默认关着的原因。 - raise quality_gate.QualityBlocked(decision.problems) - return result + ) + # 注入的测试桩通常只有 generate()。生产装配且 VideoGateway 支持 + # start_i2v 时,建单后把轮询丢进 ZSET,立刻让出 action worker。 + can_defer = getattr(gen, "can_defer_i2v", None) + if hasattr(gen, "start_video") and ( + can_defer() if callable(can_defer) else True + ): + job = gen.start_video(card, action, master, progress, canvas=canvas) + i2v_poll.schedule(task_id, job, poll_count=0) + raise ActionAwaitingVideo + generated = gen.generate(card, action, master, progress, canvas=canvas) + + return self._deliver_generated(generated, input, cons, master) finally: reset_call() + def resume_action_poll( + self, + task_id: int, + input: CharacterActionInput, + project_id: int | None = None, + *, + session: Session | None = None, + ) -> None: + """延迟队列到期后探一次 i2v。仍在跑则再挂单;完成则抽帧交付。""" + reset = None + try: + def _mark_running(s: Session) -> ProjectConstraints: + task = task_repo.get_task(s, task_id) + if task is None: + raise RuntimeError(f"任务 {task_id} 不存在") + if task.status in (TaskStatus.COMPLETED, TaskStatus.FAILED): + raise _PollSkip(f"任务 {task_id} 已终态") + return (self._fetch_constraints or _load_constraints)(s, project_id) + + try: + cons = _using_session(session, self._make_session, _mark_running) + except _PollSkip: + return + + reset = bind_call_context( + task_id=str(task_id), + start_from_model=_resolve_video_model(input.video_model), + ) + gen = self._get_generator( + _resolve_video_model(input.video_model), cons.directions + ) + reset_call = fresh_gateway_request() + try: + outcome = i2v_poll.inspect(task_id, poll_video=gen.poll_video) + if isinstance(outcome, i2v_poll.Waiting): + return + + card, action, canvas = self._action_spec(input, cons) + progress: ProgressPort = _LogProgress() + master = (self._fetch_master or self._download_master)(input) + generated = gen.finish_video( + outcome.video, card, action, master, progress, canvas=canvas + ) + result = self._deliver_generated(generated, input, cons, master) + finally: + reset_call() + i2v_poll.clear(task_id) + + def _complete(s: Session) -> None: + task_repo.update_result(s, task_id, _ACTION_RESULT, result) + _settle_credit(s, task_id, success=True) + + _using_session(session, self._make_session, _complete) + except PromptRejected as exc: + logger.info("动作任务 %s 的描述被措辞门禁拒绝: %s", task_id, exc.code.value) + error_message = user_message(exc) + + def _reject(s: Session) -> None: + task_repo.fail_task(s, task_id, error_message=error_message) + + _using_session(session, self._make_session, _reject) + except Exception as exc: # noqa: BLE001 + logger.exception("动作任务 %s 轮询失败", task_id) + if session is not None: + session.rollback() + error_message = user_message(exc) + + def _fail(s: Session) -> None: + task_repo.fail_task(s, task_id, error_message=error_message) + _settle_credit(s, task_id, success=False) + + _using_session(session, self._make_session, _fail) + finally: + if reset is not None: + reset() + + def _action_spec( + self, input: CharacterActionInput, cons: ProjectConstraints + ) -> tuple[CharacterCard, ActionSpec, tuple[int, int]]: + if cons.directions > 1: + logger.info( + "项目要求 %s 方向，本任务只负责真实源方向 %s；镜像方向由资产层复用", + cons.directions, + input.direction.value, + ) + desc_parts = [input.custom_prompt or "", direction_prompt(input.direction)] + if cons.style: + desc_parts.append(f"Art style: {cons.style}") + card = CharacterCard( + name=f"char-{input.character_id}", + desc=" ".join(desc_parts), + **({"stance": input.stance} if input.stance is not None else {}), + ) + engine_action = _to_engine_action(input.action_type) + extra: dict[str, object] = {} + if engine_action is EngineActionType.CUSTOM: + cyclic = False if input.loop is None else bool(input.loop) + extra = {"custom_action": input.custom_prompt or "", "cyclic": cyclic} + action = ActionSpec( + action=engine_action, + poses=[""] * input.num_frames, + facing=cons.facing, + direction=input.direction, + stylize=cons.stylize, + **extra, + ) + return card, action, (cons.sprite_w, cons.sprite_h) + + def _deliver_generated( + self, + generated, + input: CharacterActionInput, + cons: ProjectConstraints, + master: bytes | None, + ) -> dict: + upload = self._upload or self._upload_frame + checked = [_require_size(png, cons.sprite_w, cons.sprite_h) for png in generated.frames] + # 上传与判官都是网络 IO 且互不依赖:判官读内存 bytes。handler 线程跑 review, + # IO 池跑 PUT,避免 32 帧传完再打一次视觉模型。 + upload_futs = _submit_io(upload, checked) if len(checked) > 1 else None + decision = quality_gate.review( + self._get_judge(), checked, master, _judged_action(input) + ) + urls = ( + [fut.result() for fut in upload_futs] + if upload_futs is not None + else _upload_frames(upload, checked) + ) + frames = [ + {"index": i, "image_url": url, "duration_ms": dur} + for i, (url, dur) in enumerate(zip(urls, generated.durations)) + ] + result = { + "type": "character_action", + "action_type": input.action_type.value, + "direction": input.direction.value, + "frames": frames, + "quality": dataclasses.asdict(generated.quality), + "prompt_version": generated.prompt_version, + } + # 落位几何随产物一起交出:消费方要把帧画到画布上、判角色有没有站在地上, + # 而这条线的比例是对齐那一步的实参。前端此前抄了一份 0.92 自己算 —— 两份 + # 常数只要有一次不同步,角色就不站在地上,而没有任何一道会红。 + if generated.geometry is not None: + g = generated.geometry + result["geometry"] = { + "canvas_width": g.canvas_w, + "canvas_height": g.canvas_h, + "anchor": {"x": g.anchor_x, "y": g.anchor_y}, + "foot_y": g.foot_y, + } + if decision is not None: + result["judge"] = decision.as_payload() + if decision.blocked: + raise quality_gate.QualityBlocked(decision.problems) + return result + def _get_judge(self) -> JudgePort | None: """闸口启用时懒建判官;未启用返回 ``None``,一次调用都不发。""" if self._judge is not None or not gate_settings.enabled: @@ -599,6 +823,7 @@ def __init__( self._upload = upload self._fetch_ref = fetch_ref self._session_factory = session_factory + self._assembly_lock = threading.Lock() def run_image_task( self, @@ -608,44 +833,45 @@ def run_image_task( *, session: Session | None = None, ) -> None: - own = session is None - session = session or self._make_session() reset = None try: - task_repo.update_status(session, task_id, TaskStatus.RUNNING) - if own: - session.commit() - cons = _load_constraints(session, project_id) # 角色图也受项目约束 + def _mark_running(s: Session) -> ProjectConstraints: + task_repo.update_status(s, task_id, TaskStatus.RUNNING) + return _load_constraints(s, project_id) + + cons = _using_session(session, self._make_session, _mark_running) reset = bind_call_context(task_id=str(task_id)) urls, quality = self._produce_image(input, cons) - task_repo.update_result( - session, - task_id, - _IMAGE_RESULT, - { - "type": "character_image", - "direction": input.direction.value, - "image_urls": urls, - "quality": quality, - }, - ) - _settle_credit(session, task_id, success=True) - if own: - session.commit() + + def _complete(s: Session) -> None: + task_repo.update_result( + s, + task_id, + _IMAGE_RESULT, + { + "type": "character_image", + "direction": input.direction.value, + "image_urls": urls, + "quality": quality, + }, + ) + _settle_credit(s, task_id, success=True) + + _using_session(session, self._make_session, _complete) except Exception as exc: # noqa: BLE001 —— 兜底 logger.exception("图片任务 %s 失败", task_id) - session.rollback() - task_repo.fail_task( - session, task_id, error_message=user_message(exc), - ) - _settle_credit(session, task_id, success=False) - if own: - session.commit() + if session is not None: + session.rollback() + error_message = user_message(exc) + + def _fail(s: Session) -> None: + task_repo.fail_task(s, task_id, error_message=error_message) + _settle_credit(s, task_id, success=False) + + _using_session(session, self._make_session, _fail) finally: if reset is not None: reset() - if own: - session.close() def _produce_image( self, input: CharacterImageInput, cons: ProjectConstraints @@ -661,19 +887,39 @@ def _produce_image( refs: list[bytes] = [] has_style_ref = False + def _is_url(url: str) -> bool: + return bool(url) and url.lower() not in ("null", "none", "") + # 1. 角色参考图(用户传入,可选,做角色一致性约束) char_url = (input.reference_image_url or "").strip() - if char_url and char_url.lower() not in ("null", "none", ""): - refs.append(fetch(char_url)) - # 2. 风格参考图(项目级,有 sprite_sample_url 时走图生图模式) style_url = (cons.sprite_sample_url or "").strip() - if style_url and style_url.lower() not in ("null", "none", ""): + want_char = _is_url(char_url) + want_style = _is_url(style_url) + + def _fetch_style(url: str) -> bytes | None: try: - refs.append(fetch(style_url)) - has_style_ref = True + return fetch(url) except Exception: - pass # 风格参考图下载失败不阻断 + return None + + if want_char and want_style: + char_bytes, style_bytes = _io_map( + lambda item: fetch(item[1]) if item[0] == "char" else _fetch_style(item[1]), + (("char", char_url), ("style", style_url)), + ) + refs.append(char_bytes) + if style_bytes is not None: + refs.append(style_bytes) + has_style_ref = True + else: + if want_char: + refs.append(fetch(char_url)) + if want_style: + style_bytes = _fetch_style(style_url) + if style_bytes is not None: + refs.append(style_bytes) + has_style_ref = True # 3. 构建提示词 base = ( @@ -708,14 +954,20 @@ def _produce_image( image_gen = self._get_image() matte = self._get_matte() upload = self._upload or self._upload_image - urls: list[str] = [] - cut: list[Image.Image] = [] - for _ in range(max(1, input.num_images)): + + def _gen_one(_i: int) -> bytes: reset_call = fresh_gateway_request() try: - img = image_gen.gen_image(prompt, refs) + return image_gen.gen_image(prompt, refs) finally: reset_call() + + # 多张候选各自一次付费调用,彼此独立;ContextVar 由 _io_map 拷进 IO 线程。 + raws = _io_map(_gen_one, range(max(1, input.num_images))) + # 抠图走 ONNX,同一会话不能并发 Run;上传再并行。 + cut: list[Image.Image] = [] + pngs: list[bytes] = [] + for img in raws: # 提示词要的是浅灰底,交付的母版却必须是透明底:不在这里抠,灰底会一路带进 # 预览、也会成为下一次图生图的参考底色(#430)。抠在缩放之前 —— u2netp 按模型 # 原始分辨率分割,先缩再抠等于把一半可用像素丢掉再让它猜。 @@ -725,24 +977,31 @@ def _produce_image( # 履约"的字段。模型本身不吃宽高,所以在这里落实。 png = _fit_to(matte.cutout(img), input.width, input.height, smooth=True) cut.append(Image.open(io.BytesIO(png)).convert("RGBA")) - urls.append(upload(png)) + pngs.append(png) + urls = _upload_frames(upload, pngs) # 同一份 alpha 顺手数一次主体数(#427):此前多主体母版要等下一个动作任务才留痕, # 而那时钱已经花在错的母版上了。只记账,不在此处判成败 —— 与动作那条同一立场。 return urls, {"subject_blobs": list(subject_blobs(cut))} def _get_image(self): - if self._image is None: - from windup_framework.gateway import build_image_gateway + if self._image is not None: + return self._image + with self._assembly_lock: + if self._image is None: + from windup_framework.gateway import build_image_gateway - self._image = build_image_gateway() - return self._image + self._image = build_image_gateway() + return self._image def _get_matte(self): - if self._matte is None: - from windup_framework.providers import OnnxU2NetMatteProvider + if self._matte is not None: + return self._matte + with self._assembly_lock: + if self._matte is None: + from windup_framework.providers import OnnxU2NetMatteProvider - self._matte = OnnxU2NetMatteProvider() - return self._matte + self._matte = OnnxU2NetMatteProvider() + return self._matte def _download(self, url: str) -> bytes: # 同 _download_master:参考图 URL 由调用方给,必须走白名单取图。 @@ -771,5 +1030,6 @@ def _make_session(self) -> Session: # 默认执行器(真实依赖);bootstrap 取 run_action_task / run_image_task 注入 app.state executor = ActionTaskExecutor() run_action_task = executor.run_action_task +resume_action_poll = executor.resume_action_poll image_executor = ImageTaskExecutor() run_image_task = image_executor.run_image_task diff --git a/backend/packages/app/src/windup_app/server/orchestrator/i2v_poll.py b/backend/packages/app/src/windup_app/server/orchestrator/i2v_poll.py new file mode 100644 index 00000000..f57c7965 --- /dev/null +++ b/backend/packages/app/src/windup_app/server/orchestrator/i2v_poll.py @@ -0,0 +1,152 @@ +"""动作 i2v 等待态:挂起、探活、启动恢复。 + +Redis hash 与 ZSET 延迟队列的细节留在本模块;executor 只处理 +Completed / AwaitingVideo / Failed。 +""" + +from __future__ import annotations + +import logging +import time +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +from windup_app.server.mq.catalog import GENERATION_STREAM, MSG_TYPE_CHARACTER_ACTION_POLL +from windup_app.server.mq.i2v_state import ( + I2V_FIRST_POLL_S, + I2V_MAX_WAIT_S, + I2V_POLL_INTERVAL_S, + delete_i2v_state, + load_i2v_state, + save_i2v_state, +) +from windup_framework.mq.delayed import schedule_delayed + +logger = logging.getLogger("windup.generation.i2v_poll") + + +class ActionAwaitingVideo(Exception): + """i2v 已建单,任务保持 RUNNING,等 ZSET 到期后再探,不占 action worker。""" + + +@dataclass(frozen=True) +class Waiting: + """仍在上游生成,已重新挂单。""" + + +@dataclass(frozen=True) +class Ready: + video: bytes + route_id: str | None + + +def _job_fields(job: object) -> tuple[str, str, str]: + if isinstance(job, dict): + return ( + str(job.get("job_id") or ""), + str(job.get("route_id") or ""), + str(job.get("model") or ""), + ) + if isinstance(job, str): + return job, "", "" + return ( + str(getattr(job, "job_id", "") or ""), + str(getattr(job, "route_id", "") or ""), + str(getattr(job, "model", "") or ""), + ) + + +def _poll_dedupe(task_id: int, poll_count: int) -> str: + return f"generation:{task_id}:poll:{poll_count}" + + +def _poll_payload(task_id: int, poll_count: int) -> dict[str, Any]: + return { + "task_id": task_id, + "task_type": "character_action", + "poll_count": poll_count, + } + + +def schedule( + task_id: int, + job: object, + *, + poll_count: int, + next_wait: float | None = None, + started_at: float | None = None, +) -> None: + """写入等待态并往延迟队列挂一次探活。""" + job_id, route_id, model = _job_fields(job) + wait = I2V_FIRST_POLL_S if next_wait is None else next_wait + save_i2v_state( + task_id, + job_id=job_id, + poll_count=poll_count, + next_wait=wait, + started_at=started_at, + route_id=route_id, + model=model, + ) + schedule_delayed( + delay_s=wait, + stream=GENERATION_STREAM, + msg_type=MSG_TYPE_CHARACTER_ACTION_POLL, + payload=_poll_payload(task_id, poll_count), + dedupe_key=_poll_dedupe(task_id, poll_count), + ) + + +def inspect( + task_id: int, + *, + poll_video: Callable[..., bytes | None], +) -> Ready | Waiting: + """探一次上游。未完成则再挂单并返回 Waiting;完成则返回 Ready。""" + state = load_i2v_state(task_id) + if state is None or not state.get("job_id"): + raise RuntimeError(f"任务 {task_id} 没有 i2v 状态,无法续跑轮询") + + elapsed = time.time() - float(state["started_at"] or 0) + if elapsed >= I2V_MAX_WAIT_S: + raise RuntimeError("i2v 未取得视频 URL(超时或失败)") + + route_id = state.get("route_id") or None + video = poll_video(state["job_id"], route_id=route_id) + if video is None: + nxt = min(float(state["next_wait"]) * 2, I2V_POLL_INTERVAL_S) + schedule( + task_id, + { + "job_id": state["job_id"], + "route_id": state.get("route_id") or "", + "model": state.get("model") or "", + }, + poll_count=int(state["poll_count"]) + 1, + next_wait=nxt, + started_at=float(state["started_at"]), + ) + return Waiting() + return Ready(video=video, route_id=route_id) + + +def clear(task_id: int) -> None: + delete_i2v_state(task_id) + + +def reschedule_if_waiting(task_id: int, *, delay_s: float = 1) -> bool: + """RUNNING 且仍有等待态:补一条即将到期的探活,不当孤儿失败。""" + state = load_i2v_state(task_id) + if state is None or not state.get("job_id"): + return False + poll_count = int(state.get("poll_count") or 0) + schedule_delayed( + delay_s=delay_s, + stream=GENERATION_STREAM, + msg_type=MSG_TYPE_CHARACTER_ACTION_POLL, + payload=_poll_payload(task_id, poll_count), + dedupe_key=_poll_dedupe(task_id, poll_count), + ) + logger.info("RUNNING 任务仍在等 i2v,已补延迟轮询 | task_id=%s", task_id) + return True diff --git a/backend/packages/app/src/windup_app/server/orchestrator/recover.py b/backend/packages/app/src/windup_app/server/orchestrator/recover.py index 23a3b248..43ae7d92 100644 --- a/backend/packages/app/src/windup_app/server/orchestrator/recover.py +++ b/backend/packages/app/src/windup_app/server/orchestrator/recover.py @@ -21,6 +21,7 @@ MSG_TYPE_CHARACTER_IMAGE, ) from windup_app.server.orchestrator import billing, task_repo +from windup_app.server.orchestrator.i2v_poll import reschedule_if_waiting from windup_app.server.orchestrator.model import ( GenerationTask, GenerationType, @@ -50,6 +51,11 @@ def recover_orphaned_generation_tasks( _fail_unrecoverable(session, task) continue if task.status is TaskStatus.RUNNING: + try: + if reschedule_if_waiting(task.id): + continue + except Exception: + logger.exception("检查 i2v 延迟状态失败 | task_id=%s", task.id) if not fail_stale_running: continue updated = task.update_at diff --git a/backend/packages/app/src/windup_app/worker/consumer.py b/backend/packages/app/src/windup_app/worker/consumer.py index 724dadfa..5eec34a6 100644 --- a/backend/packages/app/src/windup_app/worker/consumer.py +++ b/backend/packages/app/src/windup_app/worker/consumer.py @@ -13,11 +13,14 @@ from typing import Any from windup_app.server.mq.catalog import ( + GENERATION_GROUP, MSG_TYPE_CHARACTER_ACTION, + MSG_TYPE_CHARACTER_ACTION_POLL, MSG_TYPE_CHARACTER_IMAGE, StreamSpec, generation_action_concurrency, generation_image_concurrency, + generation_poll_concurrency, ) from windup_app.worker.handlers import HandlerDeferred, dispatch_handler from windup_framework.db.redis import get_redis @@ -47,18 +50,30 @@ def __init__( run_image_task: Callable[..., Any], run_action_task: Callable[..., Any], stop_event: threading.Event, + resume_action_poll: Callable[..., Any] | None = None, ) -> None: self._config = config self._run_image_task = run_image_task self._run_action_task = run_action_task + self._resume_action_poll = resume_action_poll self._stop = stop_event self._consumer_name = f"{socket.gethostname()}-{threading.get_ident()}" self._executor = ThreadPoolExecutor( max_workers=config.concurrency, thread_name_prefix=f"windup-{config.group}", ) + # poll 必须有自己的线程:共享池里 image 会在 acquire 前占满 worker。 + self._poll_executor = ( + ThreadPoolExecutor( + max_workers=generation_poll_concurrency(), + thread_name_prefix="windup-generation-poll", + ) + if config.group == GENERATION_GROUP + else None + ) self._image_sem = threading.Semaphore(generation_image_concurrency()) self._action_sem = threading.Semaphore(generation_action_concurrency()) + self._poll_sem = threading.Semaphore(generation_poll_concurrency()) self._claim_cursor = "0-0" self._last_claim_at = 0.0 @@ -73,6 +88,8 @@ def start(self) -> threading.Thread: def shutdown(self, *, wait_timeout: float = 30.0) -> None: self._executor.shutdown(wait=True, cancel_futures=False) + if self._poll_executor is not None: + self._poll_executor.shutdown(wait=True, cancel_futures=False) def _loop(self) -> None: redis_client = get_redis() @@ -104,7 +121,7 @@ def _loop(self) -> None: for _stream, messages in batches: for stream_id, fields in messages: - self._executor.submit(self._process_message, stream_id, fields) + self._submit_message(stream_id, fields) def _claim_idle(self, redis_client) -> None: while not self._stop.is_set(): @@ -117,10 +134,25 @@ def _claim_idle(self, redis_client) -> None: ) self._claim_cursor = next_start for stream_id, fields in claimed: - self._executor.submit(self._process_message, stream_id, fields) + self._submit_message(stream_id, fields) if not claimed: break + def _submit_message(self, stream_id: str, fields: dict[str, str]) -> None: + self._executor_for(fields).submit(self._process_message, stream_id, fields) + + def _executor_for(self, fields: dict[str, str]) -> ThreadPoolExecutor: + poll = self._poll_executor + if poll is None: + return self._executor + try: + msg_type = str(mq_client.parse_envelope(fields)["type"]) + except Exception: + return self._executor + if msg_type == MSG_TYPE_CHARACTER_ACTION_POLL: + return poll + return self._executor + def _process_message(self, stream_id: str, fields: dict[str, str]) -> None: redis_client = get_redis() semaphores: list[threading.Semaphore] = [] @@ -165,6 +197,7 @@ def _process_message(self, stream_id: str, fields: dict[str, str]) -> None: payload, run_image_task=self._run_image_task, run_action_task=self._run_action_task, + resume_action_poll=self._resume_action_poll, ) session = SessionLocal() @@ -203,6 +236,8 @@ def _semaphore_for(self, msg_type: str) -> threading.Semaphore | None: return self._image_sem if msg_type == MSG_TYPE_CHARACTER_ACTION: return self._action_sem + if msg_type == MSG_TYPE_CHARACTER_ACTION_POLL: + return self._poll_sem return None def _defer_message(self, message_id: uuid.UUID | None) -> None: @@ -265,3 +300,21 @@ def _run() -> None: thread = threading.Thread(target=_run, name="windup-mq-relay", daemon=True) thread.start() return thread + + +def start_delayed_loop(stop_event: threading.Event) -> threading.Thread: + """把 ZSET 到期项促进到 Stream。不用 keyspace notification。""" + + def _run() -> None: + from windup_framework.mq.config import DELAYED_TICK_SECONDS + from windup_framework.mq.delayed import promote_due_messages + + while not stop_event.wait(timeout=max(1, DELAYED_TICK_SECONDS)): + try: + promote_due_messages() + except Exception: + logger.exception("延迟队列促进失败") + + thread = threading.Thread(target=_run, name="windup-mq-delayed", daemon=True) + thread.start() + return thread diff --git a/backend/packages/app/src/windup_app/worker/handlers.py b/backend/packages/app/src/windup_app/worker/handlers.py index ddc96631..75d1f44b 100644 --- a/backend/packages/app/src/windup_app/worker/handlers.py +++ b/backend/packages/app/src/windup_app/worker/handlers.py @@ -10,6 +10,7 @@ from windup_app.server.mq.catalog import ( EMAIL_HANDLER_RETRIES, MSG_TYPE_CHARACTER_ACTION, + MSG_TYPE_CHARACTER_ACTION_POLL, MSG_TYPE_CHARACTER_IMAGE, MSG_TYPE_VERIFICATION_CODE, ) @@ -133,12 +134,39 @@ def handle_generation( raise ValueError(f"未知生成任务类型: {task_type}") +def handle_action_poll( + payload: dict[str, Any], + *, + resume_action_poll: Callable[..., Any], +) -> None: + """RUNNING 是预期态:建单 worker 已 ACK,本消息只负责探一次。""" + task_id = int(payload["task_id"]) + session = SessionLocal() + try: + task = task_repo.get_task(session, task_id) + if task is None: + logger.warning("轮询任务不存在 | task_id=%d", task_id) + return + if task.status in (TaskStatus.COMPLETED, TaskStatus.FAILED): + logger.info("轮询任务已终态，跳过 | task_id=%d status=%s", task_id, task.status) + return + if not billing.has_open_freeze(session, task_id): + logger.warning("轮询任务无开放冻结，跳过 | task_id=%d", task_id) + return + input_payload = task.input_payload or {} + project_id = task.project_id + finally: + session.close() + resume_action_poll(task_id, _action_input(input_payload), project_id) + + def dispatch_handler( msg_type: str, payload: dict[str, Any], *, run_image_task: Callable[..., Any], run_action_task: Callable[..., Any], + resume_action_poll: Callable[..., Any] | None = None, ) -> None: if msg_type == MSG_TYPE_VERIFICATION_CODE: handle_verification_code(payload) @@ -150,4 +178,9 @@ def dispatch_handler( run_action_task=run_action_task, ) return + if msg_type == MSG_TYPE_CHARACTER_ACTION_POLL: + if resume_action_poll is None: + raise RuntimeError("未注入 resume_action_poll") + handle_action_poll(payload, resume_action_poll=resume_action_poll) + return raise ValueError(f"未知消息类型: {msg_type}") diff --git a/backend/packages/framework/src/windup_framework/gateway/video.py b/backend/packages/framework/src/windup_framework/gateway/video.py index f46fef1e..f3f50c35 100644 --- a/backend/packages/framework/src/windup_framework/gateway/video.py +++ b/backend/packages/framework/src/windup_framework/gateway/video.py @@ -2,7 +2,7 @@ import time import uuid -from dataclasses import replace +from dataclasses import dataclass, replace from datetime import datetime, timezone from windup_common.enums.model import ModelErrorType @@ -29,12 +29,21 @@ hash_bytes, hash_image_input, ) -from windup_framework.gateway.types import NextStep, Scene +from windup_framework.gateway.types import AdapterResult, NextStep, Scene _DEFAULT_RETRY_AFTER_S = 2.0 _SLEEP_CAP_S = 30.0 +@dataclass(frozen=True) +class SubmittedVideoJob: + """建单成功、尚未跟单。``job_id`` 绑在 ``route_id`` 那条路上。""" + + job_id: str + route_id: str + model: str + + def _utc_now() -> str: return datetime.now(timezone.utc).isoformat() @@ -51,13 +60,48 @@ def __init__(self, registry, adapter, circuit, settings, route_adapters=None) -> def _adapter_for(self, route: GatewayRoute): return lookup_adapter(self._route_adapters, route, self._adapter) + def start_i2v( + self, + first_frame: bytes, + prompt: str, + seconds: int = 5, + size: str = "1280x720", + ) -> SubmittedVideoJob: + """只建单,不跟单。把 ``job_id`` 交给延迟队列再 ``poll_i2v``。""" + result = self.i2v(first_frame, prompt, seconds, size, follow=False) + if isinstance(result, SubmittedVideoJob): + return result + raise RuntimeError("video gateway start_i2v 未返回 job") + + def poll_i2v(self, job_id: str, *, route_id: str | None = None) -> AdapterResult: + """单次探活。进行中 ``ok=False`` 且 ``error_type is None``;完成则带 mp4。""" + adapter = self._adapter + if route_id: + for route in self._routes: + if route.route_id == route_id: + adapter = self._adapter_for(route) + break + else: + mapped = self._route_adapters.get(route_id) + if mapped is not None: + adapter = mapped + if hasattr(adapter, "inspect_job"): + snap = adapter.inspect_job(job_id) + if snap.ok and snap.job_status == "completed" and snap.edge_fingerprint: + if hasattr(adapter, "download_completed"): + return adapter.download_completed(job_id, snap.edge_fingerprint) + return snap + return adapter.follow_job(job_id) + def i2v( self, first_frame: bytes, prompt: str, seconds: int = 5, size: str = "1280x720", - ) -> bytes: + *, + follow: bool = True, + ) -> bytes | SubmittedVideoJob: ctx = current_call_context() request_id = ctx.request_id or str(uuid.uuid4()) started = time.monotonic() @@ -187,7 +231,8 @@ def fail(http_status: int | None) -> None: submit_ms = int((time.monotonic() - submit_t0) * 1000) if result.ok and result.job_id: bound_job_id = result.job_id - result = adapter.follow_job(bound_job_id) + if follow: + result = adapter.follow_job(bound_job_id) elif result.ok: result = replace( result, @@ -195,8 +240,12 @@ def fail(http_status: int | None) -> None: error_type=ModelErrorType.INVALID_RESPONSE, body=b"", ) - else: + elif follow: result = adapter.follow_job(bound_job_id) + else: + result = AdapterResult( + ok=True, job_id=bound_job_id, maybe_billed=True + ) ended_at = _utc_now() attempt_latency_ms = int((time.monotonic() - attempt_t0) * 1000) @@ -259,6 +308,12 @@ def fail(http_status: int | None) -> None: model_index=model_index, maybe_billed=maybe_billed, ) + if not follow: + return SubmittedVideoJob( + job_id=bound_job_id or result.job_id or "", + route_id=route.route_id, + model=model, + ) return result.body error_type = result.error_type or ModelErrorType.UNKNOWN diff --git a/backend/packages/framework/src/windup_framework/mq/__init__.py b/backend/packages/framework/src/windup_framework/mq/__init__.py index 3a53232b..61822ec9 100644 --- a/backend/packages/framework/src/windup_framework/mq/__init__.py +++ b/backend/packages/framework/src/windup_framework/mq/__init__.py @@ -1,7 +1,14 @@ -"""Redis Stream 消息队列基础设施。""" +"""Redis Stream 消息队列 + ZSET 延迟队列。""" +from windup_framework.mq.delayed import promote_due_messages, schedule_delayed from windup_framework.mq.model import MqMessage from windup_framework.mq.publisher import MqPublisher from windup_framework.mq.relay import relay_pending_messages -__all__ = ["MqMessage", "MqPublisher", "relay_pending_messages"] +__all__ = [ + "MqMessage", + "MqPublisher", + "promote_due_messages", + "relay_pending_messages", + "schedule_delayed", +] diff --git a/backend/packages/framework/src/windup_framework/mq/config.py b/backend/packages/framework/src/windup_framework/mq/config.py index ef6d3e8b..636a5bfd 100644 --- a/backend/packages/framework/src/windup_framework/mq/config.py +++ b/backend/packages/framework/src/windup_framework/mq/config.py @@ -25,3 +25,6 @@ def _env_int(name: str, default: int) -> int: 30 * 60, ) PEL_CLAIM_INTERVAL_SECONDS = _env_int("WINDUP_MQ_PEL_CLAIM_INTERVAL_SECONDS", 30) +DELAYED_ZSET_KEY = os.getenv("WINDUP_MQ_DELAYED_ZSET", "").strip() or "windup:zset:delayed" +DELAYED_CLAIM_LIMIT = _env_int("WINDUP_MQ_DELAYED_CLAIM_LIMIT", 50) +DELAYED_TICK_SECONDS = _env_int("WINDUP_MQ_DELAYED_TICK_SECONDS", 1) diff --git a/backend/packages/framework/src/windup_framework/mq/delayed.py b/backend/packages/framework/src/windup_framework/mq/delayed.py new file mode 100644 index 00000000..41690f0c --- /dev/null +++ b/backend/packages/framework/src/windup_framework/mq/delayed.py @@ -0,0 +1,149 @@ +"""基于 Redis ZSET 的延迟队列,到期后投进现有 Redis Stream。 + +不用 keyspace notification / pubsub 过期事件:删除时机不确定、消息不持久、 +多实例会广播重复消费。对齐 Redisson DelayedQueue 的做法 —— Sorted Set 的 +score 是到期时间,扫描到期项再入就绪队列。就绪队列就是本仓已有的 Stream +(消费者组 + ACK),不另起一套。 + +参考: https://javaguide.cn/database/redis/redis-delayed-task.html + https://javaguide.cn/database/redis/redis-stream-mq.html +""" + +from __future__ import annotations + +import json +import logging +import time +import uuid +from typing import Any + +from windup_framework.db.redis import get_redis +from windup_framework.mq.config import DELAYED_CLAIM_LIMIT, DELAYED_ZSET_KEY + +logger = logging.getLogger("windup.mq.delayed") + +_PAYLOAD_HASH = f"{DELAYED_ZSET_KEY}:payload" + +# 原子:按 score 取出到期 id,从 ZSET/HASH 删掉,返回 payload。避免两个 worker 同时 +# 扫到同一条再投两次。 +_CLAIM_LUA = """ +local zkey = KEYS[1] +local hkey = KEYS[2] +local now = ARGV[1] +local limit = tonumber(ARGV[2]) +local ids = redis.call('ZRANGEBYSCORE', zkey, '-inf', now, 'LIMIT', 0, limit) +if #ids == 0 then + return {} +end +redis.call('ZREM', zkey, unpack(ids)) +local out = {} +for _, id in ipairs(ids) do + local payload = redis.call('HGET', hkey, id) + redis.call('HDEL', hkey, id) + if payload then + out[#out + 1] = payload + end +end +return out +""" + + +def schedule_delayed( + *, + delay_s: float, + stream: str, + msg_type: str, + payload: dict[str, Any], + dedupe_key: str, +) -> str: + """把一条就绪消息排到 ``now + delay_s`` 再 XADD 到 ``stream``。""" + if delay_s < 0: + delay_s = 0 + item_id = str(uuid.uuid4()) + due = time.time() + delay_s + body = json.dumps( + { + "id": item_id, + "stream": stream, + "msg_type": msg_type, + "payload": payload, + "dedupe_key": dedupe_key, + }, + ensure_ascii=False, + ) + redis_client = get_redis() + pipe = redis_client.pipeline() + pipe.zadd(DELAYED_ZSET_KEY, {item_id: due}) + pipe.hset(_PAYLOAD_HASH, item_id, body) + pipe.execute() + logger.debug( + "延迟入队 | id=%s delay_s=%.2f type=%s dedupe=%s", + item_id, + delay_s, + msg_type, + dedupe_key, + ) + return item_id + + +def claim_due(*, now: float | None = None, limit: int = DELAYED_CLAIM_LIMIT) -> list[dict[str, Any]]: + """取出并删除已到期项。返回反序列化后的就绪信封。""" + redis_client = get_redis() + raw = redis_client.eval( + _CLAIM_LUA, + 2, + DELAYED_ZSET_KEY, + _PAYLOAD_HASH, + str(now if now is not None else time.time()), + str(max(1, limit)), + ) + items: list[dict[str, Any]] = [] + for blob in raw or []: + if isinstance(blob, bytes): + blob = blob.decode() + items.append(json.loads(blob)) + return items + + +def promote_due_messages(*, now: float | None = None, limit: int = DELAYED_CLAIM_LIMIT) -> int: + """把到期项写入 outbox 并 XADD。返回成功促进条数。 + + ``publish_now`` 落库即受理;XADD 失败由既有 relay 补投,这里不把同一条 + 再塞回 ZSET,避免与 outbox 双投。 + """ + from windup_framework.db.session import SessionLocal + from windup_framework.mq.publisher import MqPublisher + + due = claim_due(now=now, limit=limit) + if not due: + return 0 + publisher = MqPublisher() + promoted = 0 + for item in due: + session = SessionLocal() + try: + publisher.publish_now( + session, + stream=str(item["stream"]), + msg_type=str(item["msg_type"]), + payload=dict(item.get("payload") or {}), + dedupe_key=str(item["dedupe_key"]), + ) + promoted += 1 + except Exception: + logger.exception("延迟项促进失败,5s 后重试 | dedupe=%s", item.get("dedupe_key")) + try: + schedule_delayed( + delay_s=5, + stream=str(item["stream"]), + msg_type=str(item["msg_type"]), + payload=dict(item.get("payload") or {}), + dedupe_key=str(item["dedupe_key"]), + ) + except Exception: + logger.exception("延迟项重新入队失败 | dedupe=%s", item.get("dedupe_key")) + finally: + session.close() + if promoted: + logger.info("延迟队列促进 %d 条到 Stream", promoted) + return promoted diff --git a/backend/packages/framework/src/windup_framework/providers/matte.py b/backend/packages/framework/src/windup_framework/providers/matte.py index aa98c46a..dbb41122 100644 --- a/backend/packages/framework/src/windup_framework/providers/matte.py +++ b/backend/packages/framework/src/windup_framework/providers/matte.py @@ -294,6 +294,11 @@ def _get_refine_session(self): ) return self._refine_session + def warmup(self) -> None: + """进程启动时把两个 ONNX 会话装进内存,避免首个任务和 overlay 缺页叠在一起。""" + self._get_session() + self._get_refine_session() + def _get_session(self): if self._session is None: try: diff --git a/backend/packages/framework/src/windup_framework/providers/sufy.py b/backend/packages/framework/src/windup_framework/providers/sufy.py index a5c9dacf..1af24d09 100644 --- a/backend/packages/framework/src/windup_framework/providers/sufy.py +++ b/backend/packages/framework/src/windup_framework/providers/sufy.py @@ -138,6 +138,31 @@ def submit_video( return _transport_result(exc) return self._protocol.parse_submit(resp) + def inspect_job(self, job_id: str) -> AdapterResult: + """单次 GET 任务状态,不 sleep、不下载。 + + completed: ``ok=True``,视频 URL 放 ``edge_fingerprint``(Gateway poll 靠这个 + 去 ``download_completed``)。进行中: ``ok=False``,``error_type is None``。 + """ + with self._client() as client: + resp = _poll_get(client, self._protocol.build_poll(job_id)) + parsed = self._protocol.parse_poll(resp, job_id) + if parsed.error_type is not None: + return parsed + if parsed.ok: + url = parsed.result_url + if not url: + return AdapterResult( + ok=False, + error_type=ModelErrorType.INVALID_RESPONSE, + job_id=job_id, + maybe_billed=True, + job_status="completed", + edge_fingerprint="completed 但没有视频 URL", + ) + return replace(parsed, edge_fingerprint=url) + return parsed + def follow_job(self, job_id: str) -> AdapterResult: """轮询已建单据 + 下载。poll GET 522/525 该次再试 1 次,不新开单。""" poll_t0 = time.monotonic() @@ -153,74 +178,70 @@ def with_poll( download_ms=download_ms, ) - with self._client() as client: - url = None - last_status: str | None = None - # 先短后长,而不是每次都睡满 ``poll_interval``。此前第一次查询也要等满一个 - # 间隔:60 秒的间隔下,一段 20 秒就绪的视频要到第 60 秒才被发现,纯白等。 - # 退避到上限后与原来一致,所以对慢任务不增加网关压力。 - # 次数与时间双上限。只用时间会让"永不完成"这类用例必须真等满预算 - # (实测把一条 0.01 秒的用例拖成 60 秒);只用次数则退避变快之后预算被提前 - # 耗光。两者取先到的那个。 - budget = max(1, int(self._max_min * 60 // self._poll)) - deadline = time.monotonic() + self._max_min * 60 - wait = self._first_poll_after - for _ in range(budget): - if time.monotonic() >= deadline: - break - time.sleep(wait) - wait = min(wait * 2, self._poll) - resp = _poll_get(client, self._protocol.build_poll(job_id)) - poll_count += 1 - parsed = self._protocol.parse_poll(resp, job_id) - if parsed.error_type is not None: - return with_poll(parsed) - last_status = parsed.job_status - if parsed.ok: - url = parsed.result_url - break - poll_ms = int((time.monotonic() - poll_t0) * 1000) - if not url: - return replace( - AdapterResult( - ok=False, - error_type=ModelErrorType.TIMEOUT, - job_id=job_id, - maybe_billed=True, - job_status=last_status or "timeout", - ), - poll_ms=poll_ms, - poll_count=poll_count, - ) - try: - download_t0 = time.monotonic() - body = _download(client, url) - download_ms = int((time.monotonic() - download_t0) * 1000) - except RuntimeError as exc: - return replace( - AdapterResult( - ok=False, - error_type=ModelErrorType.MAYBE_BILLED, - job_id=job_id, - maybe_billed=True, - job_status="completed", - edge_fingerprint=str(exc), - ), - poll_ms=poll_ms, - poll_count=poll_count, - ) + url = None + last_status: str | None = None + # 先短后长,而不是每次都睡满 ``poll_interval``。此前第一次查询也要等满一个 + # 间隔:60 秒的间隔下,一段 20 秒就绪的视频要到第 60 秒才被发现,纯白等。 + # 退避到上限后与原来一致,所以对慢任务不增加网关压力。 + # 次数与时间双上限。只用时间会让"永不完成"这类用例必须真等满预算 + # (实测把一条 0.01 秒的用例拖成 60 秒);只用次数则退避变快之后预算被提前 + # 耗光。两者取先到的那个。 + budget = max(1, int(self._max_min * 60 // self._poll)) + deadline = time.monotonic() + self._max_min * 60 + wait = self._first_poll_after + for _ in range(budget): + if time.monotonic() >= deadline: + break + time.sleep(wait) + wait = min(wait * 2, self._poll) + snap = self.inspect_job(job_id) + poll_count += 1 + last_status = snap.job_status + if snap.ok and snap.job_status == "completed": + url = snap.edge_fingerprint + break + if snap.error_type is not None: + return with_poll(snap) + poll_ms = int((time.monotonic() - poll_t0) * 1000) + if not url: return replace( AdapterResult( - ok=True, - body=body, + ok=False, + error_type=ModelErrorType.TIMEOUT, job_id=job_id, maybe_billed=True, - job_status="completed", + job_status=last_status or "timeout", ), poll_ms=poll_ms, poll_count=poll_count, - download_ms=download_ms, ) + downloaded = self.download_completed(job_id, url) + return with_poll(downloaded, download_ms=downloaded.download_ms) + + def download_completed(self, job_id: str, url: str) -> AdapterResult: + """按 inspect 拿到的 URL 下载 mp4,不轮询。""" + try: + download_t0 = time.monotonic() + with self._client() as client: + body = _download(client, url) + download_ms = int((time.monotonic() - download_t0) * 1000) + except RuntimeError as exc: + return AdapterResult( + ok=False, + error_type=ModelErrorType.MAYBE_BILLED, + job_id=job_id, + maybe_billed=True, + job_status="completed", + edge_fingerprint=str(exc), + ) + return AdapterResult( + ok=True, + body=body, + job_id=job_id, + maybe_billed=True, + job_status="completed", + download_ms=download_ms, + ) def i2v( self, first_frame: bytes, prompt: str, seconds: int = 5, size: str = "1280x720" diff --git a/backend/tests/test_env_example_keys_are_live.py b/backend/tests/test_env_example_keys_are_live.py index b486f75c..e8e5cc10 100644 --- a/backend/tests/test_env_example_keys_are_live.py +++ b/backend/tests/test_env_example_keys_are_live.py @@ -25,6 +25,10 @@ "WINDUP_MQ_MAX_CONSUME_ATTEMPTS", "WINDUP_MQ_CONSUME_LEASE_SECONDS", "WINDUP_MQ_EMAIL_HANDLER_RETRIES", "WINDUP_MQ_EMAIL_CONCURRENCY", "WINDUP_MQ_GENERATION_IMAGE_CONCURRENCY", "WINDUP_MQ_GENERATION_ACTION_CONCURRENCY", + "WINDUP_MQ_GENERATION_POLL_CONCURRENCY", + "WINDUP_MQ_DELAYED_ZSET", "WINDUP_MQ_DELAYED_CLAIM_LIMIT", + "WINDUP_MQ_DELAYED_TICK_SECONDS", + "WINDUP_IO_POOL_SIZE", "WINDUP_GENERATION_PENDING_MAX_AGE", "WINDUP_GENERATION_RUNNING_STALE_SECONDS", "WINDUP_SSE_REDIS_CHANNEL", # render3d._tc3.TencentCredentials.resolve:环境变量 → 加锁文件,不走 BaseSettings diff --git a/backend/tests/test_extract_streaming.py b/backend/tests/test_extract_streaming.py index 715bd45d..172fc2cb 100644 --- a/backend/tests/test_extract_streaming.py +++ b/backend/tests/test_extract_streaming.py @@ -104,10 +104,12 @@ def test_asking_for_more_frames_than_the_video_has_returns_all(video): def test_bytes_entry_point_streams_too(video, monkeypatch): - """公开入口是 bytes 版,它也必须走流式(它只是多包了一层临时文件)。""" + """公开入口是 bytes 版,它也必须走流式,且 pyav 成功时不落盘。""" import imageio.v3 as iio + import windup_ai_engine.slicing.extract as extract_mod monkeypatch.setattr(iio, "imread", _forbidden) + monkeypatch.setattr(extract_mod.tempfile, "TemporaryDirectory", _forbidden) with open(video, "rb") as f: assert len(extract_frames_bytes(f.read(), 4)) == 4 diff --git a/backend/tests/test_gateway_video.py b/backend/tests/test_gateway_video.py index d00ff516..43ca558e 100644 --- a/backend/tests/test_gateway_video.py +++ b/backend/tests/test_gateway_video.py @@ -289,6 +289,51 @@ def test_success_is_not_rejected_when_maybe_billed_budget_full(monkeypatch): assert body.startswith(b"\x00\x00\x00\x18ftyp") +def test_start_i2v_returns_job_without_follow(): + ad = FakeVideoAdapter( + submits={ + "kling-v2-5-turbo": [AdapterResult(ok=True, job_id="j-start", maybe_billed=True)], + "kling-v2-6": [], + }, + follows={"j-start": MP4}, + ) + job = _video_gw(ad).start_i2v(b"frame", "walk") + assert job.job_id == "j-start" + assert job.model == "kling-v2-5-turbo" + assert ad.followed == [] + + +def test_poll_i2v_pending_then_download(): + class _PollAdapter(FakeVideoAdapter): + def inspect_job(self, job_id): + self.followed.append(f"inspect:{job_id}") + return AdapterResult( + ok=True, + job_id=job_id, + job_status="completed", + edge_fingerprint="https://cdn.example.com/out.mp4", + ) + + def download_completed(self, job_id, url): + return AdapterResult(ok=True, body=b"mp4-bytes", job_id=job_id, job_status="completed") + + ad = _PollAdapter(submits={}, follows={}) + result = _video_gw(ad).poll_i2v("j-start") + assert result.ok + assert result.body == b"mp4-bytes" + + +def test_poll_i2v_still_pending(): + class _Pending(FakeVideoAdapter): + def inspect_job(self, job_id): + return AdapterResult(ok=False, job_id=job_id, maybe_billed=True, job_status="in_progress") + + result = _video_gw(_Pending(submits={}, follows={})).poll_i2v("j1") + assert not result.ok + assert result.error_type is None + assert result.job_status == "in_progress" + + def test_submit_ok_without_job_id_is_invalid_response(): no_id = AdapterResult(ok=True, body=b"") ad = FakeVideoAdapter( diff --git a/backend/tests/test_generation_orchestration.py b/backend/tests/test_generation_orchestration.py index 6c324f98..04b60668 100644 --- a/backend/tests/test_generation_orchestration.py +++ b/backend/tests/test_generation_orchestration.py @@ -7,6 +7,7 @@ from __future__ import annotations import io +import threading import pytest from PIL import Image @@ -24,7 +25,12 @@ GenerationType, TaskStatus, ) -from windup_app.server.orchestrator.executor import ActionTaskExecutor +from windup_app.server.orchestrator.executor import ( + ActionTaskExecutor, + ProjectConstraints, + _upload_frames, +) +from windup_ai_engine.ports import ActionQuality, GeneratedAction from windup_app.server.orchestrator.service import AiGenerationService from windup_ai_engine.impl import CharacterGenerator from windup_ai_engine.strategy.concrete import VideoFrameStrategy @@ -89,12 +95,88 @@ def _real_offline_generator(monkeypatch) -> CharacterGenerator: ) +def test_upload_frames_preserves_index_when_later_frames_finish_first(): + """pool.map 必须按下标对齐;完成顺序乱了也不能把 URL 填错格。""" + finished: list[int] = [] + lock = threading.Lock() + first_may_finish = threading.Event() + + def _upload(png: bytes) -> str: + i = png[0] + if i == 0: + assert first_may_finish.wait(timeout=2) + with lock: + finished.append(i) + if i != 0: + first_may_finish.set() + return f"https://cdn.example.com/f{i}.png" + + urls = _upload_frames(_upload, [bytes([i]) for i in range(4)]) + assert urls == [f"https://cdn.example.com/f{i}.png" for i in range(4)] + assert finished and finished[0] != 0 + + +def test_upload_frames_single_does_not_need_a_pool(): + assert _upload_frames(lambda png: f"u-{png.decode()}", [b"a"]) == ["u-a"] + assert _upload_frames(lambda png: "u", []) == [] + + +def test_action_task_opens_a_fresh_session_after_generate(session_factory): + """生成期间必须把连接还回池,否则 16+8 路并发会把 15 连接打满。""" + opens: list[int] = [] + inner = session_factory + + def counting(): + opens.append(1) + return inner() + + class _Gen: + def generate(self, *args, **kwargs): + assert len(opens) == 1, "produce 开始时 RUNNING 那次 session 应已关闭并归还" + png = _tiny_png() + return GeneratedAction( + frames=[png], + durations=[80], + quality=ActionQuality( + motion_scale=0.0, + dead_frames=(), + loop_seam=None, + limbs={}, + subject_blobs=(1,), + ), + prompt_version="test", + ) + + service = AiGenerationService() + executor = ActionTaskExecutor( + generator=_Gen(), + upload=lambda png: "https://cdn.example.com/f.png", + fetch_master=lambda _input: _tiny_png(), + session_factory=counting, + fetch_constraints=lambda _s, _pid: ProjectConstraints(sprite_w=64, sprite_h=96), + ) + action_input = CharacterActionInput( + character_id=1, action_type=ActionType.WALK, num_frames=1, + ) + with session_factory() as s: + task = service.generate_character_action(s, user_id=1, input=action_input) + s.commit() + task_id = task.id + executor.run_action_task(task_id, action_input) + assert len(opens) == 2 + + def test_action_task_runs_end_to_end(session_factory, monkeypatch): uploaded: list[bytes] = [] + n = 0 + lock = threading.Lock() def _upload(png: bytes) -> str: - uploaded.append(png) - return f"https://cdn.example.com/frame-{len(uploaded)}.png" + nonlocal n + with lock: + n += 1 + uploaded.append(png) + return f"https://cdn.example.com/frame-{n}.png" service = AiGenerationService() executor = ActionTaskExecutor( @@ -468,3 +550,161 @@ def test_engine_frame_of_wrong_size_fails_instead_of_being_rescaled(session_fact assert done.status is TaskStatus.FAILED, "尺寸对不上必须失败,不能悄悄缩放交付" # 尺寸不符是生成侧缺陷、用户改不动,故用户侧只给通用文案;期望尺寸留在日志里。 assert done.error_message and "512" not in done.error_message + + +def test_action_task_parks_i2v_and_stays_running(session_factory, monkeypatch): + parked: list[dict] = [] + + class _AsyncGen: + def start_video(self, *args, **kwargs): + return {"job_id": "job-park", "route_id": "primary", "model": "kling-v2-5-turbo"} + + def generate(self, *args, **kwargs): + raise AssertionError("有 start_video 时不该走阻塞 generate") + + monkeypatch.setattr( + "windup_app.server.orchestrator.i2v_poll.schedule", + lambda task_id, job, **kw: parked.append({"task_id": task_id, "job": job, **kw}), + ) + service = AiGenerationService() + executor = ActionTaskExecutor( + generator=_AsyncGen(), + fetch_master=lambda _input: _tiny_png(), + session_factory=session_factory, + fetch_constraints=lambda _s, _pid: ProjectConstraints(sprite_w=64, sprite_h=96), + ) + action_input = CharacterActionInput( + character_id=1, action_type=ActionType.WALK, num_frames=1, + ) + with session_factory() as s: + task = service.generate_character_action(s, user_id=1, input=action_input) + s.commit() + task_id = task.id + executor.run_action_task(task_id, action_input) + assert parked and parked[0]["job"]["job_id"] == "job-park" + with session_factory() as s: + done = service.get_task(s, project_id=1, task_id=task_id) + assert done.status is TaskStatus.RUNNING + + +def test_resume_action_poll_finishes_when_video_ready(session_factory, monkeypatch): + png = _tiny_png() + + class _AsyncGen: + def poll_video(self, job_id, route_id=None): + assert job_id == "job-ready" + return b"mp4" + + def finish_video(self, video, *args, **kwargs): + assert video == b"mp4" + return GeneratedAction( + frames=[png], + durations=[80], + quality=ActionQuality( + motion_scale=0.2, + dead_frames=(), + loop_seam=None, + limbs={}, + subject_blobs=(1,), + ), + prompt_version="test", + ) + + import time as _time + + monkeypatch.setattr( + "windup_app.server.orchestrator.i2v_poll.load_i2v_state", + lambda task_id: { + "job_id": "job-ready", + "poll_count": 0, + "next_wait": 5.0, + "started_at": _time.time(), + "route_id": "primary", + "model": "kling-v2-5-turbo", + }, + ) + monkeypatch.setattr("windup_app.server.orchestrator.i2v_poll.clear", lambda task_id: None) + service = AiGenerationService() + executor = ActionTaskExecutor( + generator=_AsyncGen(), + upload=lambda _png: "https://cdn.example.com/f.png", + fetch_master=lambda _input: png, + session_factory=session_factory, + fetch_constraints=lambda _s, _pid: ProjectConstraints(sprite_w=64, sprite_h=96), + ) + action_input = CharacterActionInput( + character_id=1, action_type=ActionType.WALK, num_frames=1, + ) + with session_factory() as s: + task = service.generate_character_action(s, user_id=1, input=action_input) + s.commit() + task_id = task.id + from windup_app.server.orchestrator import task_repo + task_repo.update_status(s, task_id, TaskStatus.RUNNING) + s.commit() + executor.resume_action_poll(task_id, action_input) + with session_factory() as s: + done = service.get_task(s, project_id=1, task_id=task_id) + assert done.status is TaskStatus.COMPLETED + assert done.result.frames[0].image_url.startswith("https://") + + +def test_resume_action_poll_does_not_fetch_master_while_pending( + session_factory, monkeypatch, +): + """未成片时只 inspect,取母版失败不得把仍在上游跑的任务标 FAILED。""" + fetched: list[int] = [] + parked: list[int] = [] + + class _AsyncGen: + def poll_video(self, job_id, route_id=None): + return None + + def finish_video(self, *args, **kwargs): + raise AssertionError("未成片不该 finish") + + import time as _time + + monkeypatch.setattr( + "windup_app.server.orchestrator.i2v_poll.load_i2v_state", + lambda task_id: { + "job_id": "job-pending", + "poll_count": 0, + "next_wait": 5.0, + "started_at": _time.time(), + "route_id": "primary", + "model": "kling-v2-5-turbo", + }, + ) + monkeypatch.setattr( + "windup_app.server.orchestrator.i2v_poll.schedule", + lambda task_id, job, **kw: parked.append(task_id), + ) + + def _boom(_input): + fetched.append(1) + raise RuntimeError("对象存储瞬时失败") + + service = AiGenerationService() + executor = ActionTaskExecutor( + generator=_AsyncGen(), + fetch_master=_boom, + session_factory=session_factory, + fetch_constraints=lambda _s, _pid: ProjectConstraints(sprite_w=64, sprite_h=96), + ) + action_input = CharacterActionInput( + character_id=1, action_type=ActionType.WALK, num_frames=1, + ) + with session_factory() as s: + task = service.generate_character_action(s, user_id=1, input=action_input) + s.commit() + task_id = task.id + from windup_app.server.orchestrator import task_repo + task_repo.update_status(s, task_id, TaskStatus.RUNNING) + s.commit() + executor.resume_action_poll(task_id, action_input) + assert fetched == [] + assert parked == [task_id] + with session_factory() as s: + done = service.get_task(s, project_id=1, task_id=task_id) + assert done.status is TaskStatus.RUNNING diff --git a/backend/tests/test_ground_contact_anchor.py b/backend/tests/test_ground_contact_anchor.py index 08d55fe9..3d12c587 100644 --- a/backend/tests/test_ground_contact_anchor.py +++ b/backend/tests/test_ground_contact_anchor.py @@ -261,6 +261,7 @@ def _capture_spec(**input_kw) -> ActionSpec: **input_kw, ), ProjectConstraints(sprite_w=64, sprite_h=64), + task_id=1, ) return gen.specs[0] diff --git a/backend/tests/test_i2v_poll.py b/backend/tests/test_i2v_poll.py new file mode 100644 index 00000000..f955eb7d --- /dev/null +++ b/backend/tests/test_i2v_poll.py @@ -0,0 +1,112 @@ +"""i2v 等待态应用服务:挂起 / 探活 / 恢复,不碰 Redis 细节以外的编排。""" + +from __future__ import annotations + +import time + +from windup_app.server.orchestrator.i2v_poll import ( + Ready, + Waiting, + inspect, + reschedule_if_waiting, + schedule, +) + + +def test_schedule_persists_state_and_enqueues_delayed_poll(monkeypatch): + saved: dict = {} + delayed: dict = {} + + monkeypatch.setattr( + "windup_app.server.orchestrator.i2v_poll.save_i2v_state", + lambda task_id, **kw: saved.update({"task_id": task_id, **kw}), + ) + monkeypatch.setattr( + "windup_app.server.orchestrator.i2v_poll.schedule_delayed", + lambda **kw: delayed.update(kw), + ) + + schedule(9, {"job_id": "j1", "route_id": "primary", "model": "m"}, poll_count=0) + + assert saved["task_id"] == 9 + assert saved["job_id"] == "j1" + assert delayed["msg_type"] == "character_action_poll" + assert delayed["payload"]["task_id"] == 9 + assert delayed["dedupe_key"] == "generation:9:poll:0" + + +def test_inspect_returns_waiting_and_reschedules(monkeypatch): + parked: list[int] = [] + monkeypatch.setattr( + "windup_app.server.orchestrator.i2v_poll.load_i2v_state", + lambda task_id: { + "job_id": "j1", + "poll_count": 0, + "next_wait": 5.0, + "started_at": time.time(), + "route_id": "primary", + "model": "m", + }, + ) + monkeypatch.setattr( + "windup_app.server.orchestrator.i2v_poll.schedule", + lambda task_id, job, **kw: parked.append(task_id), + ) + + out = inspect(9, poll_video=lambda *_a, **_k: None) + assert isinstance(out, Waiting) + assert parked == [9] + + +def test_inspect_returns_ready_without_reschedule(monkeypatch): + monkeypatch.setattr( + "windup_app.server.orchestrator.i2v_poll.load_i2v_state", + lambda task_id: { + "job_id": "j1", + "poll_count": 1, + "next_wait": 10.0, + "started_at": time.time(), + "route_id": "primary", + "model": "m", + }, + ) + monkeypatch.setattr( + "windup_app.server.orchestrator.i2v_poll.schedule", + lambda *_a, **_k: (_ for _ in ()).throw(AssertionError("成片不该再挂单")), + ) + + out = inspect(9, poll_video=lambda *_a, **_k: b"mp4") + assert isinstance(out, Ready) + assert out.video == b"mp4" + assert out.route_id == "primary" + + +def test_reschedule_if_waiting_false_without_state(monkeypatch): + monkeypatch.setattr( + "windup_app.server.orchestrator.i2v_poll.load_i2v_state", + lambda task_id: None, + ) + assert reschedule_if_waiting(9) is False + + +def test_reschedule_if_waiting_enqueues_soon(monkeypatch): + delayed: dict = {} + monkeypatch.setattr( + "windup_app.server.orchestrator.i2v_poll.load_i2v_state", + lambda task_id: { + "job_id": "j1", + "poll_count": 3, + "next_wait": 40.0, + "started_at": time.time(), + "route_id": "", + "model": "", + }, + ) + monkeypatch.setattr( + "windup_app.server.orchestrator.i2v_poll.schedule_delayed", + lambda **kw: delayed.update(kw), + ) + + assert reschedule_if_waiting(9) is True + assert delayed["delay_s"] == 1 + assert delayed["dedupe_key"] == "generation:9:poll:3" diff --git a/backend/tests/test_master_cutout.py b/backend/tests/test_master_cutout.py index c1f92172..b30b8efe 100644 --- a/backend/tests/test_master_cutout.py +++ b/backend/tests/test_master_cutout.py @@ -9,6 +9,7 @@ """ import io +import threading import numpy as np import pytest @@ -80,10 +81,20 @@ def _constraints(size: int = 64) -> ProjectConstraints: def _run(png: bytes, *, matte=None, num_images: int = 1, want: int = 64): got: list[bytes] = [] + n = 0 + lock = threading.Lock() + + def _upload(b: bytes) -> str: + nonlocal n + with lock: + n += 1 + got.append(b) + return f"u{n}" + ex = ImageTaskExecutor( image=_Gen(png), matte=matte or _BackgroundMatte(), - upload=lambda b: (got.append(b), f"u{len(got)}")[1], + upload=_upload, ) urls, quality = ex._produce_image( CharacterImageInput(prompt="勇者", width=want, height=want, num_images=num_images), @@ -197,3 +208,31 @@ def test_quality_reaches_the_task_result_over_http(session_factory): assert done.status is TaskStatus.COMPLETED assert done.result.quality == {"subject_blobs": [2]} + + +def test_multi_image_generation_keeps_slot_count_when_calls_overlap(): + """num_images>1 并行打模型,交付条数仍等于请求条数。""" + threads: set[int] = set() + lock = threading.Lock() + started = threading.Barrier(3) + + class _SlowGen: + def gen_image(self, prompt, refs): + del prompt, refs + with lock: + threads.add(threading.get_ident()) + started.wait(timeout=2) + return _master() + + ex = ImageTaskExecutor( + image=_SlowGen(), + matte=_BackgroundMatte(), + upload=lambda b: f"u-{id(b)}", + ) + urls, quality = ex._produce_image( + CharacterImageInput(prompt="勇者", width=64, height=64, num_images=3), + _constraints(), + ) + assert len(urls) == 3 + assert quality["subject_blobs"] == [1, 1, 1] + assert len(threads) > 1 diff --git a/backend/tests/test_mq_catalog.py b/backend/tests/test_mq_catalog.py index 8c8ccee9..7af9f353 100644 --- a/backend/tests/test_mq_catalog.py +++ b/backend/tests/test_mq_catalog.py @@ -11,6 +11,7 @@ email_stream_spec, generation_action_concurrency, generation_image_concurrency, + generation_poll_concurrency, generation_stream_spec, generation_worker_pool_size, ) @@ -28,13 +29,19 @@ def test_generation_stream_spec_aggregates_pool_size(): assert spec.stream == GENERATION_STREAM assert spec.group == GENERATION_GROUP assert spec.concurrency == generation_worker_pool_size() - assert spec.concurrency == generation_image_concurrency() + generation_action_concurrency() - assert spec.concurrency <= 10 + assert spec.concurrency == ( + generation_image_concurrency() + generation_action_concurrency() + ) -def test_default_generation_concurrency_fits_default_pool(): - """默认 handler 并发不应超过 .env.example 中的连接池容量 (5+10)。""" - assert generation_image_concurrency() + generation_action_concurrency() <= 15 +def test_default_generation_concurrency_is_the_documented_floor(monkeypatch): + """短 session 后默认不再被 15 连接卡住;16 图 + 8 动作是面向多用户的底线。""" + monkeypatch.delenv("WINDUP_MQ_GENERATION_IMAGE_CONCURRENCY", raising=False) + monkeypatch.delenv("WINDUP_MQ_GENERATION_ACTION_CONCURRENCY", raising=False) + monkeypatch.delenv("WINDUP_MQ_GENERATION_POLL_CONCURRENCY", raising=False) + assert generation_image_concurrency() == 16 + assert generation_action_concurrency() == 8 + assert generation_poll_concurrency() == 16 def test_all_stream_specs_returns_email_and_generation(): @@ -48,8 +55,10 @@ def test_catalog_respects_env_overrides(monkeypatch): monkeypatch.setenv("WINDUP_MQ_EMAIL_CONCURRENCY", "3") monkeypatch.setenv("WINDUP_MQ_GENERATION_IMAGE_CONCURRENCY", "10") monkeypatch.setenv("WINDUP_MQ_GENERATION_ACTION_CONCURRENCY", "5") + monkeypatch.setenv("WINDUP_MQ_GENERATION_POLL_CONCURRENCY", "4") assert email_stream_spec().concurrency == 3 assert generation_image_concurrency() == 10 assert generation_action_concurrency() == 5 + assert generation_poll_concurrency() == 4 assert generation_stream_spec().concurrency == 15 diff --git a/backend/tests/test_mq_delayed.py b/backend/tests/test_mq_delayed.py new file mode 100644 index 00000000..060596cf --- /dev/null +++ b/backend/tests/test_mq_delayed.py @@ -0,0 +1,129 @@ +"""Redis ZSET 延迟队列:到期促进到 Stream,不用 keyspace notification。""" + +from __future__ import annotations + +import json +from unittest.mock import MagicMock + +from windup_framework.mq.config import DELAYED_ZSET_KEY +from windup_framework.mq.delayed import claim_due, promote_due_messages, schedule_delayed + + +class _MemRedis: + """够跑 ZADD/HSET/CLAIM Lua 语义的内存桩。""" + + def __init__(self) -> None: + self.zset: dict[str, float] = {} + self.hash: dict[str, str] = {} + + def pipeline(self): + return self + + def zadd(self, key, mapping): + assert key == DELAYED_ZSET_KEY + self.zset.update(mapping) + return self + + def hset(self, key, field, value): + self.hash[field] = value + return self + + def execute(self): + return [] + + def eval(self, script, nkeys, zkey, hkey, now, limit): + now_f = float(now) + cap = int(limit) + due = sorted( + (item_id for item_id, score in self.zset.items() if score <= now_f), + key=lambda item_id: self.zset[item_id], + )[:cap] + out = [] + for item_id in due: + del self.zset[item_id] + payload = self.hash.pop(item_id, None) + if payload is not None: + out.append(payload) + return out + + +def test_schedule_and_claim_due_is_atomic(monkeypatch): + redis_mem = _MemRedis() + monkeypatch.setattr("windup_framework.mq.delayed.get_redis", lambda: redis_mem) + + schedule_delayed( + delay_s=0, + stream="windup:stream:generation", + msg_type="character_action_poll", + payload={"task_id": 7, "poll_count": 0}, + dedupe_key="generation:7:poll:0", + ) + assert len(redis_mem.zset) == 1 + first = claim_due(now=10**12) + assert len(first) == 1 + assert first[0]["msg_type"] == "character_action_poll" + assert first[0]["payload"]["task_id"] == 7 + assert claim_due(now=10**12) == [] + + +def test_future_items_are_not_claimed(monkeypatch): + redis_mem = _MemRedis() + monkeypatch.setattr("windup_framework.mq.delayed.get_redis", lambda: redis_mem) + schedule_delayed( + delay_s=30, + stream="windup:stream:generation", + msg_type="character_action_poll", + payload={"task_id": 1}, + dedupe_key="generation:1:poll:0", + ) + assert claim_due(now=0) == [] + assert len(redis_mem.zset) == 1 + + +def test_promote_due_messages_uses_publish_now(monkeypatch): + redis_mem = _MemRedis() + monkeypatch.setattr("windup_framework.mq.delayed.get_redis", lambda: redis_mem) + schedule_delayed( + delay_s=0, + stream="windup:stream:generation", + msg_type="character_action_poll", + payload={"task_id": 9}, + dedupe_key="generation:9:poll:0", + ) + + published: list[dict] = [] + + class _Pub: + def publish_now(self, session, **kwargs): + published.append(kwargs) + return "msg-id" + + monkeypatch.setattr("windup_framework.db.session.SessionLocal", lambda: MagicMock()) + monkeypatch.setattr("windup_framework.mq.publisher.MqPublisher", _Pub) + + assert promote_due_messages(now=10**12) == 1 + assert published[0]["msg_type"] == "character_action_poll" + assert published[0]["dedupe_key"] == "generation:9:poll:0" + assert json.dumps(published[0]["payload"]) + + +def test_promote_failure_requeues(monkeypatch): + redis_mem = _MemRedis() + monkeypatch.setattr("windup_framework.mq.delayed.get_redis", lambda: redis_mem) + schedule_delayed( + delay_s=0, + stream="windup:stream:generation", + msg_type="character_action_poll", + payload={"task_id": 3}, + dedupe_key="generation:3:poll:0", + ) + + class _Boom: + def publish_now(self, session, **kwargs): + raise RuntimeError("db down") + + monkeypatch.setattr("windup_framework.db.session.SessionLocal", lambda: MagicMock()) + monkeypatch.setattr("windup_framework.mq.publisher.MqPublisher", _Boom) + + assert promote_due_messages(now=10**12) == 0 + assert len(redis_mem.zset) == 1 diff --git a/backend/tests/test_mq_worker.py b/backend/tests/test_mq_worker.py index 8ccae794..9b1e22b7 100644 --- a/backend/tests/test_mq_worker.py +++ b/backend/tests/test_mq_worker.py @@ -14,6 +14,7 @@ from conftest import seed_credit_account from windup_app.server.mq.catalog import ( MSG_TYPE_CHARACTER_ACTION, + MSG_TYPE_CHARACTER_ACTION_POLL, MSG_TYPE_CHARACTER_IMAGE, MSG_TYPE_VERIFICATION_CODE, ) @@ -446,6 +447,24 @@ def test_handle_generation_unknown_type_raises(db_session, engine, monkeypatch): ) +def test_dispatch_handler_routes_action_poll(monkeypatch): + called = {"ok": False} + + monkeypatch.setattr( + "windup_app.worker.handlers.handle_action_poll", + lambda payload, **kwargs: called.update(ok=True, payload=payload), + ) + + dispatch_handler( + "character_action_poll", + {"task_id": 9, "task_type": "character_action", "poll_count": 1}, + run_image_task=MagicMock(), + run_action_task=MagicMock(), + resume_action_poll=MagicMock(), + ) + assert called["ok"] is True + + def test_dispatch_handler_routes_character_action(monkeypatch): called = {"ok": False} @@ -720,6 +739,109 @@ def test_consumer_action_message_uses_action_semaphore( redis_mock.xack.assert_called_once() +def test_consumer_routes_poll_to_reserved_executor(): + consumer = StreamConsumer( + ConsumerConfig(stream="windup:stream:generation", group="generation", concurrency=2), + run_image_task=MagicMock(), + run_action_task=MagicMock(), + stop_event=threading.Event(), + ) + try: + poll_fields = { + "data": json.dumps({ + "v": 1, + "id": str(uuid.uuid4()), + "type": MSG_TYPE_CHARACTER_ACTION_POLL, + "payload": {"task_id": 1}, + }) + } + image_fields = { + "data": json.dumps({ + "v": 1, + "id": str(uuid.uuid4()), + "type": MSG_TYPE_CHARACTER_IMAGE, + "payload": {"task_id": 1}, + }) + } + assert consumer._poll_executor is not None + assert consumer._executor_for(poll_fields) is consumer._poll_executor + assert consumer._executor_for(image_fields) is consumer._executor + finally: + consumer.shutdown() + + +def test_poll_message_runs_while_image_workers_are_busy( + engine, worker_session, monkeypatch, +): + monkeypatch.setenv("WINDUP_MQ_GENERATION_IMAGE_CONCURRENCY", "1") + monkeypatch.setenv("WINDUP_MQ_GENERATION_POLL_CONCURRENCY", "1") + _patch_worker_session_local(monkeypatch, engine) + + image_id = uuid.uuid4() + poll_id = uuid.uuid4() + mq_repo.insert_pending( + worker_session, + message_id=image_id, + dedupe_key=f"generation:image:{image_id}", + stream="windup:stream:generation", + msg_type=MSG_TYPE_CHARACTER_IMAGE, + payload={"task_id": 1, "task_type": "character_image"}, + ) + mq_repo.mark_published(worker_session, image_id, "i-0") + mq_repo.insert_pending( + worker_session, + message_id=poll_id, + dedupe_key=f"generation:poll:{poll_id}", + stream="windup:stream:generation", + msg_type=MSG_TYPE_CHARACTER_ACTION_POLL, + payload={"task_id": 2, "task_type": "character_action"}, + ) + mq_repo.mark_published(worker_session, poll_id, "p-0") + worker_session.commit() + + redis_mock = MagicMock() + monkeypatch.setattr("windup_app.worker.consumer.get_redis", lambda: redis_mock) + + image_started = threading.Event() + image_release = threading.Event() + poll_ran = threading.Event() + + def dispatch(msg_type, payload, **kwargs): + if msg_type == MSG_TYPE_CHARACTER_IMAGE: + image_started.set() + assert image_release.wait(timeout=5) + elif msg_type == MSG_TYPE_CHARACTER_ACTION_POLL: + poll_ran.set() + + monkeypatch.setattr("windup_app.worker.consumer.dispatch_handler", dispatch) + + consumer = StreamConsumer( + ConsumerConfig(stream="windup:stream:generation", group="generation", concurrency=1), + run_image_task=MagicMock(), + run_action_task=MagicMock(), + stop_event=threading.Event(), + resume_action_poll=MagicMock(), + ) + try: + consumer._submit_message("i-0", {"data": json.dumps({ + "v": 1, + "id": str(image_id), + "type": MSG_TYPE_CHARACTER_IMAGE, + "payload": {"task_id": 1, "task_type": "character_image"}, + })}) + assert image_started.wait(timeout=5) + consumer._submit_message("p-0", {"data": json.dumps({ + "v": 1, + "id": str(poll_id), + "type": MSG_TYPE_CHARACTER_ACTION_POLL, + "payload": {"task_id": 2, "task_type": "character_action"}, + })}) + assert poll_ran.wait(timeout=5), "poll 应在 image 占满共享池时仍能执行" + finally: + image_release.set() + consumer.shutdown() + + def test_release_stale_pending_tasks_unfreezes(db_session, engine, monkeypatch): from windup_app.server.mq.catalog import GENERATION_PENDING_MAX_AGE_SECONDS from windup_app.server.orchestrator.model import GenerationTaskRecord diff --git a/backend/tests/test_render3d_route_and_assets.py b/backend/tests/test_render3d_route_and_assets.py index 2dc27ea4..8ab1ea16 100644 --- a/backend/tests/test_render3d_route_and_assets.py +++ b/backend/tests/test_render3d_route_and_assets.py @@ -230,6 +230,7 @@ def test_real_server_path_reaches_render3d_when_the_outfit_has_a_model(): model_3d_url="https://cdn.example.com/outfits/hanfu.glb", ), ProjectConstraints(sprite_w=64, sprite_h=64), + task_id=1, ) assert renderer.calls == 1, "三渲二没被走到 —— 路线选择又断了" @@ -256,6 +257,7 @@ def test_real_server_path_stays_on_i2v_when_the_outfit_has_no_model(): outfit_id=OUTFIT, model_3d_url=None, ), ProjectConstraints(sprite_w=64, sprite_h=64), + task_id=1, ) assert renderer.calls == 0 diff --git a/backend/tests/test_sufy_video_download.py b/backend/tests/test_sufy_video_download.py index 87f49f9c..b8280da8 100644 --- a/backend/tests/test_sufy_video_download.py +++ b/backend/tests/test_sufy_video_download.py @@ -737,6 +737,24 @@ def h(request): return h +def test_inspect_job_pending_has_no_error_type(): + p = _video_provider(_i2v_handler({}, statuses=("in_progress",))) + snap = p.inspect_job("job-1") + assert not snap.ok + assert snap.error_type is None + assert snap.job_status == "in_progress" + assert not snap.body + + +def test_inspect_job_completed_puts_url_in_edge_fingerprint(): + p = _video_provider(_i2v_handler({}, statuses=("completed",))) + snap = p.inspect_job("job-1") + assert snap.ok + assert snap.job_status == "completed" + assert snap.edge_fingerprint.endswith("/out.mp4") + assert not snap.body + + def test_follow_job_records_poll_and_download_timings(): p = _video_provider(_i2v_handler({}, statuses=("in_progress", "completed"))) result = p.follow_job("job-1")