chengma/cmshoppe
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(product-suite): add direct job state machine

Browse Source
This commit is contained in:
chengma
2026-07-20 18:12:44 +08:00
parent 2a2cbad7bd
commit aa000ff5e6
11 changed files with 820 additions and 101 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/image_studio_generation.py
+351 -54
View File
@@ -1,4 +1,4 @@
"""cmhub hosted generation orchestration for AI image studio jobs."""
"""Multi-source generation orchestration for AI image studio jobs."""
from __future__ import annotations
@@ -6,21 +6,97 @@ import os
import threading
import time
import urllib.parse
import uuid
from concurrent.futures import CancelledError, FIRST_COMPLETED, ThreadPoolExecutor, wait
from . import ai, appconfig, image_studio
from .version import APP_VERSION
MAX_CMHUB_IMAGE_STUDIO_WORKERS = 5
_GLOBAL_IMAGE_STUDIO_SLOTS = threading.BoundedSemaphore(MAX_CMHUB_IMAGE_STUDIO_WORKERS)
MAX_IMAGE_STUDIO_WORKERS = 5
_GLOBAL_IMAGE_STUDIO_SLOTS = threading.BoundedSemaphore(MAX_IMAGE_STUDIO_WORKERS)
CMHUB_IMAGE_STUDIO_MAX_INPUT_IMAGES = 8
CMHUB_IMAGE_STUDIO_MAX_SINGLE_INPUT_BYTES = 10 * 1024 * 1024
CMHUB_IMAGE_STUDIO_MAX_TOTAL_INPUT_BYTES = 32 * 1024 * 1024
class ImageStudioGenerationError(RuntimeError):
"""Raised when AI studio hosted generation cannot complete."""
"""Raised when AI studio image generation cannot complete."""
class ImageStudioRuntimeLease:
"""A process-lifetime lock used only to make startup recovery safe."""
def __init__(self, handle, path):
self._handle = handle
self.path = path
self.session_id = "startup-" + uuid.uuid4().hex
def release(self):
handle, self._handle = self._handle, None
if handle is None:
return
try:
handle.seek(0)
if os.name == "nt":
import msvcrt
msvcrt.locking(handle.fileno(), msvcrt.LK_UNLCK, 1)
else: # pragma: no cover - Windows release is the supported path.
import fcntl
fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
finally:
handle.close()
def acquire_startup_recovery_lease(data_dir):
"""Acquire a non-blocking process lease; return None when another app owns it."""
directory = os.path.abspath(str(data_dir or ""))
if not directory:
return None
try:
os.makedirs(directory, exist_ok=True)
path = os.path.join(directory, ".image_studio_generation.lock")
handle = open(path, "a+b")
if not handle.read(1):
handle.write(b"0")
handle.flush()
handle.seek(0)
if os.name == "nt":
import msvcrt
msvcrt.locking(handle.fileno(), msvcrt.LK_NBLCK, 1)
else: # pragma: no cover - retained for developer tests outside Windows.
import fcntl
fcntl.flock(handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
return ImageStudioRuntimeLease(handle, path)
except (OSError, ImportError):
try:
handle.close()
except Exception:
pass
return None
def recover_stale_direct_jobs_at_startup(*, lease, path=None):
"""Recover interrupted direct jobs only while the process lease is held."""
if lease is None:
return {"recovered": [], "skipped": True}
recovered = image_studio.fail_stale_direct_jobs(path=path)
return {"recovered": recovered, "skipped": False}
RATIO_OUTPUT_SIZES = {
"1:1": ("1024x1024", False),
"3:4": ("1024x1536", True),
"9:16": ("1024x1536", True),
"4:3": ("1536x1024", True),
"16:9": ("1536x1024", True),
}
def _notify(callback, event):
@@ -36,16 +112,73 @@ def _runtime(config, cmhub_config_path):
return ai._cmhub_runtime(config, "image", cmhub_config_path)
def _direct_runtime(config):
ai_cfg = appconfig.ai_config(config)
models_path = appconfig.ai_models_config_path(config)
model = ai._role_model(
"image",
ai_cfg.get("default_image_model"),
models_path,
models=ai._runtime_direct_models(config),
)
error = appconfig.image_model_config_error(model)
if error:
raise ImageStudioGenerationError(error)
return {
"model": model,
"connect_timeout": ai._connect_timeout(model),
"read_timeout": ai._read_timeout(model, config),
"quality": ai._jpg_quality(ai_cfg.get("jpg_quality", 90)),
}
def _is_default_gateway_job(job):
return (
str(getattr(job, "generation_source", "") or "").strip().lower() == "cmhub"
and str(getattr(job, "provider", "") or "").strip().lower() == "cmhub"
str(getattr(job, "generation_source", "") or "").strip().lower()
== image_studio.GENERATION_SOURCE_CMHUB
and str(getattr(job, "provider", "") or "").strip().lower()
== image_studio.PROVIDER_CMHUB
)
def _is_direct_gateway_job(job):
return (
str(getattr(job, "generation_source", "") or "").strip().lower()
== image_studio.GENERATION_SOURCE_DIRECT
and str(getattr(job, "provider", "") or "").strip().lower()
== image_studio.PROVIDER_OPENAI_IMAGES_EDITS
)
def generation_source_for_config(config):
"""Return the persisted source/provider pair for one frozen run config."""
if appconfig.ai_backend(config) == "direct":
return {
"generation_source": image_studio.GENERATION_SOURCE_DIRECT,
"provider": image_studio.PROVIDER_OPENAI_IMAGES_EDITS,
}
return {
"generation_source": image_studio.GENERATION_SOURCE_CMHUB,
"provider": image_studio.PROVIDER_CMHUB,
}
def requested_output_spec(aspect_ratio):
"""Map a product-suite ratio to an explicit request size and approximation flag."""
ratio = str(aspect_ratio or "1:1")
size, approximate = RATIO_OUTPUT_SIZES.get(ratio, RATIO_OUTPUT_SIZES["1:1"])
return {
"aspect_ratio": ratio,
"requested_output_size": size,
"approximate_ratio": bool(approximate),
}
def _ensure_new_submission_allowed(config):
if appconfig.ai_backend(config) != "cmhub":
raise ImageStudioGenerationError("商品套图仅支持默认网关，请到⑤设置切换后再生成")
if appconfig.ai_backend(config) not in {"cmhub", "direct"}:
raise ImageStudioGenerationError("商品套图生成网关配置无效，请到⑤设置检查")
def _generated_kind(job_type):
@@ -104,6 +237,7 @@ def create_generation_jobs(
total = max(0, int(count or 0))
if total <= 0:
raise ImageStudioGenerationError("生成数量必须大于0")
source = generation_source_for_config(config)
jobs = []
for _ in range(total):
jobs.append(
@@ -112,8 +246,8 @@ def create_generation_jobs(
source_asset_id=source_asset_id,
job_type=job_type,
prompt=prompt,
generation_source="cmhub",
provider="cmhub",
generation_source=source["generation_source"],
provider=source["provider"],
path=path,
)
)
@@ -133,6 +267,7 @@ def generate_image_jobs(
path=None,
should_stop=None,
on_event=None,
run_session_id=None,
):
jobs = create_generation_jobs(
project_id,
@@ -151,6 +286,7 @@ def generate_image_jobs(
path=path,
should_stop=should_stop,
on_event=on_event,
run_session_id=run_session_id,
)
@@ -163,6 +299,7 @@ def resume_image_jobs(
path=None,
should_stop=None,
on_event=None,
run_session_id=None,
):
jobs = image_studio.list_resumable_jobs(
path=path,
@@ -177,6 +314,7 @@ def resume_image_jobs(
path=path,
should_stop=should_stop,
on_event=on_event,
run_session_id=run_session_id,
)
@@ -189,16 +327,18 @@ def run_jobs(
path=None,
should_stop=None,
on_event=None,
run_session_id=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": []}
rejected = [
job
for job in job_list
if getattr(job, "task_id", None) and not _is_default_gateway_job(job)
]
rejected = []
for job in job_list:
if _is_direct_gateway_job(job) and getattr(job, "task_id", None):
rejected.append(job)
elif not _is_default_gateway_job(job) and not _is_direct_gateway_job(job):
rejected.append(job)
job_list = [job for job in job_list if job not in rejected]
summary = {
"total": len(job_list) + len(rejected),
@@ -209,7 +349,7 @@ def run_jobs(
{
"job": job,
"status": "failed",
"error": "该已提交任务不属于默认网关，不能继续查询",
"error": "AI工场任务来源无效，不能执行",
}
for job in rejected
],
@@ -217,10 +357,23 @@ def run_jobs(
if not job_list:
return summary
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))
runtimes = {}
if any(_is_default_gateway_job(job) for job in job_list):
runtimes[image_studio.GENERATION_SOURCE_CMHUB] = _runtime(
cfg,
cmhub_config_path,
)
if any(_is_direct_gateway_job(job) for job in job_list):
runtimes[image_studio.GENERATION_SOURCE_DIRECT] = _direct_runtime(cfg)
run_session_id = str(run_session_id or "").strip() or "direct-" + uuid.uuid4().hex
output_spec = requested_output_spec(aspect_ratio)
max_workers = min(
MAX_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()
@@ -239,13 +392,14 @@ def run_jobs(
executor.submit(
_run_one_job_with_global_slot,
job.id,
runtime,
runtimes,
cfg,
image_root,
aspect_ratio,
output_spec,
path,
should_stop,
on_event,
run_session_id,
): job
for job in job_list
}
@@ -273,6 +427,7 @@ def run_jobs(
)
futures.pop(future, None)
record({"job": job, "status": "cancelled", "error": "用户停止"})
summary["output"] = dict(output_spec)
return summary
@@ -281,7 +436,17 @@ def _run_one_job_with_global_slot(*args):
return _run_one_job(*args)
def _run_one_job(job_id, runtime, config, image_root, aspect_ratio, db_path, should_stop, on_event):
def _run_one_job(
job_id,
runtimes,
config,
image_root,
output_spec,
db_path,
should_stop,
on_event,
run_session_id,
):
job = image_studio.get_job(job_id, path=db_path)
if job is None:
raise ImageStudioGenerationError("AI工场生图任务不存在")
@@ -305,48 +470,94 @@ def _run_one_job(job_id, runtime, config, image_root, aspect_ratio, db_path, sho
path=db_path,
)
return {"job": updated, "status": "failed", "error": "项目或源图不存在"}
is_direct = _is_direct_gateway_job(job)
if not (_is_default_gateway_job(job) or is_direct):
raise ImageStudioGenerationError("AI工场任务来源无效，不能执行")
runtime = runtimes.get(job.generation_source)
if runtime is None:
raise ImageStudioGenerationError("当前运行缺少该任务来源的配置快照")
try:
image_studio.update_job_status(job.id, "running", path=db_path)
image_studio.update_job_status(
job.id,
"running",
run_session_id=run_session_id if is_direct else None,
path=db_path,
)
reference_assets = ()
if not job.task_id:
reference_assets = _reference_assets_for_job(job, db_path)
request_result = _submit_or_resume_job(
job,
source_asset,
runtime,
config,
aspect_ratio,
db_path,
should_stop,
on_event,
reference_assets=reference_assets,
)
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)
_raise_if_stopped(should_stop)
out_path = _output_path(project, job, image_root)
saved_path = _download_and_save_job_image(
request_result,
out_path,
config,
on_event,
job.id,
should_stop=should_stop,
)
try:
if is_direct:
request_result = _submit_direct_job(
job,
source_asset,
runtime,
config,
output_spec,
should_stop,
on_event,
reference_assets=reference_assets,
)
# A submitted direct request cannot be cancelled. If bytes arrive after
# stop was requested, save them before ending the remaining queue.
saved_path = _save_direct_job_image(
request_result,
out_path,
on_event,
job.id,
)
remote_url = None
else:
request_result = _submit_or_resume_job(
job,
source_asset,
runtime,
config,
output_spec["aspect_ratio"],
db_path,
should_stop,
on_event,
reference_assets=reference_assets,
)
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,
)
_raise_if_stopped(should_stop)
except ImageStudioGenerationError:
request_result["resolution"] = output_spec["requested_output_size"]
saved_path = _download_and_save_job_image(
request_result,
out_path,
config,
on_event,
job.id,
should_stop=should_stop,
)
try:
if os.path.isfile(saved_path):
os.remove(saved_path)
finally:
raise
_raise_if_stopped(should_stop)
except ImageStudioGenerationError:
try:
if os.path.isfile(saved_path):
os.remove(saved_path)
finally:
raise
remote_url = request_result["image_url"]
rendered_width, rendered_height = _rendered_image_size(saved_path)
asset = image_studio.add_asset(
project.id,
_generated_kind(job.job_type),
remote_url=request_result["image_url"],
remote_url=remote_url,
local_path=saved_path,
aspect_ratio=aspect_ratio,
aspect_ratio=output_spec["aspect_ratio"],
requested_output_size=output_spec["requested_output_size"],
rendered_width=rendered_width,
rendered_height=rendered_height,
parent_asset_id=source_asset.id,
prompt=job.prompt,
path=db_path,
@@ -361,9 +572,17 @@ def _run_one_job(job_id, runtime, config, image_root, aspect_ratio, db_path, sho
_notify(on_event, {"job_id": job.id, "step": "job_done", "result": "success"})
return {"job": updated, "asset": asset, "status": "succeeded"}
except Exception as exc:
cancelled = isinstance(exc, CancelledError) or "停止" in str(exc)
cancelled = (
not is_direct
and (isinstance(exc, CancelledError) or "停止" in str(exc))
)
status = "cancelled" if cancelled else "failed"
error = "用户已停止，已提交任务可稍后继续查询" if cancelled else str(exc)
if cancelled:
error = "用户已停止，已提交任务可稍后继续查询"
elif is_direct:
error = _direct_failure_message(exc)
else:
error = str(exc)
current_job = image_studio.get_job(job.id, path=db_path)
updated = image_studio.update_job_status(
job.id,
@@ -398,6 +617,84 @@ def _reference_assets_for_job(job, db_path):
return assets
def _direct_image_paths(source_asset, reference_assets=()):
"""Return ordered local paths; the first source image is the subject image."""
return [_source_path(source_asset)] + [
_source_path(asset) for asset in (reference_assets or ())
]
def _submit_direct_job(
job,
source_asset,
runtime,
config,
output_spec,
should_stop,
on_event,
reference_assets=(),
):
if job.task_id:
raise ImageStudioGenerationError("自定义网关任务不能继续查询，请手动重新生成")
_raise_if_stopped(should_stop)
image_paths = _direct_image_paths(source_asset, reference_assets)
body, content_type = ai._image_edit_body(
runtime["model"],
job.prompt,
image_paths,
output_spec["requested_output_size"],
)
_notify(on_event, {"job_id": job.id, "step": "cover_submit", "result": "start"})
# Direct images/edits has no idempotency or query endpoint. Retrying an
# uncertain request could create and charge a second image, so one call only.
data = ai._call_with_retry(
runtime["model"],
body,
config,
1,
request_kind="multipart",
content_type=content_type,
)
image_bytes = ai._extract_image_bytes(data, runtime["model"], config)
_notify(on_event, {"job_id": job.id, "step": "cover_request", "result": "success"})
return {
"image_bytes": image_bytes,
"resolution": output_spec["requested_output_size"],
"quality": runtime["quality"],
}
def _save_direct_job_image(request_result, out_path, on_event, job_id):
_notify(on_event, {"job_id": job_id, "step": "cover_download", "result": "start"})
saved_path = ai._save_jpeg(
request_result["image_bytes"],
out_path,
request_result["resolution"],
request_result["quality"],
)
_notify(on_event, {"job_id": job_id, "step": "cover_download", "result": "success"})
return saved_path
def _rendered_image_size(path):
try:
from PIL import Image
with Image.open(path) as image:
width, height = image.size
return int(width), int(height)
except Exception as exc:
raise ImageStudioGenerationError("生成图片保存后无法读取尺寸") from exc
def _direct_failure_message(exc):
text = str(exc or "").lower()
if "timeout" in text or "超时" in text:
return "自定义网关生成超时，无法确认服务商是否已处理，请手动重新生成"
return "自定义网关生成失败，结果无法确认，请检查图片模型配置后手动重新生成"
def _submit_or_resume_job(
job,
source_asset,