Spaces:
Running
Running
# tasks.py | |
import asyncio | |
from typing import Dict | |
from uuid import uuid4 | |
# A dictionary to keep track of active tasks | |
tasks: Dict[str, asyncio.Task] = {} | |
def cleanup_task(task_id: str): | |
""" | |
Remove a completed or canceled task from the global `tasks` dictionary. | |
""" | |
tasks.pop(task_id, None) # Remove the task if it exists | |
def create_task(coroutine): | |
""" | |
Create a new asyncio task and add it to the global task dictionary. | |
""" | |
task_id = str(uuid4()) # Generate a unique ID for the task | |
task = asyncio.create_task(coroutine) # Create the task | |
# Add a done callback for cleanup | |
task.add_done_callback(lambda t: cleanup_task(task_id)) | |
tasks[task_id] = task | |
return task_id, task | |
def get_task(task_id: str): | |
""" | |
Retrieve a task by its task ID. | |
""" | |
return tasks.get(task_id) | |
def list_tasks(): | |
""" | |
List all currently active task IDs. | |
""" | |
return list(tasks.keys()) | |
async def stop_task(task_id: str): | |
""" | |
Cancel a running task and remove it from the global task list. | |
""" | |
task = tasks.get(task_id) | |
if not task: | |
raise ValueError(f"Task with ID {task_id} not found.") | |
task.cancel() # Request task cancellation | |
try: | |
await task # Wait for the task to handle the cancellation | |
except asyncio.CancelledError: | |
# Task successfully canceled | |
tasks.pop(task_id, None) # Remove it from the dictionary | |
return {"status": True, "message": f"Task {task_id} successfully stopped."} | |
return {"status": False, "message": f"Failed to stop task {task_id}."} | |