| 1 | # Copyright (C) 2025 AIDC-AI |
| 2 | # |
| 3 | # Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | # you may not use this file except in compliance with the License. |
| 5 | # You may obtain a copy of the License at |
| 6 | # http://www.apache.org/licenses/LICENSE-2.0 |
| 7 | # Unless required by applicable law or agreed to in writing, software |
| 8 | # distributed under the License is distributed on an "AS IS" BASIS, |
| 9 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 10 | # See the License for the specific language governing permissions and |
| 11 | # limitations under the License. |
| 12 | |
| 13 | """ |
| 14 | Task Manager |
| 15 | |
| 16 | In-memory task management for video generation jobs. |
| 17 | """ |
| 18 | |
| 19 | import asyncio |
| 20 | import uuid |
| 21 | from datetime import datetime, timedelta |
| 22 | from typing import Dict, List, Optional, Callable |
| 23 | from loguru import logger |
| 24 | |
| 25 | from api.tasks.models import Task, TaskStatus, TaskType, TaskProgress |
| 26 | from api.config import api_config |
| 27 | |
| 28 | |
| 29 | class TaskManager: |
| 30 | """ |
| 31 | Task manager for handling async video generation tasks |
| 32 | |
| 33 | Features: |
| 34 | - In-memory storage (can be replaced with Redis later) |
| 35 | - Task lifecycle management |
| 36 | - Progress tracking |
| 37 | - Auto cleanup of old tasks |
| 38 | """ |
| 39 | |
| 40 | def __init__(self): |
| 41 | self._tasks: Dict[str, Task] = {} |
| 42 | self._task_futures: Dict[str, asyncio.Task] = {} |
| 43 | self._cleanup_task: Optional[asyncio.Task] = None |
| 44 | self._running = False |
| 45 | |
| 46 | async def start(self): |
| 47 | """Start task manager and cleanup scheduler""" |
| 48 | if self._running: |
| 49 | logger.warning("Task manager already running") |
| 50 | return |
| 51 | |
| 52 | self._running = True |
| 53 | self._cleanup_task = asyncio.create_task(self._cleanup_loop()) |
| 54 | logger.info("✅ Task manager started") |
| 55 | |
| 56 | async def stop(self): |
| 57 | """Stop task manager and cancel all tasks""" |
| 58 | self._running = False |
| 59 | |
| 60 | # Cancel cleanup task |
| 61 | if self._cleanup_task: |
| 62 | self._cleanup_task.cancel() |
| 63 | try: |
| 64 | await self._cleanup_task |
| 65 | except asyncio.CancelledError: |
| 66 | pass |
| 67 | |
| 68 | # Cancel all running tasks |
| 69 | for task_id, future in self._task_futures.items(): |
| 70 | if not future.done(): |
| 71 | future.cancel() |
| 72 | logger.info(f"Cancelled task: {task_id}") |
| 73 | |
| 74 | self._tasks.clear() |
| 75 | self._task_futures.clear() |
| 76 | logger.info("✅ Task manager stopped") |
| 77 | |
| 78 | def create_task( |
| 79 | self, |
| 80 | task_type: TaskType, |
| 81 | request_params: Optional[dict] = None |
| 82 | ) -> Task: |
| 83 | """ |
| 84 | Create a new task |
| 85 | |
| 86 | Args: |
| 87 | task_type: Type of task |
| 88 | request_params: Original request parameters |
| 89 | |
| 90 | Returns: |
| 91 | Created task |
| 92 | """ |
| 93 | task_id = str(uuid.uuid4()) |
| 94 | task = Task( |
| 95 | task_id=task_id, |
| 96 | task_type=task_type, |
| 97 | status=TaskStatus.PENDING, |
| 98 | request_params=request_params, |
| 99 | ) |
| 100 | |
| 101 | self._tasks[task_id] = task |
| 102 | logger.info(f"Created task {task_id} ({task_type})") |
| 103 | return task |
| 104 | |
| 105 | async def execute_task( |
| 106 | self, |
| 107 | task_id: str, |
| 108 | coro_func: Callable, |
| 109 | *args, |
| 110 | **kwargs |
| 111 | ): |
| 112 | """ |
| 113 | Execute task asynchronously |
| 114 | |
| 115 | Args: |
| 116 | task_id: Task ID |
| 117 | coro_func: Async function to execute |
| 118 | *args: Positional arguments |
| 119 | **kwargs: Keyword arguments |
| 120 | """ |
| 121 | task = self._tasks.get(task_id) |
| 122 | if not task: |
| 123 | logger.error(f"Task {task_id} not found") |
| 124 | return |
| 125 | |
| 126 | # Create async task |
| 127 | async def _execute(): |
| 128 | try: |
| 129 | task.status = TaskStatus.RUNNING |
| 130 | task.started_at = datetime.now() |
| 131 | logger.info(f"Task {task_id} started") |
| 132 | |
| 133 | # Execute the actual work |
| 134 | result = await coro_func(*args, **kwargs) |
| 135 | |
| 136 | # Update task with result |
| 137 | task.status = TaskStatus.COMPLETED |
| 138 | task.result = result |
| 139 | task.completed_at = datetime.now() |
| 140 | logger.info(f"Task {task_id} completed") |
| 141 | |
| 142 | except Exception as e: |
| 143 | task.status = TaskStatus.FAILED |
| 144 | task.error = str(e) |
| 145 | task.completed_at = datetime.now() |
| 146 | logger.error(f"Task {task_id} failed: {e}") |
| 147 | |
| 148 | # Start execution |
| 149 | future = asyncio.create_task(_execute()) |
| 150 | self._task_futures[task_id] = future |
| 151 | |
| 152 | def get_task(self, task_id: str) -> Optional[Task]: |
| 153 | """Get task by ID""" |
| 154 | return self._tasks.get(task_id) |
| 155 | |
| 156 | def list_tasks( |
| 157 | self, |
| 158 | status: Optional[TaskStatus] = None, |
| 159 | limit: int = 100 |
| 160 | ) -> List[Task]: |
| 161 | """ |
| 162 | List tasks with optional filtering |
| 163 | |
| 164 | Args: |
| 165 | status: Filter by status |
| 166 | limit: Maximum number of tasks to return |
| 167 | |
| 168 | Returns: |
| 169 | List of tasks |
| 170 | """ |
| 171 | tasks = list(self._tasks.values()) |
| 172 | |
| 173 | if status: |
| 174 | tasks = [t for t in tasks if t.status == status] |
| 175 | |
| 176 | # Sort by created_at descending |
| 177 | tasks.sort(key=lambda t: t.created_at, reverse=True) |
| 178 | |
| 179 | return tasks[:limit] |
| 180 | |
| 181 | def update_progress( |
| 182 | self, |
| 183 | task_id: str, |
| 184 | current: int, |
| 185 | total: int, |
| 186 | message: str = "" |
| 187 | ): |
| 188 | """ |
| 189 | Update task progress |
| 190 | |
| 191 | Args: |
| 192 | task_id: Task ID |
| 193 | current: Current progress |
| 194 | total: Total steps |
| 195 | message: Progress message |
| 196 | """ |
| 197 | task = self._tasks.get(task_id) |
| 198 | if not task: |
| 199 | return |
| 200 | |
| 201 | percentage = (current / total * 100) if total > 0 else 0 |
| 202 | task.progress = TaskProgress( |
| 203 | current=current, |
| 204 | total=total, |
| 205 | percentage=percentage, |
| 206 | message=message |
| 207 | ) |
| 208 | |
| 209 | def cancel_task(self, task_id: str) -> bool: |
| 210 | """ |
| 211 | Cancel a running task |
| 212 | |
| 213 | Args: |
| 214 | task_id: Task ID |
| 215 | |
| 216 | Returns: |
| 217 | True if cancelled, False otherwise |
| 218 | """ |
| 219 | task = self._tasks.get(task_id) |
| 220 | if not task: |
| 221 | return False |
| 222 | |
| 223 | # Do not cancel already-terminal tasks |
| 224 | if task.status in [TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.CANCELLED]: |
| 225 | return False |
| 226 | |
| 227 | # Cancel future if running |
| 228 | future = self._task_futures.get(task_id) |
| 229 | if future and not future.done(): |
| 230 | future.cancel() |
| 231 | |
| 232 | # Update task status |
| 233 | task.status = TaskStatus.CANCELLED |
| 234 | task.completed_at = datetime.now() |
| 235 | logger.info(f"Cancelled task {task_id}") |
| 236 | return True |
| 237 | |
| 238 | async def _cleanup_loop(self): |
| 239 | """Periodically clean up old completed tasks""" |
| 240 | while self._running: |
| 241 | try: |
| 242 | await asyncio.sleep(api_config.task_cleanup_interval) |
| 243 | self._cleanup_old_tasks() |
| 244 | except asyncio.CancelledError: |
| 245 | break |
| 246 | except Exception as e: |
| 247 | logger.error(f"Error in cleanup loop: {e}") |
| 248 | |
| 249 | def _cleanup_old_tasks(self): |
| 250 | """Remove old completed/failed tasks""" |
| 251 | cutoff_time = datetime.now() - timedelta(seconds=api_config.task_retention_time) |
| 252 | |
| 253 | tasks_to_remove = [] |
| 254 | for task_id, task in self._tasks.items(): |
| 255 | if task.status in [TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.CANCELLED]: |
| 256 | if task.completed_at and task.completed_at < cutoff_time: |
| 257 | tasks_to_remove.append(task_id) |
| 258 | |
| 259 | for task_id in tasks_to_remove: |
| 260 | del self._tasks[task_id] |
| 261 | if task_id in self._task_futures: |
| 262 | del self._task_futures[task_id] |
| 263 | |
| 264 | if tasks_to_remove: |
| 265 | logger.info(f"Cleaned up {len(tasks_to_remove)} old tasks") |
| 266 | |
| 267 | |
| 268 | # Global task manager instance |
| 269 | task_manager = TaskManager() |
| 270 | |
| 271 |