Files

510 lines
22 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 RECOVERY_ERRORS, store, new_id
from apps.admin_console.services.task_import import case_goal, execution_title
from artemis.config.settings import settings
from artemis.runtime.task_limits import task_timeout
logger = logging.getLogger(__name__)
TERMINAL = {"completed", "failed", "cancelled", "skipped", "interrupted"}
MODEL_PAUSE_REASON = "模型请求已暂停,请先恢复模型连接再继续批次"
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"),
"execution_config": item.get("execution_config"),
"submission": {
key: item.get(key)
for key in (
"goal",
"profile",
"device_serial",
"runtime_overrides",
"expected_output",
"enable_outputter",
"verification_level",
"explorer_mode",
"locked_app_package",
"app_path",
)
},
"formats": formats if formats is not None else store.preferences()["formats"],
**extra,
}
)
def resolve_recovery_error(job: dict) -> bool:
"""Archive a superseded recovery warning without erasing execution failures."""
if job.get("error") not in RECOVERY_ERRORS:
return False
job.setdefault("recovery_events", []).append(
{"message": job.pop("error"), "resolved_at": time.time()}
)
return True
def recover_batches():
"""Recovery never silently repeats a potentially side-effecting interrupted task."""
from apps.admin_console.database.repositories.session_repository import SessionRepository
from artemis.runtime import DeviceExecutionLock
repository = SessionRepository(store.path)
live_sessions = {
str(owner.session_id)
for owner in DeviceExecutionLock.get_active_owners().values()
if owner.session_id
}
live_sessions.update(
str(row["session_id"])
for row in repository.get_all_sessions()
if row.get("status") not in TERMINAL and repository.process_is_alive(row.get("pid"))
)
live_sessions.update(
str(item["session_id"])
for item in state.queue_items
if isinstance(item, dict) and item.get("session_id") and item.get("status") not in TERMINAL
)
for job in store.jobs():
if job["status"] == "completed" and resolve_recovery_error(job):
store.save_job(job)
if str(job["session_id"]) in live_sessions:
if job["status"] not in TERMINAL and str(job["session_id"]) not in state.active_runs:
job["recovery_waiting"] = True
store.save_job(job)
logger.info("Recovery preserved live session %s", job["session_id"])
continue
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(),
)
if job["status"] == "completed":
resolve_recovery_error(job)
logger.info("Recovery resolved session %s to %s", job["session_id"], job["status"])
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"}:
if any(str(job["session_id"]) in live_sessions for job in store.jobs(batch["id"])):
continue
recovery_reason = "服务已重启;请检查中断任务,点击继续执行未开始任务"
batch.setdefault("recovery_events", []).append(
{
"message": recovery_reason,
"recovered_at": time.time(),
"previous_reason": batch.get("reason"),
"previous_pause_source": batch.get("pause_source"),
}
)
batch.update(
status="paused",
reason=batch.get("reason") or recovery_reason,
pause_source="recovery",
)
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.pop("recovery_waiting", None)
if status == "completed":
resolve_recovery_error(job)
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 report.get("error")
or (report.get("test_result") or {}).get("reason")
or "任务未完成或验证未通过,请查看报告中的具体原因"
)
elif job["status"] == "failed" and report.get("error") and not job.get("error"):
job["error"] = report["error"]
if job["status"] == "failed" and report["outcome"].get("code") == "initialization_failed":
job["failure_code"] = "initialization_failed"
batch = store.get("batch", job.get("batch_id")) if job.get("batch_id") else None
if (
batch
and (batch["status"] == "running" or batch.get("pause_source") == "model")
and any(j["status"] == "pending" for j in store.jobs(batch["id"]))
):
batch.update(
status="paused",
pause_source="initialization",
reason="任务初始化失败;请处理环境问题后继续剩余用例:"
+ (job.get("error") or "请查看任务报告"),
)
store.put("batch", batch["id"], batch)
# 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)
job["finalized_at"] = time.time()
store.save_job(job)
async def reconcile_recovered_jobs():
"""Finalize surviving workers once they exit; never resubmit their goals."""
from apps.admin_console.database.repositories.session_repository import SessionRepository
from artemis.runtime import DeviceExecutionLock
from artemis.runtime.task_result import read_worker_outcome
waiting = [job for job in store.jobs() if job.get("recovery_waiting")]
if not waiting:
return
repository = SessionRepository(store.path)
owners = {str(owner.session_id) for owner in DeviceExecutionLock.get_active_owners().values()}
for job in waiting:
sid = str(job["session_id"])
row = repository.get_session_by_id(sid) or {}
if sid in owners or repository.process_is_alive(row.get("pid")):
continue
outcome = read_worker_outcome(sid)
status = job.get("terminal_status") or outcome.get("task_status") or row.get("status")
if status not in TERMINAL:
status = "interrupted"
job["error"] = job.get("error") or "服务中断;未完成任务需手动重跑"
store.save_job(job)
await on_terminal({"session_id": sid}, status)
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
await reconcile_recovered_jobs()
from apps.admin_console.services.exploration_recovery import reconcile_explorations
reconcile_explorations()
for batch in store.records("batch"):
if batch["status"] not in {"running", "paused"}:
continue
jobs = store.jobs(batch["id"])
if (
batch["status"] == "paused"
and batch.get("pause_source") == "model"
and not state.is_paused
):
batch.update(status="running", reason="")
batch.pop("pause_source", None)
store.put("batch", batch["id"], batch)
active = [j for j in jobs if j["status"] in {"running", "queued", "finalizing"}]
if active:
if state.is_paused and batch["status"] == "running":
batch.update(status="paused", reason=MODEL_PAUSE_REASON, pause_source="model")
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
from artemis.runtime.task_result import read_worker_outcome
if read_worker_outcome(str(job["session_id"])):
continue # Execution has ended; the worker supervisor bounds cleanup.
limit = task_timeout({}, job, batch)
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)
from artemis.runtime.task_result import save_failure_result
try:
await asyncio.to_thread(
save_failure_result,
str(job["session_id"]),
job["error"],
code="task_timeout",
)
finally:
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=MODEL_PAUSE_REASON, pause_source="model")
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")
computer = options.get("execution_target") == "computer"
devices = await device_pool.try_list_devices_async() if not computer else []
rejection = (
None
if computer
or (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), pause_source="device")
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=execution_title(job),
)
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:
repeat_count = payload.get("repeat_count", 1)
if type(repeat_count) is not int or repeat_count < 1:
raise ValueError("循环次数必须为正整数")
cases = payload["cases"]
total = len(cases) * repeat_count
if not cases or total > settings.ARTEMIS_BATCH_MAX_EXECUTIONS:
raise ValueError(
f"请选择用例;用例数 × 循环次数不能超过 {settings.ARTEMIS_BATCH_MAX_EXECUTIONS} 次执行"
)
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 settings.ARTEMIS_TASK_TIMEOUT_SECONDS,
"stop_on_failure": payload.get("stop_on_failure", False),
"request_id": payload.get("request_id"),
"repeat_count": repeat_count,
"case_count": len(cases),
"total_executions": total,
}
jobs = []
for iteration in range(1, repeat_count + 1):
for number, case in enumerate(cases, 1):
jobs.append(
{
"session_id": new_id(),
"batch_id": bid,
"position": len(jobs) + 1,
"iteration": iteration,
"repeat_count": repeat_count,
"case_position": number,
"case": case,
"status": "pending",
"formats": batch["formats"],
"created_at": batch["created_at"],
"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(
"batch:" + batch_id, *("session:" + j["session_id"] for j in jobs)
):
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="用户暂停;当前任务继续,后续任务等待", pause_source="user"
)
elif action == "resume":
if batch["status"] not in {"draft", "paused"}:
raise ValueError("只能开始草稿或继续已暂停批次")
batch.update(status="running", reason="")
batch.pop("pause_source", None)
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],
"repeat_count": 1,
"request_id": None,
}
)
return revised
else:
raise ValueError("不支持的批次操作")
return store.put("batch", batch_id, batch)