"""Batch task scheduler for multi-task simulation (P4-M3). Priority queue, parallel execution, checkpoint/resume, real-time stats. """ import json import os import threading from datetime import datetime from typing import Any, Callable, Dict, List, Optional from collections import deque class BatchScheduler: """Batch simulation task scheduler with priority queue.""" def __init__(self, state_file: Optional[str] = None): self._queue: deque = deque() self._running: Dict[str, Dict[str, Any]] = {} self._completed: List[Dict[str, Any]] = [] # B1 fix: use RLock (reentrant) because _notify_callbacks -> get_statistics # acquires the same lock, and _notify_callbacks is called from within # locked sections (get_next_task, cancel_task). self._lock = threading.RLock() self._max_parallel = int(os.environ.get("MAX_PARALLEL_TASKS", "2")) self._state_file = state_file or os.path.join( os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "output", "scheduler_state.json" ) self._callbacks: List[Callable] = [] self._load_state() def add_task(self, task_id: str, task_name: str, priority: int = 5, parameters: Optional[List[Dict[str, Any]]] = None, plan_data: Optional[Dict[str, Any]] = None, wait_for: Optional[List[str]] = None) -> Dict[str, Any]: task_info = { "task_id": task_id, "task_name": task_name, "priority": priority, "status": "queued", "parameters": parameters or [], "plan_data": plan_data or {}, "wait_for": wait_for or [], "enqueued_at": datetime.now().isoformat(), "started_at": None, "completed_at": None, "current_point": 0, "total_points": len(parameters or []), "error": None, } with self._lock: inserted = False for i, existing in enumerate(self._queue): if priority > existing["priority"]: self._queue.insert(i, task_info) inserted = True break if not inserted: self._queue.append(task_info) self._save_state_locked() self._notify_callbacks() return task_info def get_next_task(self) -> Optional[Dict[str, Any]]: with self._lock: if len(self._running) >= self._max_parallel: return None for task in list(self._queue): if task["wait_for"]: waiting = set(task["wait_for"]) done = {t["task_id"] for t in self._completed if t["status"] == "completed"} if not waiting.issubset(done): continue self._queue.remove(task) task["status"] = "running" task["started_at"] = datetime.now().isoformat() self._running[task["task_id"]] = task self._save_state_locked() self._notify_callbacks() return task return None def update_task_progress(self, task_id: str, current_point: int, total_points: Optional[int] = None) -> None: with self._lock: if task_id in self._running: self._running[task_id]["current_point"] = current_point if total_points: self._running[task_id]["total_points"] = total_points self._save_state_locked() self._notify_callbacks() def complete_task(self, task_id: str, status: str = "completed", error: Optional[str] = None) -> None: with self._lock: if task_id in self._running: task = self._running.pop(task_id) task["status"] = status task["completed_at"] = datetime.now().isoformat() task["error"] = error self._completed.append(task) if len(self._completed) > 100: self._completed = self._completed[-100:] self._save_state_locked() self._notify_callbacks() def cancel_task(self, task_id: str) -> bool: with self._lock: for task in list(self._queue): if task["task_id"] == task_id: self._queue.remove(task) task["status"] = "cancelled" task["completed_at"] = datetime.now().isoformat() self._completed.append(task) self._save_state_locked() self._notify_callbacks() return True if task_id in self._running: task = self._running.pop(task_id) task["status"] = "cancelled" task["completed_at"] = datetime.now().isoformat() self._completed.append(task) self._save_state_locked() self._notify_callbacks() return True return False def get_statistics(self) -> Dict[str, Any]: with self._lock: queued = list(self._queue) running = list(self._running.values()) completed = list(self._completed) tpq = sum(t["total_points"] for t in queued) tpr = sum(t["total_points"] for t in running) cpr = sum(t["current_point"] for t in running) tc = len([t for t in completed if t["status"] == "completed"]) tf = len([t for t in completed if t["status"] == "failed"]) total_all = tpq + tpr done_all = cpr + sum(t["total_points"] for t in completed if t["status"] == "completed") overall = round((done_all / total_all * 100), 1) if total_all > 0 else 0 return { "queued_count": len(queued), "running_count": len(running), "completed_count": len(completed), "successful_count": tc, "failed_count": tf, "max_parallel": self._max_parallel, "total_points_queued": tpq, "total_points_running": tpr, "completed_points_running": cpr, "overall_progress": overall, "queued_tasks": [self._summary(t) for t in queued[:20]], "running_tasks": [self._summary(t) for t in running], "recent_completed": [self._summary(t) for t in completed[-10:]], "timestamp": datetime.now().isoformat(), } def _summary(self, task: Dict[str, Any]) -> Dict[str, Any]: progress = round((task["current_point"] / task["total_points"] * 100), 1) if task["total_points"] > 0 else 0 return { "task_id": task["task_id"], "task_name": task["task_name"], "status": task["status"], "priority": task["priority"], "current_point": task["current_point"], "total_points": task["total_points"], "progress": progress, "enqueued_at": task.get("enqueued_at"), "started_at": task.get("started_at"), "completed_at": task.get("completed_at"), "error": task.get("error"), } def register_callback(self, callback: Callable) -> None: self._callbacks.append(callback) def _notify_callbacks(self) -> None: stats = self.get_statistics() for cb in self._callbacks: try: cb(stats) except Exception: pass def _save_state_locked(self) -> None: try: os.makedirs(os.path.dirname(self._state_file), exist_ok=True) state = {"queue": list(self._queue), "running": self._running, "completed": self._completed[-50:], "saved_at": datetime.now().isoformat()} with open(self._state_file, "w", encoding="utf-8") as f: json.dump(state, f, ensure_ascii=False, indent=2) except Exception: pass def _load_state(self) -> None: try: if os.path.exists(self._state_file): with open(self._state_file, "r", encoding="utf-8") as f: state = json.load(f) with self._lock: self._queue = deque(state.get("queue", [])) for task in state.get("running", {}).values(): task["status"] = "queued" task["error"] = "Recovered from previous session" self._queue.append(task) self._completed = state.get("completed", []) except Exception: pass _scheduler: Optional[BatchScheduler] = None def get_scheduler() -> BatchScheduler: global _scheduler if _scheduler is None: _scheduler = BatchScheduler() return _scheduler