"""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"} 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 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 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.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 "任务执行结束，但业务检查未通过" 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 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"]) 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 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="模型请求已暂停，请先恢复模型连接再继续批次") 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)) 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: 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], "repeat_count": 1, "request_id": None, } ) return revised else: raise ValueError("不支持的批次操作") return store.put("batch", batch_id, batch)