fix(product-suite): finalize and cancel generation
Tests / Python 3.11 / Windows (push) Has been cancelled

This commit is contained in:
chengma
2026-07-16 17:12:41 +08:00
parent 2675f85598
commit b68e22c094
11 changed files with 1094 additions and 85 deletions
+102 -22
View File
@@ -1447,6 +1447,7 @@ def _download_cmhub_image_with_retry(
use_system_proxy=False,
download_with_curl="false",
on_step=None,
should_stop=None,
attempts=CMHUB_IMAGE_DOWNLOAD_ATTEMPTS,
slow_threshold=CMHUB_IMAGE_SLOW_DOWNLOAD_SECONDS,
):
@@ -1454,6 +1455,7 @@ def _download_cmhub_image_with_retry(
total_started = time.perf_counter()
last_exc = None
for index in range(total_attempts):
_raise_if_download_cancelled(should_stop)
try:
image_bytes = _download_cmhub_image(
url,
@@ -1461,7 +1463,9 @@ def _download_cmhub_image_with_retry(
read_timeout=read_timeout,
use_system_proxy=use_system_proxy,
download_with_curl=download_with_curl,
should_stop=should_stop,
)
_raise_if_download_cancelled(should_stop)
elapsed = time.perf_counter() - total_started
if elapsed >= float(slow_threshold or 0):
_notify_step_event(
@@ -1473,6 +1477,8 @@ def _download_cmhub_image_with_retry(
level="warning",
)
return image_bytes, elapsed
except CancelledError:
raise
except Exception as exc:
last_exc = exc
if index + 1 >= total_attempts or not _cmhub_download_retryable(exc):
@@ -1486,7 +1492,7 @@ def _download_cmhub_image_with_retry(
attempt=index + 1,
attempts=total_attempts,
)
time.sleep(min(2.0, 0.5 * (index + 1)))
_sleep_download_retry(min(2.0, 0.5 * (index + 1)), should_stop)
if total_attempts > 1 and _cmhub_download_retryable(last_exc):
raise AIError(
"下载 cmhub 图片失败(已尝试 %s 次): %s"
@@ -1739,7 +1745,9 @@ def _download_cmhub_image(
max_bytes=CMHUB_IMAGE_MAX_BYTES,
use_system_proxy=False,
download_with_curl="false",
should_stop=None,
):
_raise_if_download_cancelled(should_stop)
_assert_public_http_url(url)
if _should_use_curl_for_cmhub_download(download_with_curl):
try:
@@ -1749,7 +1757,10 @@ def _download_cmhub_image(
read_timeout=read_timeout,
max_bytes=max_bytes,
use_system_proxy=use_system_proxy,
should_stop=should_stop,
)
except CancelledError:
raise
except AIError:
pass
return _download_cmhub_image_with_requests(
@@ -1757,32 +1768,53 @@ def _download_cmhub_image(
connect_timeout=connect_timeout,
read_timeout=read_timeout,
max_bytes=max_bytes,
should_stop=should_stop,
)
def _download_cmhub_image_with_requests(url, connect_timeout, read_timeout, max_bytes):
def _download_cmhub_image_with_requests(
url,
connect_timeout,
read_timeout,
max_bytes,
should_stop=None,
):
response = None
try:
_raise_if_download_cancelled(should_stop)
response = _cmhub_session().get(
url,
stream=True,
timeout=(max(1, int(connect_timeout)), max(1, int(read_timeout))),
)
_raise_if_download_cancelled(should_stop)
status = getattr(response, "status_code", 200)
if status >= 400:
raise AIError("下载 cmhub 图片失败: HTTP %s" % status)
chunks = []
total = 0
iterator = (
response.iter_content(chunk_size=65536)
if hasattr(response, "iter_content")
else [response.content]
)
for chunk in iterator:
_raise_if_download_cancelled(should_stop)
if not chunk:
continue
total += len(chunk)
if total > max_bytes:
raise AIError("下载 cmhub 图片失败: 图片超过大小上限")
chunks.append(chunk)
_raise_if_download_cancelled(should_stop)
return b"".join(chunks)
except CancelledError:
raise
except requests.exceptions.RequestException as exc:
raise AIError("下载 cmhub 图片失败: %s" % exc) from exc
status = getattr(response, "status_code", 200)
if status >= 400:
raise AIError("下载 cmhub 图片失败: HTTP %s" % status)
chunks = []
total = 0
iterator = response.iter_content(chunk_size=65536) if hasattr(response, "iter_content") else [response.content]
for chunk in iterator:
if not chunk:
continue
total += len(chunk)
if total > max_bytes:
raise AIError("下载 cmhub 图片失败: 图片超过大小上限")
chunks.append(chunk)
return b"".join(chunks)
finally:
if response is not None and hasattr(response, "close"):
response.close()
def _should_use_curl_for_cmhub_download(mode):
@@ -1846,6 +1878,7 @@ def _download_cmhub_image_with_curl(
read_timeout,
max_bytes,
use_system_proxy=False,
should_stop=None,
):
curl_path = _find_system_curl()
if not curl_path:
@@ -1881,20 +1914,36 @@ def _download_cmhub_image_with_curl(
]
if not bool(use_system_proxy):
args.extend(["--noproxy", "*"])
process = None
try:
completed = subprocess.run(
process = subprocess.Popen(
args,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
timeout=max(2, int(connect_timeout) + int(read_timeout) + 10),
check=False,
shell=False,
**_subprocess_hidden_window_kwargs(),
)
except (OSError, subprocess.TimeoutExpired) as exc:
deadline = time.monotonic() + max(
2,
int(connect_timeout) + int(read_timeout) + 10,
)
while process.poll() is None:
try:
_raise_if_download_cancelled(should_stop)
except CancelledError:
_stop_download_process(process)
raise
if time.monotonic() >= deadline:
_stop_download_process(process)
raise AIError("下载 cmhub 图片失败: curl 执行超时")
time.sleep(0.1)
process.communicate()
except CancelledError:
raise
except OSError as exc:
raise AIError("下载 cmhub 图片失败: curl 执行失败") from exc
if completed.returncode != 0:
raise AIError("下载 cmhub 图片失败: curl 退出码 %s" % completed.returncode)
if process.returncode != 0:
raise AIError("下载 cmhub 图片失败: curl 退出码 %s" % process.returncode)
size = os.path.getsize(temp_output_path)
if size > max_bytes:
raise AIError("下载 cmhub 图片失败: 图片超过大小上限")
@@ -1909,6 +1958,37 @@ def _download_cmhub_image_with_curl(
pass
def _stop_download_process(process):
if process is None or process.poll() is not None:
return
try:
process.terminate()
process.wait(timeout=2)
except (OSError, subprocess.TimeoutExpired):
try:
process.kill()
process.wait(timeout=2)
except (OSError, subprocess.TimeoutExpired):
pass
def _raise_if_download_cancelled(should_stop):
try:
stopped = bool(should_stop and should_stop())
except Exception:
stopped = False
if stopped:
raise CancelledError()
def _sleep_download_retry(delay_seconds, should_stop):
deadline = time.monotonic() + max(0.0, float(delay_seconds or 0.0))
while time.monotonic() < deadline:
_raise_if_download_cancelled(should_stop)
time.sleep(min(0.1, max(0.0, deadline - time.monotonic())))
_raise_if_download_cancelled(should_stop)
def _curl_config_quote(value):
text = str(value or "")
return '"' + text.replace("\\", "\\\\").replace('"', '\\"') + '"'
+393 -27
View File
@@ -5,9 +5,20 @@ from __future__ import annotations
import os
import re
import time
import uuid
from dataclasses import dataclass, field
from PySide6.QtCore import QByteArray, QBuffer, QIODevice, QRect, QSize, Qt, QTimer, Signal
from PySide6.QtCore import (
QByteArray,
QBuffer,
QIODevice,
QRect,
QSize,
Qt,
QTimer,
Signal,
Slot,
)
from PySide6.QtGui import QColor, QIcon, QImage, QImageReader, QKeySequence, QPainter, QPixmap
from PySide6.QtWidgets import (
QApplication,
@@ -688,6 +699,9 @@ class SuiteTaskState:
show_history: bool = False
worker: object = None
thread: object = None
generation_run_token: str = ""
generation_stop_requested: bool = False
generation_terminal_streak: int = 0
pull_worker: object = None
pull_thread: object = None
import_worker: object = None
@@ -739,6 +753,7 @@ class ProductSuiteTab(QWidget):
self._next_serial = 1
self._displayed_state = None
self._prompt_save_timers = {}
self._generation_run_states = {}
self._original_list_context = None
self._loading = False
self._result_refresh_pending = False
@@ -759,6 +774,10 @@ class ProductSuiteTab(QWidget):
self.elapsed_timer.setInterval(1000)
self.elapsed_timer.timeout.connect(self._refresh_elapsed)
self.elapsed_timer.start()
self.generation_watchdog_timer = QTimer(self)
self.generation_watchdog_timer.setInterval(750)
self.generation_watchdog_timer.timeout.connect(self._check_generation_watchdogs)
self.generation_watchdog_timer.start()
if self._prompt_template_init_error:
self._status(self._prompt_template_init_error, "danger")
@@ -1324,6 +1343,7 @@ class ProductSuiteTab(QWidget):
destructive=True,
):
return
state.generation_stop_requested = True
state.worker.cancel()
if state.ai_worker is not None:
state.ai_worker.cancel()
@@ -1347,6 +1367,8 @@ class ProductSuiteTab(QWidget):
except Exception as exc:
self._status("清理空临时草稿失败:%s" % _user_error(exc), "danger")
self._release_prompt_save_timer(state)
if state.generation_run_token:
self._generation_run_states.pop(state.generation_run_token, None)
self._retired_states.append(state)
self._states.pop(state.key, None)
self.task_tabs.removeTab(index)
@@ -1884,6 +1906,10 @@ class ProductSuiteTab(QWidget):
token = id(thread)
_PRODUCT_SUITE_THREAD_REFS[token] = (thread, worker)
thread.finished.connect(lambda token=token: _PRODUCT_SUITE_THREAD_REFS.pop(token, None))
run_token = str(getattr(worker, "run_token", "") or "")
if run_token:
thread.setProperty("productSuiteRunToken", run_token)
thread.finished.connect(self._on_generation_thread_finished_signal)
thread.start()
return thread
@@ -2491,14 +2517,22 @@ class ProductSuiteTab(QWidget):
if state is None:
return
if state.generation_running():
if state.generation_stop_requested:
self._status("正在停止当前套图任务", "warning")
return
if self._confirm(
"停止生成套图",
"确认取消当前任务吗?已提交任务会在安全边界停止。",
destructive=True,
):
state.generation_stop_requested = True
state.worker.cancel()
self.generate_button.setText("正在停止...")
self.generate_button.setEnabled(False)
self._log_generation_lifecycle(
state,
state.generation_run_token,
"stop_requested",
)
self._apply_running_state(state)
self._status("已请求停止当前套图任务", "warning")
return
self.start_generation(state)
@@ -2547,26 +2581,38 @@ class ProductSuiteTab(QWidget):
):
return False
self._persist_state(state)
run_token = uuid.uuid4().hex
worker = ProductSuiteGenerateWorker(
state.project_id,
specs,
run_token=run_token,
aspect_ratio=state.settings["ratio"],
db_path=self.db_path,
config=self.config,
cmhub_config_path=self.cmhub_config_path,
)
state.worker = worker
state.generation_run_token = run_token
state.generation_stop_requested = False
state.generation_terminal_streak = 0
state.done = 0
state.failed = 0
state.total = len(specs)
state.started_at = time.monotonic()
state.current_job_ids = []
state.show_history = False
worker.progress.connect(lambda payload, state=state: self._on_generation_progress(state, payload))
worker.finished.connect(lambda result, state=state: self._on_generation_finished(state, result))
worker.cancelled.connect(lambda result, state=state: self._on_generation_finished(state, result))
worker.failed.connect(lambda row, error, state=state: self._on_generation_failed(state, error))
self._generation_run_states[run_token] = state.key
worker.progress.connect(self._on_generation_progress_signal)
worker.finished.connect(self._on_generation_finished_signal)
worker.cancelled.connect(self._on_generation_finished_signal)
worker.failed.connect(self._on_generation_failed_signal)
state.thread = self._start_thread(worker, "商品套图生成")
self._log_generation_lifecycle(
state,
run_token,
"started",
{"total": len(specs), "job_ids": 0},
)
if state is self._displayed_state:
self._loading = True
try:
@@ -2578,52 +2624,367 @@ class ProductSuiteTab(QWidget):
self._status("商品套图生成已开始,共%d张;可切换到其他任务" % len(specs), "info")
return True
def _generation_signal_token(self, payload=None):
token = str((payload or {}).get("run_token") or "")
sender = self.sender()
return token or str(getattr(sender, "run_token", "") or "")
def _generation_state(self, run_token):
token = str(run_token or "")
state = self._states.get(self._generation_run_states.get(token))
if state is None or state.generation_run_token != token:
return None
return state
@Slot(dict)
def _on_generation_progress_signal(self, payload):
token = self._generation_signal_token(payload)
state = self._generation_state(token)
if state is None:
return
self._on_generation_progress(state, payload)
def _on_generation_progress(self, state, payload):
state.total = int(payload.get("total", state.total) or state.total)
if "total" in payload:
state.total = max(0, int(payload.get("total") or 0))
state.done = int(payload.get("done", state.done) or 0)
state.failed = int(payload.get("failed", state.failed) or 0)
job_ids = payload.get("job_ids")
if job_ids:
if job_ids is not None:
state.current_job_ids = [int(job_id) for job_id in job_ids]
state.generation_terminal_streak = 0
if state is self._displayed_state:
self._refresh_results(state)
self._refresh_elapsed()
def _on_generation_failed(self, state, error):
@Slot(int, str)
def _on_generation_failed_signal(self, row, error):
token = self._generation_signal_token()
state = self._generation_state(token)
if state is None:
return
self._log_generation_lifecycle(
state,
token,
"worker_failed",
{"has_error": True},
level="ERROR",
)
self._status("商品套图生成失败:%s" % _user_error(error), "danger")
def _on_generation_finished(self, state, result):
@Slot(dict)
def _on_generation_finished_signal(self, result):
token = self._generation_signal_token(result)
state = self._generation_state(token)
if state is None:
return
self._log_generation_lifecycle(
state,
token,
"worker_finished",
{
"ok": result.get("ok", True),
"success": result.get("success", 0),
"failed": result.get("failed", 0),
"cancelled": result.get(
"cancelled_count",
result.get("cancelled", 0),
),
},
)
if result.get("ok") is False:
self._reconcile_generation_jobs(state, "生成线程异常结束")
if (
state.generation_stop_requested
and not self._generation_job_ids(state)
and result.get("cancelled") is True
):
result = dict(result)
result["total"] = state.total
result["cancelled_count"] = state.total
self._finalize_generation(state, token, result, source="worker")
@Slot()
def _on_generation_thread_finished_signal(self):
sender = self.sender()
token = str(
sender.property("productSuiteRunToken")
if sender is not None
else ""
)
self._handle_generation_thread_finished(token)
def _handle_generation_thread_finished(self, run_token):
state = self._generation_state(run_token)
if state is None:
return
snapshot = self._generation_job_snapshot(state)
had_active = bool(snapshot["active"])
self._log_generation_lifecycle(
state,
run_token,
"thread_finished_fallback",
snapshot,
level="WARNING" if snapshot["active"] else "INFO",
)
if had_active:
self._reconcile_generation_jobs(state, "生成线程已结束")
snapshot = self._generation_job_snapshot(state)
result = self._generation_result_from_snapshot(snapshot)
if not snapshot["job_ids"] and state.generation_stop_requested:
result.update(
{
"total": state.total,
"cancelled_count": state.total,
}
)
elif had_active:
result.update(
{
"ok": False,
"error": "生成线程已结束,部分任务可稍后继续查询",
}
)
self._finalize_generation(
state,
run_token,
result,
source="thread_finished",
)
def _check_generation_watchdogs(self):
for state in list(self._states.values()):
token = state.generation_run_token
if not token or not state.generation_running():
continue
snapshot = self._generation_job_snapshot(state)
if snapshot["all_terminal"]:
state.generation_terminal_streak += 1
else:
state.generation_terminal_streak = 0
if state.generation_terminal_streak < 2:
continue
self._log_generation_lifecycle(
state,
token,
"terminal_watchdog_finalize",
snapshot,
level="WARNING",
)
self._finalize_generation(
state,
token,
self._generation_result_from_snapshot(snapshot),
source="terminal_watchdog",
)
def _generation_job_ids(self, state):
job_ids = list(state.current_job_ids)
if not job_ids and state.worker is not None:
worker_job_ids = getattr(state.worker, "job_ids", [])
if not isinstance(worker_job_ids, (list, tuple, set)):
worker_job_ids = []
job_ids = [
int(job_id)
for job_id in list(worker_job_ids or [])
]
if job_ids:
state.current_job_ids = job_ids
return job_ids
def _generation_job_snapshot(self, state):
job_ids = self._generation_job_ids(state)
counts = {
"success": 0,
"failed": 0,
"cancelled": 0,
"active": 0,
"job_ids": len(job_ids),
"all_terminal": False,
}
if not job_ids:
counts["active"] = max(0, int(state.total or 0))
return counts
for job_id in job_ids:
try:
job = image_studio.get_job(job_id, path=self.db_path)
except Exception:
job = None
status = str(getattr(job, "status", "") or "")
if status == "succeeded":
counts["success"] += 1
elif status in {"failed", "expired"}:
counts["failed"] += 1
elif status == "cancelled":
counts["cancelled"] += 1
else:
counts["active"] += 1
counts["all_terminal"] = (
len(job_ids) == int(state.total or 0)
and counts["active"] == 0
)
return counts
def _generation_result_from_snapshot(self, snapshot):
return {
"total": int(snapshot.get("job_ids", 0) or 0),
"success": int(snapshot.get("success", 0) or 0),
"failed": int(snapshot.get("failed", 0) or 0),
"cancelled_count": int(snapshot.get("cancelled", 0) or 0),
}
def _reconcile_generation_jobs(self, state, reason):
for job_id in self._generation_job_ids(state):
try:
job = image_studio.get_job(job_id, path=self.db_path)
if job is None or job.status in {
"succeeded",
"failed",
"expired",
"cancelled",
}:
continue
recovery = (
image_studio.JOB_RECOVERY_RESUME
if job.task_id
else image_studio.JOB_RECOVERY_REGENERATE
)
image_studio.update_job_status(
job.id,
"cancelled",
error="%s,任务可稍后继续处理" % reason,
recovery_action=recovery,
path=self.db_path,
)
except Exception as exc:
self._status(
"商品套图任务状态收尾失败:%s" % _user_error(exc),
"danger",
)
def _finalize_generation(self, state, run_token, result, *, source):
if self._generation_state(run_token) is not state:
return False
snapshot = self._generation_job_snapshot(state)
if snapshot["job_ids"] and (
snapshot["all_terminal"]
or source in {"thread_finished", "terminal_watchdog"}
):
result = dict(result or {})
result.update(self._generation_result_from_snapshot(snapshot))
result = dict(result or {})
stop_requested = state.generation_stop_requested
success = int(result.get("success", 0) or 0)
failed = int(result.get("failed", 0) or 0)
cancelled = int(
result.get("cancelled_count", result.get("cancelled", 0)) or 0
)
total = int(result.get("total", state.total) or state.total)
active = max(0, total - success - failed - cancelled)
elapsed = (
int(max(0, time.monotonic() - state.started_at))
if state.started_at
else 0
)
self._generation_run_states.pop(run_token, None)
state.generation_run_token = ""
state.generation_stop_requested = False
state.generation_terminal_streak = 0
state.worker = None
state.thread = None
state.done = int(result.get("success", 0) or 0) + int(result.get("failed", 0) or 0) + int(
result.get("cancelled", 0) or 0
)
state.failed = int(result.get("failed", state.failed) or 0)
state.total = int(result.get("total", state.total) or state.total)
elapsed = int(max(0, time.monotonic() - state.started_at)) if state.started_at else 0
state.done = success + failed + cancelled
state.failed = failed
state.total = total
state.started_at = None
if state is self._displayed_state:
self._apply_running_state(state)
self._refresh_results(state)
self._refresh_elapsed()
self._log_generation_lifecycle(
state,
run_token,
"finalized",
{
"source": source,
"total": total,
"success": success,
"failed": failed,
"cancelled": cancelled,
"active": active,
"elapsed_seconds": elapsed,
},
level="WARNING" if active or result.get("ok") is False else "INFO",
)
if state.key not in self._states:
return
if result.get("ok") is False:
return True
if active or result.get("ok") is False:
message = str(
result.get("error")
or "生成线程已结束,部分任务可稍后继续查询"
)
if state is self._displayed_state:
self._message("商品套图生成失败", _user_error(result.get("error")))
self._message("商品套图生成未完整结束", _user_error(message))
else:
self._status("套图任务%d生成失败" % state.serial, "danger")
return
success = int(result.get("success", 0) or 0)
cancelled = int(result.get("cancelled", 0) or 0)
self._status(
"套图任务%d生成未完整结束" % state.serial,
"danger",
)
return True
if stop_requested or cancelled:
if state is self._displayed_state:
self._message(
"商品套图生成已停止",
"本轮共%d张:成功%d张,失败%d张,停止%d张;"
"已提交任务可稍后继续查询;总用时%d秒。"
% (total, success, failed, cancelled, elapsed),
icon=QMessageBox.Information,
)
self._status(
"商品套图生成已停止:成功%d张,失败%d张,停止%d张"
% (success, failed, cancelled),
"warning",
)
return True
if state is self._displayed_state:
self._message(
"商品套图生成完成",
"本轮共%d张:成功%d张,失败%d张,停止%d张;总用时%d秒。"
% (state.total, success, state.failed, cancelled, elapsed),
% (total, success, failed, cancelled, elapsed),
icon=QMessageBox.Information,
)
self._status("商品套图生成完成:成功%d张,失败%d张" % (success, state.failed), "success")
self._status(
"商品套图生成完成:成功%d张,失败%d张" % (success, failed),
"success",
)
return True
def _log_generation_lifecycle(
self,
state,
run_token,
event,
payload=None,
*,
level="INFO",
):
data = {
"run_token": str(run_token or "")[:8],
"project_id": getattr(state, "project_id", None),
"event": str(event or ""),
}
data.update(dict(payload or {}))
try:
diagnostics.write_diagnostic_log(
"商品套图生成生命周期",
level=level,
step="product_suite_generation",
task_id=getattr(state, "project_id", None),
item_id=getattr(state, "item_id", None),
payload=data,
log_dir=appconfig.diagnostic_log_dir(self.config),
)
except Exception:
pass
def _apply_running_state(self, state):
generation_running = state.generation_running()
@@ -2648,7 +3009,11 @@ class ProductSuiteTab(QWidget):
row.set_controls_enabled(not generation_running)
self.generate_button.setEnabled(True)
if generation_running:
self.generate_button.setText("停止生成")
self.generate_button.setText(
"正在停止..."
if state.generation_stop_requested
else "停止生成"
)
self.generate_button.setStyleSheet(
"QPushButton { background: #cf222e; color: white; border-color: #a40e26; font-weight: 600; }"
"QPushButton:hover { background: #a40e26; }"
@@ -2909,4 +3274,5 @@ class ProductSuiteTab(QWidget):
worker.cancel()
for state in list(self._states.values()) + list(self._retired_states):
self._release_prompt_save_timer(state)
self._generation_run_states.clear()
super().closeEvent(event)
+11 -5
View File
@@ -303,6 +303,7 @@ class ProductSuiteGenerateWorker(BaseWorker):
project_id,
job_specs,
*,
run_token="",
aspect_ratio="1:1",
db_path=None,
config=None,
@@ -311,10 +312,12 @@ class ProductSuiteGenerateWorker(BaseWorker):
super().__init__()
self.project_id = int(project_id)
self.job_specs = [dict(spec) for spec in (job_specs or [])]
self.run_token = str(run_token or "")
self.aspect_ratio = str(aspect_ratio or "1:1")
self.db_path = db_path
self.config = config
self.cmhub_config_path = cmhub_config_path
self.job_ids = []
self._done = 0
self._failed = 0
self._lock = threading.Lock()
@@ -325,8 +328,6 @@ class ProductSuiteGenerateWorker(BaseWorker):
raise ValueError("商品套图生成任务不能为空")
jobs = []
for spec in self.job_specs:
if self.should_cancel():
break
jobs.append(
image_studio.create_job(
self.project_id,
@@ -338,12 +339,14 @@ class ProductSuiteGenerateWorker(BaseWorker):
path=self.db_path,
)
)
self.job_ids = [job.id for job in jobs]
self.progress.emit(
{
"run_token": self.run_token,
"total": len(jobs),
"done": 0,
"failed": 0,
"job_ids": [job.id for job in jobs],
"job_ids": list(self.job_ids),
}
)
@@ -356,10 +359,11 @@ class ProductSuiteGenerateWorker(BaseWorker):
if event.get("result") != "success":
self._failed += 1
progress = {
"run_token": self.run_token,
"total": len(jobs),
"done": self._done,
"failed": self._failed,
"job_ids": [job.id for job in jobs],
"job_ids": list(self.job_ids),
}
self.progress.emit(progress)
@@ -373,7 +377,9 @@ class ProductSuiteGenerateWorker(BaseWorker):
on_event=on_event,
)
summary["project_id"] = self.project_id
summary["job_ids"] = [job.id for job in jobs]
summary["job_ids"] = list(self.job_ids)
summary["cancelled_count"] = int(summary.get("cancelled", 0) or 0)
summary["run_token"] = self.run_token
return summary
+60 -15
View File
@@ -6,7 +6,7 @@ import os
import threading
import time
import urllib.parse
from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
from concurrent.futures import CancelledError, FIRST_COMPLETED, ThreadPoolExecutor, wait
from . import ai, appconfig, image_studio
from .version import APP_VERSION
@@ -185,7 +185,11 @@ def run_jobs(
for job in job_list
}
while futures:
done, _ = wait(set(futures), return_when=FIRST_COMPLETED)
done, _ = wait(
set(futures),
timeout=0.2,
return_when=FIRST_COMPLETED,
)
for future in done:
futures.pop(future)
try:
@@ -252,7 +256,14 @@ def _run_one_job(job_id, runtime, config, image_root, aspect_ratio, db_path, sho
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)
saved_path = _download_and_save_job_image(
request_result,
out_path,
config,
on_event,
job.id,
should_stop=should_stop,
)
try:
_raise_if_stopped(should_stop)
except ImageStudioGenerationError:
@@ -281,17 +292,27 @@ 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:
status = "cancelled" if "停止" in str(exc) else "failed"
cancelled = isinstance(exc, CancelledError) or "停止" in str(exc)
status = "cancelled" if cancelled else "failed"
error = "用户已停止,已提交任务可稍后继续查询" if cancelled else str(exc)
current_job = image_studio.get_job(job.id, path=db_path)
updated = image_studio.update_job_status(
job.id,
status,
error=str(exc),
error=error,
recovery_action=_recovery_action_for_job(current_job),
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)}
_notify(
on_event,
{
"job_id": job.id,
"step": "job_done",
"result": status,
"detail": error,
},
)
return {"job": updated, "status": status, "error": error}
def _submit_or_resume_job(
@@ -395,6 +416,7 @@ def _poll_job(job_id, task_id, runtime, request_result, db_path, should_stop, on
read_timeout=ai.CMHUB_IMAGE_POLL_READ_TIMEOUT_SECONDS,
headers_extra={"X-Client-Version": str(APP_VERSION)},
)
_raise_if_stopped(should_stop)
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})
@@ -424,21 +446,44 @@ def _poll_job(job_id, task_id, runtime, request_result, db_path, should_stop, on
raise ImageStudioGenerationError("cmhub 生图任务状态返回格式错误")
def _download_and_save_job_image(request_result, out_path, config, on_event, job_id):
def _download_and_save_job_image(
request_result,
out_path,
config,
on_event,
job_id,
should_stop=None,
):
_raise_if_stopped(should_stop)
_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"),
)
try:
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"),
should_stop=should_stop,
)
except CancelledError as exc:
raise ImageStudioGenerationError(
"用户已停止,已提交任务可稍后继续查询"
) from exc
_raise_if_stopped(should_stop)
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),
)
try:
_raise_if_stopped(should_stop)
except ImageStudioGenerationError:
try:
if os.path.isfile(saved_path):
os.remove(saved_path)
finally:
raise
_notify(on_event, {"job_id": job_id, "step": "cover_download", "result": "success"})
return saved_path