pcb-afm-simulation-system/web/backend/app/services/batch_scheduler.py

"""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

from ..services.task_contract import merge_adaptive_fields


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,
                 task_type: str = "scan",
                 loop_id: Optional[str] = None,
                 batch_id: Optional[int] = None,
                 point_ids: Optional[List] = None,
                 dynamic: bool = False) -> 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,
        }
        # P3-M4: adaptive-batch fields align scheduler tasks with Task ORM.
        task_info = merge_adaptive_fields(
            task_info, task_type=task_type, loop_id=loop_id,
            batch_id=batch_id, point_ids=point_ids, dynamic=dynamic,
        )
        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"),
            "task_type": task.get("task_type", "scan"),
            "loop_id": task.get("loop_id"),
            "batch_id": task.get("batch_id"),
            "point_ids": task.get("point_ids") or [],
            "dynamic": bool(task.get("dynamic", False)),
        }

    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