Prefect introduction

README.md

@@ -9,7 +9,7 @@ Both modes require the frontend to be running on the host. See [Frontend](#front
 
 ### Containerized mode (supported)
 
-Runs the full backend stack in Docker (Postgres, Redis, Celery worker, and the Python API).
+Runs the full backend stack in Docker (Postgres, Redis, Prefect worker, and the Python API).
 This is the supported way to verify the app works end to end.
 
 **Requirements:** [Docker](https://docs.docker.com/get-docker/) with Compose
@@ -29,7 +29,7 @@ and faster iteration during development.
 - Python ≥ 3.13 + [uv](https://docs.astral.sh/uv/)
 - [Docker](https://docs.docker.com/get-docker/) with Compose
 
-**1. Start Postgres, Redis, and the Celery worker**
+**1. Start Postgres, Redis, and the Prefect worker**
 
 ```bash
 docker compose -f compose/compose.services.yaml --env-file compose/.env.local up

compose/.env.local

@@ -59,7 +59,14 @@ SNAKEDISPATCH_SYNC_INTERVAL=10.0
 CALLBACK_URL_ALLOWED_DOMAINS=
 
 # Comma-separated list of Snakedispatch backends in name=url format (e.g. cluster-a=http://sd-a:8000,cluster-b=http://sd-b:8000)
-# SNAKEDISPATCH_BACKENDS=
+SNAKEDISPATCH_BACKENDS=local=http://host.docker.internal:8010
+
+# Prefect
+# -------
+
+# Prefect server API URL for background task orchestration
+PREFECT_API_URL=http://localhost:4200/api
+PREFECT_LOCAL_STORAGE_PATH=${DATA_DIR}/prefect-results
 
 # Redis
 # -----

compose/compose.services.yaml

@@ -1,4 +1,4 @@
-# Backend services for development: Postgres + Redis + Celery worker in Docker.
+# Backend services for development: Postgres + Redis + Prefect worker in Docker.
 # Values come from compose/.env.local. Run with:
 #   docker compose -f compose/compose.services.yaml --env-file compose/.env.local up
 
@@ -17,6 +17,7 @@ services:
     - "${POSTGRES_PORT:-5432}:5432"
     volumes:
     - postgres_data:/var/lib/postgresql
+    - ./postgres-init:/docker-entrypoint-initdb.d:ro
     healthcheck:
       test: ["CMD-SHELL", "pg_isready -U ${POSTGRES_USER} -d ${POSTGRES_DB}"]
       interval: 5s
@@ -36,20 +37,43 @@ services:
       timeout: 3s
       retries: 5
 
-  celery-worker:
+  prefect-server:
+    image: prefecthq/prefect:3-latest
+    container_name: pypsa-app-prefect-server
+    command: prefect server start --host 0.0.0.0
+    environment:
+      PREFECT_API_DATABASE_CONNECTION_URL: postgresql+asyncpg://${POSTGRES_USER}:${POSTGRES_PASSWORD}@postgres:5432/prefect
+    ports:
+    - "4200:4200"
+    depends_on:
+      postgres:
+        condition: service_healthy
+    healthcheck:
+      test: ["CMD", "python", "-c", "import urllib.request; urllib.request.urlopen('http://localhost:4200/api/health')"]
+      interval: 5s
+      timeout: 5s
+      retries: 10
+      start_period: 30s
+
+  prefect-worker:
     build:
       context: ..
       dockerfile: Dockerfile
       target: backend
-    container_name: pypsa-app-celery-worker
+    container_name: pypsa-app-prefect-worker
     env_file: .env.local
     environment:
       DATABASE_URL: postgresql://${POSTGRES_USER}:${POSTGRES_PASSWORD}@postgres:5432/${POSTGRES_DB}
       REDIS_URL: redis://redis:6379/0
       DATA_DIR: /data
+      PREFECT_API_URL: http://prefect-server:4200/api
+      PREFECT_WORKER_CONCURRENCY: "8"
+      PREFECT_LOCAL_STORAGE_PATH: /data/prefect-results
       # The host-run backend owns migrations on startup; avoid racing it here.
       SKIP_MIGRATIONS: "true"
-    command: ["celery", "-A", "pypsa_app.backend.task_queue.task_app", "worker", "--loglevel=info", "--concurrency=2"]
+    command: ["python", "-m", "pypsa_app.worker"]
+    extra_hosts:
+    - "host.docker.internal:host-gateway"
     volumes:
     - ../src:/app/src
     - ../data:/data
@@ -59,6 +83,8 @@ services:
         condition: service_healthy
       redis:
         condition: service_healthy
+      prefect-server:
+        condition: service_healthy
 
 volumes:
   postgres_data:

compose/compose.yaml

@@ -1,4 +1,4 @@
-# Containerized mode: runs the full backend stack in Docker (postgres/redis/celery-worker + API).
+# Containerized mode: runs the full backend stack in Docker (postgres/redis/prefect + API).
 # This is the supported way to verify the app works end to end. Run with:
 #   docker compose -f compose/compose.yaml --env-file compose/.env.local up --build
 
@@ -22,15 +22,21 @@ services:
       REDIS_URL: redis://redis:6379/0
       DATA_DIR: /data
       BASE_URL: http://localhost:${APP_PORT:-8000}
+      PREFECT_API_URL: http://prefect-server:4200/api
+      PREFECT_LOCAL_STORAGE_PATH: /data/prefect-results
     ports:
     - "${APP_PORT:-8000}:8000"
+    extra_hosts:
+    - "host.docker.internal:host-gateway"
     volumes:
     - ../data:/data
     depends_on:
       postgres:
         condition: service_healthy
       redis:
         condition: service_healthy
+      prefect-server:
+        condition: service_healthy
     healthcheck:
       test: ["CMD", "curl", "-f", "http://localhost:8000/health"]
       interval: 10s

compose/postgres-init/01-create-prefect-db.sql

@@ -0,0 +1 @@
+SELECT 'CREATE DATABASE prefect' WHERE NOT EXISTS (SELECT FROM pg_database WHERE datname = 'prefect')\gexec

pyproject.toml

@@ -42,7 +42,7 @@ dependencies = [
 full = [
     "psycopg2-binary>=2.9",
     "redis>=6.4,<7",
-    "celery[redis]>=5.6,<6",
+    "prefect>=3.0,<4",
 ]
 
 dev = [

src/pypsa_app/backend/api/routes/networks_remote.py

@@ -1,5 +1,6 @@
 """Network import routes (upload, from-url)."""
 
+import asyncio
 from collections.abc import Iterator
 from pathlib import PurePosixPath
 
@@ -11,13 +12,13 @@
     stream_to_file_with_limit,
 )
 from pypsa_app.backend.api.utils.task_utils import queue_task
+from pypsa_app.backend.flows import (
+    import_network_from_file_flow,
+    import_network_from_url_flow,
+)
 from pypsa_app.backend.models import Permission, User
 from pypsa_app.backend.schemas.task import TaskQueuedResponse
 from pypsa_app.backend.settings import settings
-from pypsa_app.backend.tasks import (
-    import_network_from_file_task,
-    import_network_from_url_task,
-)
 from pypsa_app.backend.utils.validation import ExternalHttpUrl
 
 router = APIRouter()
@@ -29,7 +30,7 @@ def _iter_upload_file(file: UploadFile, chunk_size: int = 8192) -> Iterator[byte
 
 
 @router.post("/", response_model=TaskQueuedResponse, status_code=202)
-def upload_network(
+async def upload_network(
     file: UploadFile,
     user: User = Depends(require_permission(Permission.NETWORKS_MODIFY)),
 ) -> dict:
@@ -45,24 +46,30 @@ def upload_network(
     max_bytes = settings.max_upload_size_mb * 1024 * 1024
     tmp = make_temp_path(user.id)
     try:
-        stream_to_file_with_limit(_iter_upload_file(file), tmp, max_bytes)
+        await asyncio.to_thread(
+            stream_to_file_with_limit, _iter_upload_file(file), tmp, max_bytes
+        )
     except ValueError as exc:
         tmp.unlink(missing_ok=True)
         raise HTTPException(413, str(exc)) from exc
     except Exception:
         tmp.unlink(missing_ok=True)
         raise
 
-    return queue_task(
-        import_network_from_file_task,
-        tmp_path=str(tmp),
-        original_filename=safe_filename,
-        user_id=str(user.id),
-    )
+    try:
+        return await queue_task(
+            import_network_from_file_flow,
+            tmp_path=str(tmp),
+            original_filename=safe_filename,
+            user_id=str(user.id),
+        )
+    except Exception:
+        tmp.unlink(missing_ok=True)
+        raise
 
 
 @router.post("/from-url", response_model=TaskQueuedResponse, status_code=202)
-def create_network_from_url(
+async def create_network_from_url(
     url: ExternalHttpUrl = Body(..., embed=True),
     user: User = Depends(require_permission(Permission.NETWORKS_MODIFY)),
 ) -> dict:
@@ -73,8 +80,8 @@ def create_network_from_url(
     if not url_str.lower().endswith(".nc"):
         raise HTTPException(400, "URL must point to a .nc file")
 
-    return queue_task(
-        import_network_from_url_task,
+    return await queue_task(
+        import_network_from_url_flow,
         url=url_str,
         user_id=str(user.id),
     )

src/pypsa_app/backend/api/routes/plots.py

@@ -5,20 +5,20 @@
 
 from pypsa_app.backend.api.deps import get_db, get_networks, require_permission
 from pypsa_app.backend.api.utils.task_utils import queue_task
+from pypsa_app.backend.flows import get_explore_flow, get_plot_flow
 from pypsa_app.backend.models import Permission, User
 from pypsa_app.backend.ratelimit import limiter
 from pypsa_app.backend.schemas.plot import ExploreRequest, PlotParams
 from pypsa_app.backend.schemas.task import TaskQueuedResponse
 from pypsa_app.backend.settings import settings
-from pypsa_app.backend.tasks import get_explore_task, get_plot_task
 
 router = APIRouter()
 logger = logging.getLogger(__name__)
 
 
 @router.post("/generate", response_model=TaskQueuedResponse)
 @limiter.limit(settings.ratelimit_expensive)
-def generate_plot(
+async def generate_plot(
     request: Request,
     response: Response,
     body: PlotParams,
@@ -29,8 +29,8 @@ def generate_plot(
     networks = get_networks(db, body.network_ids, user)
     file_paths = [net.file_path for net in networks]
 
-    return queue_task(
-        get_plot_task,
+    return await queue_task(
+        get_plot_flow,
         file_paths=file_paths,
         statistic=body.statistic,
         plot_type=body.plot_type,
@@ -40,7 +40,7 @@ def generate_plot(
 
 @router.post("/explore", response_model=TaskQueuedResponse)
 @limiter.limit(settings.ratelimit_expensive)
-def generate_explore(
+async def generate_explore(
     request: Request,
     response: Response,
     body: ExploreRequest,
@@ -57,8 +57,8 @@ def generate_explore(
     if body.geometry:
         parameters["geometry"] = True
 
-    return queue_task(
-        get_explore_task,
+    return await queue_task(
+        get_explore_flow,
         file_paths=file_paths,
         parameters=parameters,
     )

src/pypsa_app/backend/api/routes/runs.py

@@ -2,13 +2,12 @@
 
 import logging
 import re
-import threading
 import urllib.parse
 import uuid
 from pathlib import PurePosixPath
 from typing import Any, Literal
 
-from fastapi import APIRouter, Depends, HTTPException, Query
+from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException, Query
 from fastapi import Path as PathParam
 from fastapi.responses import StreamingResponse
 from sqlalchemy import or_, select
@@ -25,6 +24,7 @@
     apply_pagination,
     list_meta,
 )
+from pypsa_app.backend.api.utils.task_utils import submit_flow
 from pypsa_app.backend.cache import cache
 from pypsa_app.backend.filters import (
     FieldMap,
@@ -33,6 +33,7 @@
     enum_coercer,
     name_to_id,
 )
+from pypsa_app.backend.flows import import_run_outputs_flow
 from pypsa_app.backend.models import (
     Permission,
     Run,
@@ -272,6 +273,7 @@ def _distinct_vals(col: Any) -> list:
 
 @router.get("/{run_id}", response_model=RunResponse)
 def get_run(
+    background_tasks: BackgroundTasks,
     auth: Authorized[Run] = Depends(require_run("read")),
     db: Session = Depends(get_db),
 ) -> RunResponse:
@@ -283,18 +285,16 @@ def get_run(
         if client:
             try:
                 job = client.get_job(str(run.job_id))
-                needs_callback = sync_run_from_job(run, job, db)
+                needs_callback, needs_import = sync_run_from_job(run, job, db)
                 db.commit()
                 if needs_callback and run.callback_url:
-                    # TODO: replace with proper async callback or
-                    # FastAPI BackgroundTasks.
-                    url = str(run.callback_url)
-                    payload = _build_payload(run)
-                    threading.Thread(
-                        target=post_callback_sync,
-                        args=(url, payload),
-                        daemon=True,
-                    ).start()
+                    background_tasks.add_task(
+                        post_callback_sync, str(run.callback_url), _build_payload(run)
+                    )
+                if needs_import:
+                    background_tasks.add_task(
+                        submit_flow, import_run_outputs_flow, job_id=str(run.job_id)
+                    )
             except SnakedispatchError:
                 pass
 
@@ -421,6 +421,7 @@ def download_run_output(
 
 @router.post("/{run_id}/cancel", response_model=MessageResponse)
 def cancel_run(
+    background_tasks: BackgroundTasks,
     auth: Authorized[Run] = Depends(require_run("modify")),
     db: Session = Depends(get_db),
 ) -> dict:
@@ -430,31 +431,25 @@ def cancel_run(
     sd_client = _get_client_for_run(run)
     try:
         result = sd_client.cancel_job(str(run.job_id))
-        needs_callback = sync_run_from_job(run, result, db)
+        needs_callback, needs_import = sync_run_from_job(run, result, db)
         db.commit()
         if needs_callback and run.callback_url:
-            # TODO: replace with proper async callback or
-            # FastAPI BackgroundTasks.
-            url = str(run.callback_url)
-            payload = _build_payload(run)
-            threading.Thread(
-                target=post_callback_sync, args=(url, payload), daemon=True
-            ).start()
+            background_tasks.add_task(
+                post_callback_sync, str(run.callback_url), _build_payload(run)
+            )
+        if needs_import:
+            background_tasks.add_task(
+                submit_flow, import_run_outputs_flow, job_id=str(run.job_id)
+            )
     except SnakedispatchError as e:
         if e.status_code in (404, 409):
             if run.status not in SYNCED_STATUSES:
                 run.status = RunStatus.CANCELLED
                 db.commit()
                 if run.callback_url:
-                    # TODO: replace with proper async callback or
-                    # FastAPI BackgroundTasks.
-                    url = str(run.callback_url)
-                    payload = _build_payload(run)
-                    threading.Thread(
-                        target=post_callback_sync,
-                        args=(url, payload),
-                        daemon=True,
-                    ).start()
+                    background_tasks.add_task(
+                        post_callback_sync, str(run.callback_url), _build_payload(run)
+                    )
         else:
             raise