"""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. """ 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 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", ) -> 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 Returns: Created task dict """ task_uuid = str(uuid.uuid4())[:8] task_name = task_name or f"task_{task_uuid}" 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), } 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(), ) 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. Args: task_id: Task UUID 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 != "pending": raise ValueError(f"Task {task_id} is not pending (status: {task.status})") task.status = "dispatched" task.dispatched_at = datetime.now() db.commit() 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) 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 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, } 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 # 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