"""Task management service for Web-Local system dispatch (P4-M2). Handles task creation, dispatch, status tracking, progress updates, and result reception from local simulation executor. P0-5: When a task completes and has a plan_id, results are automatically synced to the SimulationResult table so the Plan detail page can display them. """ import json import os import uuid from datetime import datetime from typing import Dict, List, Optional, Any from pathlib import Path from ..database import SessionLocal from ..models.task import Task from ..models.simulation_plan import SimulationPlan from ..models.simulation_result import SimulationResult from ..metrics_constants import METRIC_KEYS class TaskManager: """Manages simulation tasks between Web and Local systems.""" TASK_STATUSES = [ "pending", # Created, waiting for dispatch "dispatched", # Sent to local executor "running", # Local executor is running "completed", # All points completed successfully "failed", # Task failed "cancelled", # User cancelled ] def __init__(self, output_dir: Optional[str] = None): self.output_dir = output_dir or os.path.join( os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "output", "tasks" ) os.makedirs(self.output_dir, exist_ok=True) def create_task( self, plan_id: Optional[int], plan_data: Dict[str, Any], parameters: List[Dict[str, Any]], task_name: Optional[str] = None, priority: int = 5, created_by: str = "web", task_type: str = "scan", loop_id: Optional[str] = None, batch_id: Optional[int] = None, point_ids: Optional[List] = None, dynamic: bool = False, thermal_mode: str = "steady", ) -> Dict[str, Any]: """Create a new simulation task. Args: plan_id: Associated plan ID (optional) plan_data: Full plan data (boundary conditions, topology, etc.) parameters: List of parameter sets to simulate task_name: Optional task name priority: Task priority (1-10, higher = more urgent) created_by: Creator identifier thermal_mode: Thermal simulation mode for every point of this task. "off" (electromagnetic only) / "steady" (EM + steady-state thermal, default) / "coupled" (magnetic-thermal coupled). Returns: Created task dict """ task_uuid = str(uuid.uuid4())[:8] task_name = task_name or f"task_{task_uuid}" # Validate thermal_mode (task-level switch, see plan/executor docs). _valid_thermal_modes = ("off", "steady", "coupled") if thermal_mode not in _valid_thermal_modes: raise ValueError( "thermal_mode must be one of %s, got %r" % (_valid_thermal_modes, thermal_mode) ) task_dir = os.path.join(self.output_dir, f"{datetime.now().strftime('%Y%m%d_%H%M%S')}_{task_name}") os.makedirs(task_dir, exist_ok=True) # Write task file for local executor task_file = os.path.join(task_dir, "task.json") task_payload = { "task_id": task_uuid, "task_name": task_name, "plan_id": plan_id, "plan_data": plan_data, "parameters": parameters, "priority": priority, "created_at": datetime.now().isoformat(), "total_points": len(parameters), "task_type": task_type, "loop_id": loop_id, "batch_id": batch_id, "point_ids": point_ids or [], "dynamic": bool(dynamic), "thermal_mode": thermal_mode, } with open(task_file, "w", encoding="utf-8") as f: json.dump(task_payload, f, ensure_ascii=False, indent=2) # Save to database with SessionLocal() as db: db_task = Task( task_id=task_uuid, task_name=task_name, plan_id=plan_id, status="pending", priority=priority, total_points=len(parameters), completed_points=0, task_dir=task_dir, task_file=task_file, created_by=created_by, created_at=datetime.now(), task_type=task_type, loop_id=loop_id, batch_id=batch_id, point_ids=json.dumps(point_ids or [], ensure_ascii=False), dynamic=1 if dynamic else 0, ) db.add(db_task) db.commit() db.refresh(db_task) return self._task_to_dict(db_task) def dispatch_task(self, task_id: str) -> Dict[str, Any]: """Mark task as dispatched and ready for local executor. Uses an atomic conditional UPDATE so exactly one executor instance can claim a pending task; a concurrent second claim gets a ValueError. Args: task_id: Task UUID Returns: Updated task dict Raises: ValueError: task not found, or already claimed / not pending. """ from sqlalchemy import update with SessionLocal() as db: task = db.query(Task).filter(Task.task_id == task_id).first() if not task: raise ValueError(f"Task {task_id} not found") result = db.execute( update(Task) .where(Task.task_id == task_id, Task.status == "pending") .values(status="dispatched", dispatched_at=datetime.now()) ) db.commit() if result.rowcount == 0: # another instance already claimed it, or the state moved on raise ValueError( "Task %s is not pending (status: %s)" % (task_id, task.status) ) db.refresh(task) return self._task_to_dict(task) def update_progress( self, task_id: str, current_point: int, total_points: Optional[int] = None, current_params: Optional[Dict[str, Any]] = None, elapsed_time: Optional[float] = None, ) -> Dict[str, Any]: """Update task progress from local executor. Args: task_id: Task UUID current_point: Current point index (0-based) total_points: Total points (optional, will use stored value) current_params: Current parameter values being simulated elapsed_time: Elapsed time in seconds Returns: Updated task dict """ with SessionLocal() as db: task = db.query(Task).filter(Task.task_id == task_id).first() if not task: raise ValueError(f"Task {task_id} not found") if task.status in ("dispatched", "running"): task.status = "running" task.started_at = task.started_at or datetime.now() task.completed_points = current_point if total_points: task.total_points = total_points # Update progress metadata progress_data = { "current_point": current_point, "current_params": current_params, "elapsed_time": elapsed_time, "updated_at": datetime.now().isoformat(), } existing_progress = json.loads(task.progress_data or "{}") existing_progress.update(progress_data) task.progress_data = json.dumps(existing_progress, ensure_ascii=False) db.commit() db.refresh(task) return self._task_to_dict(task) def report_results( self, task_id: str, results: List[Dict[str, Any]], metrics: Optional[Dict[str, Any]] = None, logs: Optional[str] = None, duration: Optional[float] = None, status: str = "completed", ) -> Dict[str, Any]: """Report final results from local executor. Args: task_id: Task UUID results: List of simulation result dicts metrics: Aggregated metrics logs: Program logs duration: Total duration in seconds status: Final status (completed/failed) Returns: Updated task dict """ with SessionLocal() as db: task = db.query(Task).filter(Task.task_id == task_id).first() if not task: raise ValueError(f"Task {task_id} not found") task.status = status task.completed_at = datetime.now() task.completed_points = len(results) if duration: task.duration = duration # Save results to file results_file = os.path.join(task.task_dir, "results.json") with open(results_file, "w", encoding="utf-8") as f: json.dump({"results": results, "metrics": metrics, "logs": logs}, f, ensure_ascii=False, indent=2) task.results_file = results_file if metrics: task.result_metrics = json.dumps(metrics, ensure_ascii=False) db.commit() db.refresh(task) # P0-5: Auto-sync results to plan's SimulationResult table plan_id = getattr(task, 'plan_id', None) if plan_id is not None: self._sync_results_to_plan(plan_id, results) return self._task_to_dict(task) def get_task(self, task_id: str) -> Optional[Dict[str, Any]]: """Get task by ID.""" with SessionLocal() as db: task = db.query(Task).filter(Task.task_id == task_id).first() if not task: return None return self._task_to_dict(task) def list_tasks( self, status: Optional[str] = None, plan_id: Optional[int] = None, limit: int = 50, offset: int = 0, ) -> Dict[str, Any]: """List tasks with filters.""" with SessionLocal() as db: query = db.query(Task) if status: query = query.filter(Task.status == status) if plan_id: query = query.filter(Task.plan_id == plan_id) total = query.count() tasks = query.order_by(Task.created_at.desc()).offset(offset).limit(limit).all() return { "total": total, "limit": limit, "offset": offset, "tasks": [self._task_to_dict(t) for t in tasks], } def cancel_task(self, task_id: str) -> Dict[str, Any]: """Cancel a pending or running task.""" with SessionLocal() as db: task = db.query(Task).filter(Task.task_id == task_id).first() if not task: raise ValueError(f"Task {task_id} not found") if task.status in ("completed", "failed", "cancelled"): raise ValueError(f"Task {task_id} already finished (status: {task.status})") task.status = "cancelled" task.completed_at = datetime.now() db.commit() db.refresh(task) return self._task_to_dict(task) def delete_task(self, task_id: str) -> None: """Delete a finished task (terminal status only) and its task file. Only completed/failed/cancelled tasks may be deleted; active tasks must be cancelled first so a running executor is never orphaned. """ with SessionLocal() as db: task = db.query(Task).filter(Task.task_id == task_id).first() if not task: raise ValueError(f"Task {task_id} not found") if task.status not in ("completed", "failed", "cancelled"): raise ValueError( f"Task {task_id} is {task.status}; cancel it before deleting" ) task_file = task.task_file db.delete(task) db.commit() # Best-effort removal of the on-disk task file after the row is gone. if task_file and os.path.exists(task_file): try: os.remove(task_file) except OSError: pass def get_task_results(self, task_id: str) -> Optional[Dict[str, Any]]: """Get task results file content.""" task = self.get_task(task_id) if not task or not task.get("results_file"): return None if not os.path.exists(task["results_file"]): return None with open(task["results_file"], "r", encoding="utf-8") as f: return json.load(f) def _task_to_dict(self, task: Task) -> Dict[str, Any]: """Convert Task ORM object to dict.""" result = { "id": task.id, "task_id": task.task_id, "task_name": task.task_name, "plan_id": task.plan_id, "status": task.status, "priority": task.priority, "total_points": task.total_points, "completed_points": task.completed_points, "progress": round((task.completed_points / task.total_points * 100), 1) if task.total_points else 0, "task_dir": task.task_dir, "task_file": task.task_file, "results_file": task.results_file, "created_by": task.created_by, "created_at": task.created_at.isoformat() if task.created_at else None, "dispatched_at": task.dispatched_at.isoformat() if task.dispatched_at else None, "started_at": task.started_at.isoformat() if task.started_at else None, "completed_at": task.completed_at.isoformat() if task.completed_at else None, "duration": task.duration, "task_type": task.task_type or "scan", "loop_id": task.loop_id, "batch_id": task.batch_id, "point_ids": json.loads(task.point_ids) if task.point_ids else [], "dynamic": bool(task.dynamic), } if task.progress_data: try: result["progress_data"] = json.loads(task.progress_data) except (json.JSONDecodeError, Exception): result["progress_data"] = None if task.result_metrics: try: result["result_metrics"] = json.loads(task.result_metrics) except (json.JSONDecodeError, Exception): result["result_metrics"] = None return result def _sync_results_to_plan(self, plan_id: int, results: List[Dict[str, Any]]) -> None: """Sync task results to SimulationResult table (P0-5). Clears existing results for this plan and inserts new ones. Each result has params (scan params + fixed params) and metrics. """ try: with SessionLocal() as db: plan = db.query(SimulationPlan).filter(SimulationPlan.id == plan_id).first() if not plan: return # Clear existing results db.query(SimulationResult).filter(SimulationResult.plan_id == plan_id).delete() metric_keys = set(METRIC_KEYS) for idx, r in enumerate(results): params = r.get("params", {}) or {} # Metrics may be at top level or under "metrics" key metrics = r.get("metrics", {}) or {} if not metrics: # Extract metric fields from top level for k, v in r.items(): if k in metric_keys and isinstance(v, (int, float)): metrics[k] = float(v) status = r.get("status", "OK") solve_time = float(r.get("solve_time_s", r.get("seconds", 0)) or 0) error_msg = r.get("error", "") or "" result = SimulationResult( plan_id=plan_id, run_index=idx + 1, status=status, solve_time_s=solve_time, error_message=error_msg[:500] if error_msg else "", ) result.set_params({k: float(v) for k, v in params.items() if isinstance(v, (int, float))}) result.set_metrics({k: float(v) for k, v in metrics.items() if isinstance(v, (int, float))}) # Lossless full export archive (all EM+thermal fields # with units) for the complete-data Excel/CSV export. raw_flat = r.get("raw_flat") if isinstance(raw_flat, list) and raw_flat: result.set_raw(raw_flat) db.add(result) # Update plan status based on results ok_count = sum(1 for r in results if r.get("status") == "OK") if ok_count == len(results) and len(results) > 0: plan.status = "completed" elif ok_count > 0: plan.status = "completed" elif len(results) > 0: plan.status = "failed" db.commit() except Exception as e: # Sync failure should not break task completion print(f"[TaskManager] Failed to sync results to plan {plan_id}: {e}") # ------------------------------------------------------------------ # Overview & Executor monitoring (P1-9) # ------------------------------------------------------------------ def get_overview(self) -> Dict[str, Any]: """Get task overview statistics for monitoring dashboard.""" with SessionLocal() as db: total = db.query(Task).count() pending = db.query(Task).filter(Task.status == "pending").count() dispatched = db.query(Task).filter(Task.status == "dispatched").count() running = db.query(Task).filter(Task.status == "running").count() completed = db.query(Task).filter(Task.status == "completed").count() failed = db.query(Task).filter(Task.status == "failed").count() # Recent tasks (last 20) recent = ( db.query(Task) .order_by(Task.created_at.desc()) .limit(20) .all() ) recent_list = [self._task_to_dict(t) for t in recent] # Active (pending + dispatched + running) active = ( db.query(Task) .filter(Task.status.in_(["pending", "dispatched", "running"])) .order_by(Task.priority.desc(), Task.created_at.asc()) .all() ) active_list = [self._task_to_dict(t) for t in active] return { "total": total, "pending": pending, "dispatched": dispatched, "running": running, "completed": completed, "failed": failed, "active_tasks": active_list, "recent_tasks": recent_list, } def executor_heartbeat(self, executor_id: str, status: str = "idle", current_task: Optional[str] = None, progress: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: """Record executor heartbeat and return next pending task if available. Args: executor_id: Unique executor identifier. status: executor status (idle / running / error). current_task: task_id currently being executed, if any. progress: current progress data. Returns: Dict with executor status and optionally a task to claim. """ now = datetime.now() # Store heartbeat in memory (simple approach) if not hasattr(self, "_executors"): self._executors = {} self._executors[executor_id] = { "executor_id": executor_id, "status": status, "current_task": current_task, "last_heartbeat": now.isoformat(), "progress": progress, } return {"status": "ok", "executor_id": executor_id, "server_time": now.isoformat()} def get_executor_status(self) -> Dict[str, Any]: """Get all registered executors and their status.""" if not hasattr(self, "_executors"): self._executors = {} now = datetime.now() executors = [] for eid, info in self._executors.items(): last_hb = info.get("last_heartbeat") online = False if last_hb: try: hb_time = datetime.fromisoformat(last_hb) online = (now - hb_time).total_seconds() < 60 except Exception: online = False executors.append({**info, "online": online}) return {"executors": executors, "count": len(executors)} # Global singleton _task_manager: Optional[TaskManager] = None def get_task_manager() -> TaskManager: """Get or create global TaskManager singleton.""" global _task_manager if _task_manager is None: _task_manager = TaskManager() return _task_manager