tech: add async runner celery

Author: stephenprn

.github/workflows/deploy-to-scaleway.yml

@@ -38,9 +38,16 @@ jobs:
             source .env || exit 1
             set +a # Disable auto-export of variables
             
+            xargs -a Aptfile sudo apt install -y || exit 1
+            
+            sudo systemctl start redis-server || exit 1
+            sudo systemctl enable redis-server || exit 1
+
             source venv/bin/activate || exit 1
             python -m pip install -r requirements.txt || exit 1
             python manage.py migrate || exit 1
+
+            python -m celery -A aigle worker --loglevel=info || exit 1
             sudo systemctl restart gunicorn_aigle || exit 1
             
             echo "Deployment completed successfully!"

.vscode/settings.json

@@ -1,3 +1,4 @@
 gdal-bin
 libgdal-dev
 python3-gdal
+redis-server

Aptfile

@@ -0,0 +1,3 @@
+from .celery import app as celery_app
+
+__all__ = ("celery_app",)

aigle/__init__.py

@@ -0,0 +1,8 @@
+import os
+from celery import Celery
+
+os.environ.setdefault("DJANGO_SETTINGS_MODULE", "aigle.settings")
+
+app = Celery("aigle")
+app.config_from_object("django.conf:settings", namespace="CELERY")
+app.autodiscover_tasks()

aigle/celery.py

@@ -20,6 +20,8 @@
 import logging_loki  # noqa: F401
 from core.utils.logs import scaleway_logger  # noqa: F401
 
+from celery import Celery  # noqa: F401
+
 DEPLOYMENT_DATETIME = datetime.now()
 
 # Build paths inside the project like this: BASE_DIR / 'subdir'.
@@ -339,3 +341,17 @@
 DEFAULT_AUTO_FIELD = "django.db.models.BigAutoField"
 
 MEDIA_ROOT = os.path.join(BASE_DIR, "media")
+
+
+# Celery Configuration
+CELERY_BROKER_URL = "redis://localhost:6379/0"
+CELERY_RESULT_BACKEND = "redis://localhost:6379/0"
+CELERY_ACCEPT_CONTENT = ["json"]
+CELERY_TASK_SERIALIZER = "json"
+CELERY_RESULT_SERIALIZER = "json"
+CELERY_TIMEZONE = "UTC"
+
+# Optional: Task routing
+CELERY_ROUTES = {
+    "myapp.tasks.run_management_command": {"queue": "management_commands"},
+}

aigle/settings.py

@@ -0,0 +1,37 @@
+# management/commands/run_async.py
+
+from typing import Any, List, Optional
+from django.core.management.base import BaseCommand, CommandParser
+
+from core.utils.tasks import AsyncCommandService
+
+
+class Command(BaseCommand):
+    help = "Run Django management commands asynchronously"
+
+    def add_arguments(self, parser: CommandParser) -> None:
+        parser.add_argument("command", type=str, help="Command to run")
+        parser.add_argument("--args", nargs="*", default=[], help="Command arguments")
+        parser.add_argument("--wait", action="store_true", help="Wait for completion")
+
+    def handle(self, *args: Any, **options: Any) -> None:
+        command_name: str = options["command"]
+        command_args: List[str] = options.get("args", [])
+        wait: bool = options.get("wait", False)
+
+        task_id: str = AsyncCommandService.run_command_async(
+            command_name, *command_args
+        )
+        self.stdout.write(f"Task started with ID: {task_id}")
+
+        if wait:
+            self.stdout.write("Waiting for completion...")
+            result: Optional[Any] = AsyncCommandService.get_task_result(task_id)
+            while result is None:
+                import time
+
+                time.sleep(1)
+                result = AsyncCommandService.get_task_result(task_id)
+
+            self.stdout.write("Task completed:")
+            self.stdout.write(str(result))

core/management/commands/run_async.py

