feat(ai-studio): add hosted image generation jobs
This commit is contained in:
chengma
@@ -0,0 +1,392 @@
|
||||
"""cmhub hosted generation orchestration for AI image studio jobs."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
import urllib.parse
|
||||
from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
|
||||
|
||||
from . import ai, appconfig, image_studio
|
||||
from .version import APP_VERSION
|
||||
|
||||
|
||||
MAX_CMHUB_IMAGE_STUDIO_WORKERS = 5
|
||||
|
||||
|
||||
class ImageStudioGenerationError(RuntimeError):
|
||||
"""Raised when AI studio hosted generation cannot complete."""
|
||||
|
||||
|
||||
def _notify(callback, event):
|
||||
if callback is None:
|
||||
return
|
||||
try:
|
||||
callback(dict(event))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _runtime(config, cmhub_config_path):
|
||||
return ai._cmhub_runtime(config, "image", cmhub_config_path)
|
||||
|
||||
|
||||
def _generated_kind(job_type):
|
||||
return "generated_detail" if str(job_type) == "detail" else "generated_main"
|
||||
|
||||
|
||||
def _output_path(project, job, image_root):
|
||||
dirs = image_studio.project_image_dirs(image_root, project)
|
||||
filename = "job_%06d_%s.jpg" % (int(job.id), _generated_kind(job.job_type))
|
||||
return os.path.join(dirs["generated"], filename)
|
||||
|
||||
|
||||
def _source_path(source_asset):
|
||||
path = str(getattr(source_asset, "local_path", "") or "").strip()
|
||||
if not path or not os.path.isfile(path):
|
||||
raise ImageStudioGenerationError("源图尚未下载到本地，不能生成")
|
||||
return path
|
||||
|
||||
|
||||
def create_generation_jobs(
|
||||
project_id,
|
||||
source_asset_id,
|
||||
prompt,
|
||||
count,
|
||||
*,
|
||||
job_type="main",
|
||||
path=None,
|
||||
):
|
||||
total = max(0, int(count or 0))
|
||||
if total <= 0:
|
||||
raise ImageStudioGenerationError("生成数量必须大于0")
|
||||
jobs = []
|
||||
for _ in range(total):
|
||||
jobs.append(
|
||||
image_studio.create_job(
|
||||
project_id,
|
||||
source_asset_id=source_asset_id,
|
||||
job_type=job_type,
|
||||
prompt=prompt,
|
||||
generation_source="cmhub",
|
||||
provider="cmhub",
|
||||
path=path,
|
||||
)
|
||||
)
|
||||
return jobs
|
||||
|
||||
|
||||
def generate_image_jobs(
|
||||
project_id,
|
||||
source_asset_id,
|
||||
prompt,
|
||||
count,
|
||||
*,
|
||||
job_type="main",
|
||||
aspect_ratio="1:1",
|
||||
config=None,
|
||||
cmhub_config_path=None,
|
||||
path=None,
|
||||
should_stop=None,
|
||||
on_event=None,
|
||||
):
|
||||
jobs = create_generation_jobs(
|
||||
project_id,
|
||||
source_asset_id,
|
||||
prompt,
|
||||
count,
|
||||
job_type=job_type,
|
||||
path=path,
|
||||
)
|
||||
return run_jobs(
|
||||
jobs,
|
||||
aspect_ratio=aspect_ratio,
|
||||
config=config,
|
||||
cmhub_config_path=cmhub_config_path,
|
||||
path=path,
|
||||
should_stop=should_stop,
|
||||
on_event=on_event,
|
||||
)
|
||||
|
||||
|
||||
def resume_image_jobs(
|
||||
*,
|
||||
project_id=None,
|
||||
aspect_ratio="1:1",
|
||||
config=None,
|
||||
cmhub_config_path=None,
|
||||
path=None,
|
||||
should_stop=None,
|
||||
on_event=None,
|
||||
):
|
||||
jobs = image_studio.list_resumable_jobs(path=path, project_id=project_id)
|
||||
return run_jobs(
|
||||
jobs,
|
||||
aspect_ratio=aspect_ratio,
|
||||
config=config,
|
||||
cmhub_config_path=cmhub_config_path,
|
||||
path=path,
|
||||
should_stop=should_stop,
|
||||
on_event=on_event,
|
||||
)
|
||||
|
||||
|
||||
def run_jobs(
|
||||
jobs,
|
||||
*,
|
||||
aspect_ratio="1:1",
|
||||
config=None,
|
||||
cmhub_config_path=None,
|
||||
path=None,
|
||||
should_stop=None,
|
||||
on_event=None,
|
||||
):
|
||||
cfg = appconfig.load_config() if config is None else config
|
||||
job_list = list(jobs or [])
|
||||
if not job_list:
|
||||
return {"total": 0, "success": 0, "failed": 0, "cancelled": 0, "jobs": []}
|
||||
runtime = _runtime(cfg, cmhub_config_path)
|
||||
ai_cfg = appconfig.ai_config(cfg)
|
||||
image_root = appconfig.image_dir(cfg)
|
||||
max_workers = min(MAX_CMHUB_IMAGE_STUDIO_WORKERS, max(1, int(ai_cfg.get("image_concurrency", 1) or 1)), len(job_list))
|
||||
should_stop = should_stop or (lambda: False)
|
||||
lock = threading.Lock()
|
||||
summary = {"total": len(job_list), "success": 0, "failed": 0, "cancelled": 0, "jobs": []}
|
||||
|
||||
def record(result):
|
||||
with lock:
|
||||
summary["jobs"].append(result)
|
||||
if result.get("status") == "succeeded":
|
||||
summary["success"] += 1
|
||||
elif result.get("status") == "cancelled":
|
||||
summary["cancelled"] += 1
|
||||
else:
|
||||
summary["failed"] += 1
|
||||
|
||||
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||
futures = {
|
||||
executor.submit(
|
||||
_run_one_job,
|
||||
job.id,
|
||||
runtime,
|
||||
cfg,
|
||||
image_root,
|
||||
aspect_ratio,
|
||||
path,
|
||||
should_stop,
|
||||
on_event,
|
||||
): job
|
||||
for job in job_list
|
||||
}
|
||||
while futures:
|
||||
done, _ = wait(set(futures), return_when=FIRST_COMPLETED)
|
||||
for future in done:
|
||||
futures.pop(future)
|
||||
try:
|
||||
record(future.result())
|
||||
except Exception as exc:
|
||||
record({"status": "failed", "error": str(exc)})
|
||||
if should_stop():
|
||||
for future, job in list(futures.items()):
|
||||
if future.cancel():
|
||||
image_studio.update_job_status(job.id, "cancelled", error="用户停止", path=path)
|
||||
futures.pop(future, None)
|
||||
record({"job": job, "status": "cancelled", "error": "用户停止"})
|
||||
return summary
|
||||
|
||||
|
||||
def _run_one_job(job_id, runtime, config, image_root, aspect_ratio, db_path, should_stop, on_event):
|
||||
job = image_studio.get_job(job_id, path=db_path)
|
||||
if job is None:
|
||||
raise ImageStudioGenerationError("AI工场生图任务不存在")
|
||||
if should_stop():
|
||||
updated = image_studio.update_job_status(job.id, "cancelled", error="用户停止", path=db_path)
|
||||
return {"job": updated, "status": "cancelled", "error": "用户停止"}
|
||||
project = image_studio.get_project(job.project_id, path=db_path)
|
||||
source_asset = image_studio.get_asset(job.source_asset_id, path=db_path)
|
||||
if project is None or source_asset is None:
|
||||
updated = image_studio.update_job_status(job.id, "failed", error="项目或源图不存在", path=db_path)
|
||||
return {"job": updated, "status": "failed", "error": "项目或源图不存在"}
|
||||
try:
|
||||
image_studio.update_job_status(job.id, "running", path=db_path)
|
||||
request_result = _submit_or_resume_job(
|
||||
job,
|
||||
source_asset,
|
||||
runtime,
|
||||
config,
|
||||
db_path,
|
||||
should_stop,
|
||||
on_event,
|
||||
)
|
||||
image_studio.update_job_status(job.id, "running", path=db_path)
|
||||
request_result = _poll_job(job.id, request_result["task_id"], runtime, request_result, db_path, should_stop, on_event)
|
||||
out_path = _output_path(project, job, image_root)
|
||||
saved_path = _download_and_save_job_image(request_result, out_path, config, on_event, job.id)
|
||||
asset = image_studio.add_asset(
|
||||
project.id,
|
||||
_generated_kind(job.job_type),
|
||||
remote_url=request_result["image_url"],
|
||||
local_path=saved_path,
|
||||
aspect_ratio=aspect_ratio,
|
||||
parent_asset_id=source_asset.id,
|
||||
prompt=job.prompt,
|
||||
path=db_path,
|
||||
)
|
||||
updated = image_studio.update_job_status(
|
||||
job.id,
|
||||
"succeeded",
|
||||
output_asset_id=asset.id,
|
||||
points_balance=request_result.get("points_balance"),
|
||||
path=db_path,
|
||||
)
|
||||
_notify(on_event, {"job_id": job.id, "step": "job_done", "result": "success"})
|
||||
return {"job": updated, "asset": asset, "status": "succeeded"}
|
||||
except Exception as exc:
|
||||
status = "cancelled" if "停止" in str(exc) else "failed"
|
||||
updated = image_studio.update_job_status(job.id, status, error=str(exc), path=db_path)
|
||||
_notify(on_event, {"job_id": job.id, "step": "job_done", "result": status, "detail": str(exc)})
|
||||
return {"job": updated, "status": status, "error": str(exc)}
|
||||
|
||||
|
||||
def _submit_or_resume_job(job, source_asset, runtime, config, db_path, should_stop, on_event):
|
||||
if job.task_id:
|
||||
_notify(on_event, {"job_id": job.id, "step": "cover_request", "result": "resume", "task_id": job.task_id})
|
||||
return _request_result(job.task_id, runtime, config)
|
||||
_raise_if_stopped(should_stop)
|
||||
source_path = _source_path(source_asset)
|
||||
ai_cfg = appconfig.ai_config(config)
|
||||
resolution = str(ai_cfg.get("resolution", "1k") or "1k")
|
||||
payload = {
|
||||
"prompt": str(job.prompt or ""),
|
||||
"model": runtime["alias"],
|
||||
"image_base64": ai._image_data_url(source_path),
|
||||
"resolution": ai._normalize_cmhub_resolution(resolution),
|
||||
"aspect_ratio": "1:1",
|
||||
}
|
||||
_notify(on_event, {"job_id": job.id, "step": "cover_submit", "result": "start"})
|
||||
data = ai._cmhub_call_with_retry(
|
||||
"POST",
|
||||
appconfig.cmhub_request_url(runtime["base_url"], "/api/v1/generate/image/tasks"),
|
||||
runtime["api_key"],
|
||||
payload=payload,
|
||||
connect_timeout=runtime["connect_timeout"],
|
||||
read_timeout=ai.CMHUB_IMAGE_SUBMIT_READ_TIMEOUT_SECONDS,
|
||||
attempts=max(1, int(ai_cfg.get("retry", 0) or 0) + 1),
|
||||
headers_extra={
|
||||
"Idempotency-Key": job.task_key,
|
||||
"X-Client-Version": str(APP_VERSION),
|
||||
},
|
||||
)
|
||||
image_task_id = str(data.get("task_id") or "").strip()
|
||||
if not image_task_id:
|
||||
raise ImageStudioGenerationError("cmhub 生图任务提交返回格式错误")
|
||||
submitted = image_studio.set_job_submitted(
|
||||
job.id,
|
||||
image_task_id,
|
||||
call_id=data.get("call_id"),
|
||||
points_cost=data.get("points_cost"),
|
||||
points_balance=data.get("points_balance"),
|
||||
path=db_path,
|
||||
)
|
||||
_notify(
|
||||
on_event,
|
||||
{
|
||||
"job_id": job.id,
|
||||
"step": "cover_submit",
|
||||
"result": "success",
|
||||
"task_id": submitted.task_id,
|
||||
"call_id": submitted.call_id,
|
||||
"points_cost": submitted.points_cost,
|
||||
"points_balance": submitted.points_balance,
|
||||
},
|
||||
)
|
||||
return _request_result(image_task_id, runtime, config, points_balance=data.get("points_balance"))
|
||||
|
||||
|
||||
def _request_result(task_id, runtime, config, points_balance=None):
|
||||
ai_cfg = appconfig.ai_config(config)
|
||||
resolution = str(ai_cfg.get("resolution", "1k") or "1k")
|
||||
return {
|
||||
"task_id": task_id,
|
||||
"connect_timeout": runtime["connect_timeout"],
|
||||
"read_timeout": ai.CMHUB_IMAGE_READ_TIMEOUT_SECONDS,
|
||||
"resolution": resolution,
|
||||
"quality": ai._jpg_quality(ai_cfg.get("jpg_quality", 90)),
|
||||
"use_system_proxy": runtime["use_system_proxy"],
|
||||
"download_with_curl": runtime["download_with_curl"],
|
||||
"points_balance": points_balance,
|
||||
}
|
||||
|
||||
|
||||
def _poll_job(job_id, task_id, runtime, request_result, db_path, should_stop, on_event):
|
||||
poll_url = appconfig.cmhub_request_url(
|
||||
runtime["base_url"],
|
||||
"/api/v1/generate/image/tasks/%s" % urllib.parse.quote(str(task_id), safe=""),
|
||||
)
|
||||
started = time.perf_counter()
|
||||
deadline = started + max(1, int(ai.CMHUB_IMAGE_READ_TIMEOUT_SECONDS))
|
||||
poll_index = 0
|
||||
while True:
|
||||
_raise_if_stopped(should_stop)
|
||||
if time.perf_counter() >= deadline:
|
||||
raise ImageStudioGenerationError("等待 cmhub 生图任务完成超时，下次可继续查询")
|
||||
_notify(on_event, {"job_id": job_id, "step": "cover_poll", "result": "start", "task_id": task_id})
|
||||
data = ai._cmhub_call_once(
|
||||
"GET",
|
||||
poll_url,
|
||||
runtime["api_key"],
|
||||
payload=None,
|
||||
connect_timeout=runtime["connect_timeout"],
|
||||
read_timeout=ai.CMHUB_IMAGE_POLL_READ_TIMEOUT_SECONDS,
|
||||
headers_extra={"X-Client-Version": str(APP_VERSION)},
|
||||
)
|
||||
status = str(data.get("status") or "").strip().lower()
|
||||
if status in {"queued", "running"}:
|
||||
_notify(on_event, {"job_id": job_id, "step": "cover_poll", "result": status, "task_id": task_id})
|
||||
ai._sleep_cmhub_poll(poll_index, should_stop)
|
||||
poll_index += 1
|
||||
continue
|
||||
if status == "succeeded":
|
||||
image_url = ai._extract_cmhub_image_url(data, runtime["base_url"])
|
||||
if not image_url:
|
||||
raise ImageStudioGenerationError("cmhub 生图任务成功但没有图片地址")
|
||||
merged = dict(request_result)
|
||||
merged["image_url"] = image_url
|
||||
merged["points_balance"] = data.get("points_balance", merged.get("points_balance"))
|
||||
_notify(on_event, {"job_id": job_id, "step": "cover_poll", "result": "success", "task_id": task_id})
|
||||
return merged
|
||||
if status in {"failed", "expired"}:
|
||||
error = data.get("error") if isinstance(data.get("error"), dict) else {}
|
||||
message = str(error.get("message") or error.get("code") or status)
|
||||
image_studio.update_job_status(job_id, status, error=message, path=db_path)
|
||||
raise ImageStudioGenerationError(message)
|
||||
raise ImageStudioGenerationError("cmhub 生图任务状态返回格式错误")
|
||||
|
||||
|
||||
def _download_and_save_job_image(request_result, out_path, config, on_event, job_id):
|
||||
_notify(on_event, {"job_id": job_id, "step": "cover_download", "result": "start"})
|
||||
image_bytes, _ = ai._download_cmhub_image_with_retry(
|
||||
request_result["image_url"],
|
||||
connect_timeout=request_result["connect_timeout"],
|
||||
read_timeout=request_result["read_timeout"],
|
||||
use_system_proxy=request_result.get("use_system_proxy", False),
|
||||
download_with_curl=request_result.get("download_with_curl", "auto"),
|
||||
)
|
||||
saved_path = ai._save_jpeg(
|
||||
image_bytes,
|
||||
out_path,
|
||||
request_result.get("resolution") or appconfig.ai_config(config).get("resolution", "1k"),
|
||||
request_result.get("quality") or appconfig.ai_config(config).get("jpg_quality", 90),
|
||||
)
|
||||
_notify(on_event, {"job_id": job_id, "step": "cover_download", "result": "success"})
|
||||
return saved_path
|
||||
|
||||
|
||||
def _raise_if_stopped(should_stop):
|
||||
try:
|
||||
stopped = bool(should_stop and should_stop())
|
||||
except Exception:
|
||||
stopped = False
|
||||
if stopped:
|
||||
raise ImageStudioGenerationError("用户已停止，已提交任务可稍后继续查询")
|
||||
+9
-2
@@ -3,7 +3,7 @@ id: T-590
|
||||
title: AI工场托管模式多图异步生成编排与重启续查
|
||||
phase: 7
|
||||
deps: [T-586, T-564]
|
||||
status: TODO
|
||||
status: DONE
|
||||
created: 2026-07-11
|
||||
---
|
||||
|
||||
@@ -37,4 +37,11 @@ AI工场一次可生成 N 张，且每张都要独立计费、失败、重试、
|
||||
|
||||
## 执行记录
|
||||
|
||||
（完成后记录抽取边界、任务生命周期和回归测试。）
|
||||
- 2026-07-11：完成 AI工场 cmhub 托管多图异步生成 service。
|
||||
- 新增 `app/image_studio_generation.py`，不改现有② AI生成路径；AI工场单独使用 `image_studio_jobs/assets` 持久化 job/key/task_id/计费/状态和输出 asset。
|
||||
- `generate_image_jobs()` 对同一源 asset + 完整提示词创建 N 条独立 job；每条 job 单独提交 `/api/v1/generate/image/tasks`，使用自身 `task_key` 作为 `Idempotency-Key`，submit 成功后立即写 `task_id/call_id/points_cost/points_balance`。
|
||||
- `resume_image_jobs()` 对已有 `task_id` 的 `submitted/running` job 只 GET poll，不再 POST，避免重启续查重复扣点；poll 成功后下载保存 JPEG 到项目 `generated/` 并创建 `generated_main/generated_detail` asset，关联父源图、比例、提示词和远程 `image_url`。
|
||||
- 单张失败只更新该 job 为 `failed`，不回滚其他成功图；下载失败发生在 poll 成功之后，不重新 submit、不创建输出 asset；停止语义保留已提交 job 的 task_id，未开始/可取消 job 标记 `cancelled`。
|
||||
- 新增 `tests/test_image_studio_generation.py` 覆盖 N=3 独立 job/key/task、submit 后 poll 前已落库、续查不 POST、下载失败不重提交/不创建 asset、服务端单张 failed 不影响其他成功。
|
||||
- 验证：在只套用 T-590 diff 的 clean worktree 中运行 `py -3.10 -m unittest tests.test_image_studio_generation tests.test_ai tests.test_db`（66 tests）、`python -m ruff check app tests main.py`、`py -3.10 -m compileall app main.py`、`py -3.10 -m unittest discover -s tests`（360 tests）和 `git diff --check`，全部通过。
|
||||
- 本任务未接 GUI，也未打真实 cmhub 收费接口；T-591 接 UI 时应在 worker/QThread 中调用该 service，避免阻塞主线程。
|
||||
|
||||
@@ -0,0 +1,282 @@
|
||||
import io
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
from unittest import mock
|
||||
|
||||
sys.path.insert(0, os.path.dirname(__file__))
|
||||
|
||||
from _helpers import TempDirMixin
|
||||
|
||||
from app import db, image_studio, image_studio_generation
|
||||
|
||||
|
||||
class ImageStudioGenerationTests(TempDirMixin, unittest.TestCase):
|
||||
def _png_bytes(self):
|
||||
from PIL import Image
|
||||
|
||||
output = io.BytesIO()
|
||||
Image.new("RGB", (32, 32), (120, 80, 160)).save(output, format="PNG")
|
||||
return output.getvalue()
|
||||
|
||||
def _config(self, temp_dir):
|
||||
return {
|
||||
"data_dir": temp_dir,
|
||||
"db_path": os.path.join(temp_dir, "cmshopee.db"),
|
||||
"image_dir": os.path.join(temp_dir, "images"),
|
||||
"ai": {
|
||||
"backend": "cmhub",
|
||||
"image_concurrency": 1,
|
||||
"retry": 0,
|
||||
"resolution": "1k",
|
||||
"jpg_quality": 90,
|
||||
"cmhub": {
|
||||
"base_url": "https://cmhub.example.com",
|
||||
"image_alias": "image-hd",
|
||||
"connect_timeout": 3,
|
||||
"download_with_curl": "false",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
def _project_source(self, temp_dir):
|
||||
cfg = self._config(temp_dir)
|
||||
db.init_db(cfg["db_path"])
|
||||
project = image_studio.create_or_get_project(
|
||||
account_alias="alias",
|
||||
account_slug="alias_slug",
|
||||
item_id="51100639510",
|
||||
path=cfg["db_path"],
|
||||
)
|
||||
source_path = os.path.join(temp_dir, "source.png")
|
||||
with open(source_path, "wb") as fh:
|
||||
fh.write(self._png_bytes())
|
||||
source = image_studio.add_asset(
|
||||
project.id,
|
||||
image_studio.ASSET_KIND_ORIGINAL,
|
||||
remote_url="https://cdn.example.com/source.png",
|
||||
local_path=source_path,
|
||||
source_order=1,
|
||||
path=cfg["db_path"],
|
||||
)
|
||||
return cfg, project, source
|
||||
|
||||
def _runtime(self):
|
||||
return {
|
||||
"base_url": "https://cmhub.example.com",
|
||||
"api_key": "sk-test",
|
||||
"alias": "image-hd",
|
||||
"connect_timeout": 3,
|
||||
"use_system_proxy": False,
|
||||
"download_with_curl": "false",
|
||||
}
|
||||
|
||||
def test_generate_image_jobs_creates_independent_jobs_and_assets(self):
|
||||
with self.make_temp_dir() as temp_dir:
|
||||
cfg, project, source = self._project_source(temp_dir)
|
||||
submitted = []
|
||||
poll_seen_persisted = []
|
||||
|
||||
def fake_submit(method, url, api_key, **kwargs):
|
||||
self.assertEqual("POST", method)
|
||||
task_id = "cmhub-task-%d" % (len(submitted) + 1)
|
||||
submitted.append((task_id, kwargs["headers_extra"]["Idempotency-Key"]))
|
||||
return {
|
||||
"task_id": task_id,
|
||||
"status": "queued",
|
||||
"call_id": "call-%d" % len(submitted),
|
||||
"points_cost": 2,
|
||||
"points_balance": 100 - len(submitted) * 2,
|
||||
}
|
||||
|
||||
def fake_poll(method, url, api_key, **kwargs):
|
||||
self.assertEqual("GET", method)
|
||||
task_id = url.rsplit("/", 1)[-1]
|
||||
jobs = image_studio.list_resumable_jobs(path=cfg["db_path"], project_id=project.id)
|
||||
poll_seen_persisted.append(any(job.task_id == task_id for job in jobs))
|
||||
return {
|
||||
"task_id": task_id,
|
||||
"status": "succeeded",
|
||||
"result": {"image_url": "https://cdn.example.com/%s.png" % task_id},
|
||||
}
|
||||
|
||||
with mock.patch("app.image_studio_generation._runtime", return_value=self._runtime()), \
|
||||
mock.patch(
|
||||
"app.image_studio_generation.ai._cmhub_call_with_retry",
|
||||
side_effect=fake_submit,
|
||||
), \
|
||||
mock.patch(
|
||||
"app.image_studio_generation.ai._cmhub_call_once",
|
||||
side_effect=fake_poll,
|
||||
), \
|
||||
mock.patch(
|
||||
"app.image_studio_generation.ai._download_cmhub_image_with_retry",
|
||||
return_value=(self._png_bytes(), 0.1),
|
||||
), \
|
||||
mock.patch("app.image_studio_generation.ai._sleep_cmhub_poll"):
|
||||
summary = image_studio_generation.generate_image_jobs(
|
||||
project.id,
|
||||
source.id,
|
||||
"完整提示词",
|
||||
3,
|
||||
config=cfg,
|
||||
path=cfg["db_path"],
|
||||
)
|
||||
|
||||
self.assertEqual(3, summary["success"])
|
||||
self.assertEqual(0, summary["failed"])
|
||||
self.assertEqual(3, len(submitted))
|
||||
self.assertEqual(3, len({key for _, key in submitted}))
|
||||
self.assertEqual([True, True, True], poll_seen_persisted)
|
||||
jobs = image_studio.list_resumable_jobs(path=cfg["db_path"], project_id=project.id)
|
||||
self.assertEqual([], jobs)
|
||||
all_jobs = [
|
||||
image_studio.get_job(result["job"].id, path=cfg["db_path"])
|
||||
for result in summary["jobs"]
|
||||
]
|
||||
self.assertEqual(["succeeded", "succeeded", "succeeded"], [job.status for job in all_jobs])
|
||||
self.assertTrue(all(job.task_id for job in all_jobs))
|
||||
self.assertTrue(all(job.task_key for job in all_jobs))
|
||||
self.assertEqual([2, 2, 2], [job.points_cost for job in all_jobs])
|
||||
assets = image_studio.list_assets(project.id, kind="generated_main", path=cfg["db_path"])
|
||||
self.assertEqual(3, len(assets))
|
||||
self.assertTrue(all(os.path.isfile(asset.local_path) for asset in assets))
|
||||
self.assertTrue(all(asset.parent_asset_id == source.id for asset in assets))
|
||||
self.assertTrue(all(asset.prompt == "完整提示词" for asset in assets))
|
||||
|
||||
self.assert_removed(temp_dir)
|
||||
|
||||
def test_resume_existing_job_polls_without_new_submit(self):
|
||||
with self.make_temp_dir() as temp_dir:
|
||||
cfg, project, source = self._project_source(temp_dir)
|
||||
job = image_studio.create_job(
|
||||
project.id,
|
||||
source_asset_id=source.id,
|
||||
job_type="main",
|
||||
prompt="续查提示词",
|
||||
path=cfg["db_path"],
|
||||
)
|
||||
image_studio.set_job_submitted(job.id, "cmhub-task-resume", path=cfg["db_path"])
|
||||
|
||||
def fake_poll(method, url, api_key, **kwargs):
|
||||
return {
|
||||
"task_id": "cmhub-task-resume",
|
||||
"status": "succeeded",
|
||||
"result": {"image_url": "https://cdn.example.com/resume.png"},
|
||||
}
|
||||
|
||||
with mock.patch("app.image_studio_generation._runtime", return_value=self._runtime()), \
|
||||
mock.patch("app.image_studio_generation.ai._cmhub_call_with_retry") as submit, \
|
||||
mock.patch("app.image_studio_generation.ai._cmhub_call_once", side_effect=fake_poll), \
|
||||
mock.patch(
|
||||
"app.image_studio_generation.ai._download_cmhub_image_with_retry",
|
||||
return_value=(self._png_bytes(), 0.1),
|
||||
):
|
||||
summary = image_studio_generation.resume_image_jobs(
|
||||
project_id=project.id,
|
||||
config=cfg,
|
||||
path=cfg["db_path"],
|
||||
)
|
||||
|
||||
self.assertEqual(1, summary["success"])
|
||||
submit.assert_not_called()
|
||||
updated = image_studio.get_job(job.id, path=cfg["db_path"])
|
||||
self.assertEqual("succeeded", updated.status)
|
||||
self.assertEqual("cmhub-task-resume", updated.task_id)
|
||||
|
||||
self.assert_removed(temp_dir)
|
||||
|
||||
def test_download_failure_does_not_submit_again_or_create_asset(self):
|
||||
with self.make_temp_dir() as temp_dir:
|
||||
cfg, project, source = self._project_source(temp_dir)
|
||||
submit_count = 0
|
||||
|
||||
def fake_submit(method, url, api_key, **kwargs):
|
||||
nonlocal submit_count
|
||||
submit_count += 1
|
||||
return {"task_id": "cmhub-task-1", "status": "queued"}
|
||||
|
||||
def fake_poll(method, url, api_key, **kwargs):
|
||||
return {
|
||||
"task_id": "cmhub-task-1",
|
||||
"status": "succeeded",
|
||||
"result": {"image_url": "https://cdn.example.com/fail.png"},
|
||||
}
|
||||
|
||||
with mock.patch("app.image_studio_generation._runtime", return_value=self._runtime()), \
|
||||
mock.patch(
|
||||
"app.image_studio_generation.ai._cmhub_call_with_retry",
|
||||
side_effect=fake_submit,
|
||||
), \
|
||||
mock.patch("app.image_studio_generation.ai._cmhub_call_once", side_effect=fake_poll), \
|
||||
mock.patch(
|
||||
"app.image_studio_generation.ai._download_cmhub_image_with_retry",
|
||||
side_effect=RuntimeError("download failed"),
|
||||
):
|
||||
summary = image_studio_generation.generate_image_jobs(
|
||||
project.id,
|
||||
source.id,
|
||||
"完整提示词",
|
||||
1,
|
||||
config=cfg,
|
||||
path=cfg["db_path"],
|
||||
)
|
||||
|
||||
self.assertEqual(0, summary["success"])
|
||||
self.assertEqual(1, summary["failed"])
|
||||
self.assertEqual(1, submit_count)
|
||||
jobs = [result["job"] for result in summary["jobs"]]
|
||||
self.assertEqual("failed", image_studio.get_job(jobs[0].id, path=cfg["db_path"]).status)
|
||||
self.assertEqual([], image_studio.list_assets(project.id, kind="generated_main", path=cfg["db_path"]))
|
||||
|
||||
self.assert_removed(temp_dir)
|
||||
|
||||
def test_failed_cmhub_task_marks_only_that_job_failed(self):
|
||||
with self.make_temp_dir() as temp_dir:
|
||||
cfg, project, source = self._project_source(temp_dir)
|
||||
submitted = []
|
||||
|
||||
def fake_submit(method, url, api_key, **kwargs):
|
||||
task_id = "cmhub-task-%d" % (len(submitted) + 1)
|
||||
submitted.append(task_id)
|
||||
return {"task_id": task_id, "status": "queued"}
|
||||
|
||||
def fake_poll(method, url, api_key, **kwargs):
|
||||
task_id = url.rsplit("/", 1)[-1]
|
||||
if task_id.endswith("-2"):
|
||||
return {"task_id": task_id, "status": "failed", "error": {"message": "上游失败"}}
|
||||
return {
|
||||
"task_id": task_id,
|
||||
"status": "succeeded",
|
||||
"result": {"image_url": "https://cdn.example.com/%s.png" % task_id},
|
||||
}
|
||||
|
||||
with mock.patch("app.image_studio_generation._runtime", return_value=self._runtime()), \
|
||||
mock.patch(
|
||||
"app.image_studio_generation.ai._cmhub_call_with_retry",
|
||||
side_effect=fake_submit,
|
||||
), \
|
||||
mock.patch("app.image_studio_generation.ai._cmhub_call_once", side_effect=fake_poll), \
|
||||
mock.patch(
|
||||
"app.image_studio_generation.ai._download_cmhub_image_with_retry",
|
||||
return_value=(self._png_bytes(), 0.1),
|
||||
):
|
||||
summary = image_studio_generation.generate_image_jobs(
|
||||
project.id,
|
||||
source.id,
|
||||
"完整提示词",
|
||||
2,
|
||||
config=cfg,
|
||||
path=cfg["db_path"],
|
||||
)
|
||||
|
||||
self.assertEqual(1, summary["success"])
|
||||
self.assertEqual(1, summary["failed"])
|
||||
assets = image_studio.list_assets(project.id, kind="generated_main", path=cfg["db_path"])
|
||||
self.assertEqual(1, len(assets))
|
||||
|
||||
self.assert_removed(temp_dir)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user