338 lines
14 KiB
Python
338 lines
14 KiB
Python
"""Persistent batch coordinator layered on the existing device queue."""
|
|
|
|
import asyncio
|
|
import logging
|
|
import os
|
|
import shutil
|
|
import time
|
|
from uuid import UUID
|
|
|
|
from apps.admin_console.core.state import state
|
|
from apps.admin_console.services.automation_store import store, new_id
|
|
from apps.admin_console.services.task_import import case_goal, case_title
|
|
|
|
logger = logging.getLogger(__name__)
|
|
TERMINAL = {"completed", "failed", "cancelled", "skipped", "interrupted"}
|
|
background_tasks: set[asyncio.Task] = set()
|
|
|
|
|
|
def spawn(coro):
|
|
task = asyncio.create_task(coro)
|
|
background_tasks.add(task)
|
|
task.add_done_callback(background_tasks.discard)
|
|
return task
|
|
|
|
|
|
def register_submission(item: dict, case: dict | None = None, formats=None, **extra):
|
|
sid = str(item["session_id"])
|
|
if store.job(sid):
|
|
return
|
|
store.save_job(
|
|
{
|
|
"session_id": sid,
|
|
"case": case or {"name": item.get("goal", sid)[:120], "goal": item.get("goal", sid)},
|
|
"status": "queued",
|
|
"created_at": time.time(),
|
|
"device_serial": item.get("device_serial"),
|
|
"formats": formats if formats is not None else store.preferences()["formats"],
|
|
**extra,
|
|
}
|
|
)
|
|
|
|
|
|
def recover_batches():
|
|
"""Recovery never silently repeats a potentially side-effecting interrupted task."""
|
|
for job in store.jobs():
|
|
if job["status"] in {"running", "queued", "finalizing", "task_paused", "resuming"}:
|
|
job.update(
|
|
status=job.get("terminal_status", "interrupted"),
|
|
error=job.get("error")
|
|
or ("" if job.get("terminal_status") else "服务中断；未完成任务需手动重跑"),
|
|
ended_at=job.get("ended_at") or time.time(),
|
|
)
|
|
store.save_job(job)
|
|
for batch in store.records("batch"):
|
|
if has_legacy_replay(batch["id"]):
|
|
for job in store.jobs(batch["id"]):
|
|
if job["status"] not in TERMINAL:
|
|
job.update(status="cancelled", error="操作回放功能已移除", ended_at=time.time())
|
|
store.save_job(job)
|
|
if batch["status"] not in TERMINAL:
|
|
batch.update(status="cancelled", reason="操作回放功能已移除；保留历史供查看")
|
|
store.put("batch", batch["id"], batch)
|
|
continue
|
|
if batch["status"] in {"running", "pausing", "paused"}:
|
|
batch.update(
|
|
status="paused", reason="服务已重启；请检查中断任务，点击继续执行未开始任务"
|
|
)
|
|
for job in store.jobs(batch["id"]):
|
|
if job["status"] in {"running", "queued", "finalizing", "task_paused", "resuming"}:
|
|
job.update(
|
|
status="interrupted", error="服务中断，需手动重跑", ended_at=time.time()
|
|
)
|
|
store.save_job(job)
|
|
store.put("batch", batch["id"], batch)
|
|
|
|
|
|
async def on_terminal(item: dict, status: str):
|
|
from apps.admin_console.services.report_export import export_session, session_report
|
|
|
|
sid = str(item["session_id"])
|
|
job = store.job(sid)
|
|
if not job:
|
|
register_submission(item)
|
|
job = store.job(sid)
|
|
if job is None:
|
|
raise RuntimeError("任务元数据写入失败")
|
|
if job.get("timed_out"):
|
|
status = "failed"
|
|
job.update(status="finalizing", terminal_status=status, ended_at=time.time())
|
|
store.save_job(job)
|
|
try:
|
|
job["status"] = status
|
|
report = await asyncio.to_thread(session_report, sid)
|
|
if status != "interrupted" and (
|
|
report["verdict"] == "验证失败"
|
|
or report["outcome"].get("task_status")
|
|
in {
|
|
"blocked",
|
|
"partial",
|
|
"failed",
|
|
}
|
|
):
|
|
job["status"] = "failed"
|
|
job["error"] = job.get("error") or "任务执行结束，但业务检查未通过"
|
|
elif job["status"] == "failed" and report.get("error") and not job.get("error"):
|
|
job["error"] = report["error"]
|
|
# Persist terminal classification before export, but the coordinator waits for finalization.
|
|
store.save_job({**job, "status": "finalizing", "terminal_status": job["status"]})
|
|
await asyncio.to_thread(export_session, sid, job["formats"])
|
|
job["report_ready"] = True
|
|
except Exception as exc:
|
|
logger.exception("Report finalization failed for %s", sid)
|
|
job["report_error"] = str(exc)[:500]
|
|
job.pop("terminal_status", None)
|
|
store.save_job(job)
|
|
|
|
|
|
async def tick_batches():
|
|
from apps.admin_console.server import task_queue_service
|
|
from artemis.runtime import device_pool
|
|
from apps.admin_console.services.report_export import export_batch
|
|
|
|
for batch in store.records("batch"):
|
|
if batch["status"] not in {"running", "paused"}:
|
|
continue
|
|
jobs = store.jobs(batch["id"])
|
|
active = [j for j in jobs if j["status"] in {"running", "queued", "finalizing"}]
|
|
if active:
|
|
if state.is_paused:
|
|
batch.update(status="paused", reason="模型请求已暂停，请先恢复模型连接再继续批次")
|
|
store.put("batch", batch["id"], batch)
|
|
for job in active:
|
|
# The queue records started_at at submission, then replaces it
|
|
# when the worker actually starts. Waiting for a device must
|
|
# not consume this case's execution budget.
|
|
if job["status"] != "running":
|
|
continue
|
|
limit = int(
|
|
(job.get("case", {}).get("task_config") or {}).get("timeout_seconds")
|
|
or batch.get("timeout_seconds", 1800)
|
|
)
|
|
task_state = getattr(state, "active_runs", {}).get(str(job["session_id"])) or {}
|
|
excluded = float(task_state.get("non_execution_seconds") or 0)
|
|
if task_state.get("non_execution_started_at"):
|
|
excluded += time.time() - float(task_state["non_execution_started_at"])
|
|
if (
|
|
not job.get("timed_out")
|
|
and time.time() - job.get("started_at", time.time()) - excluded > limit
|
|
):
|
|
job.update(
|
|
error=f"任务执行超时（{limit} 秒，不含暂停/排队时间）", timed_out=True
|
|
)
|
|
store.save_job(job)
|
|
task_queue_service.stop_tasks(session_id=job["session_id"])
|
|
# on_terminal records the failed case. The next pending
|
|
# case then runs normally unless stop_on_failure is set.
|
|
continue
|
|
if batch["status"] != "running":
|
|
continue
|
|
pending = [j for j in jobs if j["status"] == "pending"]
|
|
failed = [j for j in jobs if j["status"] in {"failed", "cancelled", "interrupted"}]
|
|
if batch.get("stop_on_failure") and failed and pending:
|
|
for job in pending:
|
|
job.update(status="skipped", error="本批次设置为失败即停止", ended_at=time.time())
|
|
store.save_job(job)
|
|
pending = []
|
|
if not pending and any(
|
|
j["status"] in {"task_paused", "resuming", "running", "queued", "finalizing"}
|
|
for j in jobs
|
|
):
|
|
continue
|
|
if not pending:
|
|
batch.update(status="completed", ended_at=time.time())
|
|
try:
|
|
await asyncio.to_thread(export_batch, batch["id"])
|
|
batch["report_ready"] = True
|
|
except Exception as exc:
|
|
logger.exception("Batch report generation failed")
|
|
batch["report_error"] = str(exc)[:500]
|
|
store.put("batch", batch["id"], batch)
|
|
continue
|
|
if state.is_paused:
|
|
batch.update(status="paused", reason="模型请求已暂停，请先恢复模型连接再继续批次")
|
|
store.put("batch", batch["id"], batch)
|
|
continue
|
|
job = pending[0]
|
|
options = job["case"].get("task_config") or {}
|
|
serial = options.get("device_serial") or batch.get("device_serial")
|
|
devices = await device_pool.try_list_devices_async()
|
|
rejection = (
|
|
None
|
|
if serial and any(d.serial == serial and d.state == "device" for d in devices or [])
|
|
else "无法确认所选设备在线且已授权，请恢复设备连接后继续批次"
|
|
)
|
|
if rejection:
|
|
batch.update(status="paused", reason=str(rejection))
|
|
store.put("batch", batch["id"], batch)
|
|
continue
|
|
if (store.get("batch", batch["id"]) or {}).get("status") != "running":
|
|
continue
|
|
job.update(status="queued", started_at=time.time())
|
|
store.save_job(job)
|
|
try:
|
|
response = await task_queue_service.enqueue_tasks(
|
|
[case_goal(job["case"])],
|
|
profile=options.get("profile") or batch.get("profile", "flash"),
|
|
expected_output="用简体中文报告实际结果与证据。明确区分执行完成、验证通过、验证失败和前置条件未满足。",
|
|
enable_outputter=True,
|
|
device_serial=serial,
|
|
verification_level=options.get("verification_level"),
|
|
explorer_mode=options.get("explorer_mode"),
|
|
locked_app_package=options.get("locked_app_package"),
|
|
runtime_overrides=options,
|
|
session_id=job["session_id"],
|
|
ingress="batch",
|
|
display_name=case_title(job["case"]),
|
|
)
|
|
if response.get("status") == "rejected" or not response.get("tasks"):
|
|
raise ValueError(response.get("error") or "队列拒绝接收任务")
|
|
bound = response["tasks"][0].get("device_serial")
|
|
if bound:
|
|
job["device_serial"] = bound
|
|
store.save_job(job)
|
|
except Exception as exc:
|
|
job.update(status="pending", error=str(exc)[:500])
|
|
store.save_job(job)
|
|
batch.update(status="paused", reason=str(exc)[:500])
|
|
store.put("batch", batch["id"], batch)
|
|
|
|
|
|
async def batch_loop():
|
|
recover_batches()
|
|
while True:
|
|
try:
|
|
await tick_batches()
|
|
except Exception:
|
|
logger.exception("Batch coordinator failed; persisted work is retained")
|
|
await asyncio.sleep(1)
|
|
|
|
|
|
def has_legacy_replay(batch_id: str) -> bool:
|
|
return any(
|
|
j.get("workflow_id") or j.get("execution_mode") == "strict" or j.get("adaptive_source")
|
|
for j in store.jobs(batch_id)
|
|
)
|
|
|
|
|
|
def create_batch(payload: dict) -> dict:
|
|
bid = new_id()
|
|
batch = {
|
|
"id": bid,
|
|
"name": payload["name"],
|
|
"status": "draft",
|
|
"created_at": time.time(),
|
|
"formats": payload.get("formats") or store.preferences()["formats"],
|
|
"profile": payload.get("profile", "flash"),
|
|
"device_serial": payload.get("device_serial"),
|
|
"timeout_seconds": payload.get("timeout_seconds")
|
|
or int(os.getenv("ARTEMIS_BATCH_TASK_TIMEOUT", "1800")),
|
|
"stop_on_failure": payload.get("stop_on_failure", False),
|
|
"request_id": payload.get("request_id"),
|
|
}
|
|
jobs = []
|
|
for number, case in enumerate(payload["cases"], 1):
|
|
jobs.append(
|
|
{
|
|
"session_id": new_id(),
|
|
"batch_id": bid,
|
|
"position": number,
|
|
"case": case,
|
|
"status": "pending",
|
|
"formats": batch["formats"],
|
|
"created_at": time.time(),
|
|
"device_serial": batch["device_serial"],
|
|
}
|
|
)
|
|
return store.create_batch(batch, jobs, payload)
|
|
|
|
|
|
def delete_batch(batch_id: str) -> bool:
|
|
"""Delete an inactive batch, its sessions and its generated reports."""
|
|
from artemis.data_engine.storage import StorageManager
|
|
from apps.admin_console.services import report_export
|
|
|
|
jobs = store.begin_delete_batch(batch_id)
|
|
if jobs is None:
|
|
return False
|
|
storage = StorageManager(store.path, report_export.TRACES_PATH)
|
|
mcp_exports = report_export.EXPORT_ROOT.parent / "mcp_exports"
|
|
with report_export.EXPORT_LOCK:
|
|
for job in jobs:
|
|
session_id = str(UUID(job["session_id"]))
|
|
storage.delete_session(UUID(session_id))
|
|
report_dir = report_export.EXPORT_ROOT / "sessions" / session_id
|
|
if report_dir.is_dir():
|
|
shutil.rmtree(report_dir)
|
|
mcp_session_dir = mcp_exports / "sessions" / session_id
|
|
if mcp_session_dir.is_dir():
|
|
shutil.rmtree(mcp_session_dir)
|
|
batch_dir = report_export.EXPORT_ROOT / "batches" / batch_id
|
|
if batch_dir.is_dir():
|
|
shutil.rmtree(batch_dir)
|
|
mcp_batch_dir = mcp_exports / "batches" / batch_id
|
|
if mcp_batch_dir.is_dir():
|
|
shutil.rmtree(mcp_batch_dir)
|
|
return store.delete_batch(batch_id)
|
|
|
|
|
|
def batch_action(batch_id: str, action: str) -> dict:
|
|
batch = store.get("batch", batch_id)
|
|
if not batch:
|
|
raise ValueError("批次不存在")
|
|
if has_legacy_replay(batch_id):
|
|
raise ValueError("操作回放功能已移除；历史记录仅供查看")
|
|
if action == "pause":
|
|
if batch["status"] not in {"running", "paused"}:
|
|
raise ValueError("只能暂停正在执行的批次")
|
|
batch.update(status="paused", reason="用户暂停；当前任务继续，后续任务等待")
|
|
elif action == "resume":
|
|
if batch["status"] not in {"draft", "paused"}:
|
|
raise ValueError("只能开始草稿或继续已暂停批次")
|
|
batch.update(status="running", reason="")
|
|
elif action == "retry_failed":
|
|
if batch["status"] == "running":
|
|
raise ValueError("请先暂停批次")
|
|
failed = [
|
|
j for j in store.jobs(batch_id) if j["status"] in {"failed", "cancelled", "interrupted"}
|
|
]
|
|
if not failed:
|
|
raise ValueError("没有可重跑的失败或中断任务")
|
|
revised = create_batch(
|
|
{**batch, "name": batch["name"] + " · 失败项重跑", "cases": [j["case"] for j in failed]}
|
|
)
|
|
return revised
|
|
else:
|
|
raise ValueError("不支持的批次操作")
|
|
return store.put("batch", batch_id, batch)
|