feat: improved in-memory task queue as celery alternative

Author: lkstrp

src/pypsa_app/backend/api/utils/task_utils.py

@@ -1,17 +1,27 @@
-"""Shared utilities for Celery task status responses"""
+"""Shared utilities for task status responses"""
 
 import logging
-from celery.result import AsyncResult
 
-from pypsa_app.backend.celery_app import celery_app
+from pypsa_app.backend.task_queue import task_app
 from pypsa_app.backend.settings import settings
 
 logger = logging.getLogger(__name__)
 
 
 def get_task_status_response(task_id: str) -> dict:
-    """Get standardized Celery task status response"""
-    task = AsyncResult(task_id, app=celery_app)
+    """Get standardized task status response"""
+    # Check if using Celery or in-memory queue
+    try:
+        # Try using Celery's AsyncResult if Celery is the backend
+        from celery.result import AsyncResult
+
+        task = AsyncResult(task_id, app=task_app)
+    except (ImportError, AttributeError):
+        # Fall back to in-memory AsyncResult
+        from pypsa_app.backend.task_queue import InMemoryAsyncResult
+
+        task = InMemoryAsyncResult(task_id)
+
     response = {"task_id": task_id, "state": task.state}
 
     match task.state:

src/pypsa_app/backend/celery_app.py

@@ -0,0 +1,140 @@
+"""Celery configuration and in-memory task queue fallback"""
+
+import logging
+import threading
+import uuid
+from concurrent.futures import ThreadPoolExecutor
+from datetime import datetime, timedelta, timezone
+
+from pypsa_app.backend.settings import settings
+
+logger = logging.getLogger(__name__)
+
+_tasks = {}
+_lock = threading.Lock()
+_pool = ThreadPoolExecutor(max_workers=1, thread_name_prefix="task")
+
+
+class InMemoryAsyncResult:
+    def __init__(self, task_id):
+        self.id = task_id
+
+    @property
+    def state(self):
+        with _lock:
+            return _tasks.get(self.id, {}).get("state", "PENDING")
+
+    @property
+    def result(self):
+        with _lock:
+            t = _tasks.get(self.id, {})
+            return t.get("result") if t.get("state") == "SUCCESS" else None
+
+    @property
+    def info(self):
+        with _lock:
+            t = _tasks.get(self.id, {})
+            if t.get("state") == "FAILURE":
+                return t.get("exception")
+            if t.get("state") == "PROGRESS":
+                return t.get("meta", {})
+            return None
+
+
+class InMemoryTaskQueue:
+    def task(self, *args, **kwargs):
+        bind = kwargs.get("bind", False)
+
+        def decorator(func):
+            def apply_async(args=(), kwargs=None, **options):
+                tid = str(uuid.uuid4())
+                now = datetime.now(timezone.utc)
+
+                with _lock:
+                    _tasks[tid] = {"state": "PENDING", "created_at": now}
+                    cutoff = now - timedelta(hours=24)
+                    for k in list(_tasks.keys()):
+                        if _tasks[k].get("created_at", cutoff) < cutoff:
+                            del _tasks[k]
+
+                class Task:
+                    request = type("Request", (), {"id": tid})()
+
+                    @staticmethod
+                    def update_state(state=None, meta=None):
+                        with _lock:
+                            if tid in _tasks and state:
+                                _tasks[tid]["state"] = state
+                            if tid in _tasks and meta:
+                                _tasks[tid]["meta"] = meta
+
+                def run():
+                    try:
+                        res = (
+                            func(Task(), *args, **(kwargs or {}))
+                            if bind
+                            else func(*args, **(kwargs or {}))
+                        )
+                        with _lock:
+                            _tasks[tid].update({"state": "SUCCESS", "result": res})
+                    except Exception as e:
+                        with _lock:
+                            _tasks[tid].update(
+                                {"state": "FAILURE", "exception": str(e)}
+                            )
+                        logger.error(
+                            "Task failed",
+                            extra={"task_id": tid, "error": str(e)},
+                            exc_info=True,
+                        )
+
+                _pool.submit(run)
+                return InMemoryAsyncResult(tid)
+
+            func.apply_async = apply_async
+            func.name = kwargs.get("name", func.__name__)
+            return func
+
+        return decorator
+
+
+# Try to use Celery with Redis, fall back to in-memory task queue
+try:
+    from celery import Celery
+
+    # Only use real Celery if Redis URL is configured
+    if not settings.redis_url:
+        logger.warning(
+            "Redis URL not configured - using in-memory task queue",
+            extra={"backend": "in-memory", "background_tasks_enabled": True},
+        )
+        task_app = InMemoryTaskQueue()
+    else:
+        task_app = Celery(
+            "pypsa_app",
+            broker=settings.redis_url,
+            backend=settings.redis_url,
+            include=["pypsa_app.backend.tasks"],
+        )
+
+        task_app.conf.update(
+            accept_content=["json"],
+            result_expires=86400,
+            worker_prefetch_multiplier=1,
+            worker_max_tasks_per_child=10,
+            task_soft_time_limit=3600,
+            task_time_limit=7200,
+            task_acks_late=True,
+        )
+
+        logger.info(
+            "Initialized Celery with Redis backend",
+            extra={"redis_url": settings.redis_url, "backend": "celery"},
+        )
+
+except ImportError:
+    logger.warning(
+        "Celery not installed - using in-memory task queue",
+        extra={"backend": "in-memory", "background_tasks_enabled": True},
+    )
+    task_app = InMemoryTaskQueue()