@@ -0,0 +1,30 @@
+from rest_framework import serializers
+
+
+class RunCommandSerializer(serializers.Serializer):
+    command = serializers.CharField(max_length=100, required=True)
+    args = serializers.ListField(
+        child=serializers.CharField(), required=False, default=list
+    )
+    kwargs = serializers.DictField(required=False, default=dict)
+
+    def validate_command(self, value: str) -> str:
+        if not value.strip():
+            raise serializers.ValidationError("Command name cannot be empty")
+        return value.strip()
+
+
+class TaskStatusSerializer(serializers.Serializer):
+    task_id = serializers.CharField(read_only=True)
+    status = serializers.CharField(read_only=True)
+    result = serializers.JSONField(read_only=True, allow_null=True)
+    traceback = serializers.CharField(read_only=True, allow_null=True)
+
+
+class CancelTaskSerializer(serializers.Serializer):
+    task_id = serializers.CharField(required=True)
+
+    def validate_task_id(self, value: str) -> str:
+        if not value.strip():
+            raise serializers.ValidationError("Task ID cannot be empty")
+        return value.strip()

core/serializers/run_command.py

@@ -11,6 +11,7 @@
 from core.views.object_type import ObjectTypeViewSet
 from core.views.object_type_category import ObjectTypeCategoryViewSet
 from core.views.parcel import ParcelViewSet
+from core.views.run_command import CommandAsyncViewSet
 from core.views.statistics.validation_status_evolution import (
     StatisticsValidationStatusEvolutionView,
 )
@@ -62,6 +63,7 @@
     "detection-object", DetectionObjectViewSet, basename="DetectionObjectViewSet"
 )
 router.register("detection-data", DetectionDataViewSet, basename="DetectionDataViewSet")
+router.register("run-command", CommandAsyncViewSet, basename="CommandAsyncViewSet")
 
 urlpatterns = router.urls
 

core/urls.py

@@ -0,0 +1,69 @@
+from io import StringIO
+from celery import shared_task
+from celery.result import AsyncResult
+from django.core.management import call_command
+from django.core.management.base import CommandError
+from typing import Dict, Any, Optional, Union
+
+
+@shared_task(bind=True)
+def run_management_command(
+    self, command_name: str, *args: Any, **kwargs: Any
+) -> Dict[str, Union[str, Any]]:
+    try:
+        output = StringIO()
+        call_command(command_name, *args, stdout=output, stderr=output, **kwargs)
+        return {
+            "status": "success",
+            "output": output.getvalue(),
+            "task_id": self.request.id,
+        }
+    except CommandError as e:
+        return {"status": "error", "error": str(e), "task_id": self.request.id}
+    except Exception as e:
+        self.retry(countdown=60, max_retries=3)
+        return {
+            "status": "error",
+            "error": f"Unexpected error: {str(e)}",
+            "task_id": self.request.id,
+        }
+
+
+@shared_task
+def run_custom_command(command_name: str, **options: Any) -> str:
+    output = StringIO()
+    try:
+        call_command(command_name, stdout=output, stderr=output, **options)
+        return output.getvalue()
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+
+class AsyncCommandService:
+    @staticmethod
+    def run_command_async(command_name: str, *args: Any, **kwargs: Any) -> str:
+        task = run_management_command.delay(command_name, *args, **kwargs)
+        return task.id
+
+    @staticmethod
+    def get_task_status(task_id: str) -> Dict[str, Union[str, Any, None]]:
+        result = AsyncResult(task_id)
+        return {
+            "task_id": task_id,
+            "status": result.status,
+            "result": result.result if result.ready() else None,
+            "traceback": result.traceback if result.failed() else None,
+        }
+
+    @staticmethod
+    def get_task_result(task_id: str) -> Optional[Any]:
+        result = AsyncResult(task_id)
+        if result.ready():
+            return result.result
+        return None
+
+    @staticmethod
+    def cancel_task(task_id: str) -> bool:
+        result = AsyncResult(task_id)
+        result.revoke(terminate=True)
+        return True