src/pypsa_app/backend/task_queue.py

@@ -5,7 +5,7 @@
 from typing import Any, Callable, Dict
 
 from pypsa_app.backend.cache import cache
-from pypsa_app.backend.celery_app import celery_app
+from pypsa_app.backend.task_queue import task_app
 from pypsa_app.backend.schemas.task import TaskResult
 from pypsa_app.backend.services.map import extract_geographic_layer
 from pypsa_app.backend.services.network import scan_networks
@@ -46,28 +46,28 @@ def _execute_task(self, name: str, func: Callable, **kwargs) -> Dict[str, Any]:
         ).model_dump()
 
 
-@celery_app.task(bind=True, name="tasks.get_statistics")
+@task_app.task(bind=True, name="tasks.get_statistics")
 def get_statistics_task(self, **kwargs):
-    """Celery task for statistics generation"""
+    """Background task for statistics generation"""
     func = cache("statistics", ttl=settings.plot_cache_ttl)(get_statistics_service)
     return _execute_task(self, "Statistics generation", func, **kwargs)
 
 
-@celery_app.task(bind=True, name="tasks.get_plot")
+@task_app.task(bind=True, name="tasks.get_plot")
 def get_plot_task(self, **kwargs):
-    """Celery task for plot generation"""
+    """Background task for plot generation"""
     func = cache("plot", ttl=settings.plot_cache_ttl)(get_plot_service)
     return _execute_task(self, "Plot generation", func, **kwargs)
 
 
-@celery_app.task(bind=True, name="tasks.extract_geographic_layer")
+@task_app.task(bind=True, name="tasks.extract_geographic_layer")
 def extract_geographic_layer_task(self, **kwargs):
-    """Celery task for geographic layer extraction"""
+    """Background task for geographic layer extraction"""
     func = cache("map_data", ttl=settings.map_cache_ttl)(extract_geographic_layer)
     return _execute_task(self, "Geographic layer extraction", func, **kwargs)
 
 
-@celery_app.task(bind=True, name="tasks.scan_networks")
+@task_app.task(bind=True, name="tasks.scan_networks")
 def scan_networks_task(self, **kwargs):
-    """Celery task for network scanning (no caching)"""
+    """Background task for network scanning (no caching)"""
     return _execute_task(self, "Network scan", scan_networks, **kwargs)