Merge pull request #14 from alphagov/ACW-37/non-blocking

ACW-37 Implement tracking for running and background jobs with resilience

Author: daymohlar

app.py

@@ -11,7 +11,14 @@
 from flask import Flask, request, jsonify, render_template
 from src.visualiser_graph_generator import generate_graph, generate_output_path
 from src.visualiser_graph_loader import load_json_file, extract_path_parts, visualiser_graph_file_path
-from src.utils import update_job_status, read_job_status
+from src.utils import (
+    update_job_status, 
+    read_job_status, 
+    get_job_id_for_path,
+    get_active_job_status,
+    background_run_extraction,
+    resume_interrupted_jobs
+)
 from werkzeug.exceptions import BadRequest
 
 
@@ -73,8 +80,18 @@ async def extract_quotes():
                 return jsonify({"error": "Missing 'source_path' query parameter"}), 400
 
             input_path, output_path = generate_output_path(source_path)
-            job_id = str(uuid.uuid4())
+            job_id = get_job_id_for_path(source_path)
 
+            active_status = get_active_job_status(job_id)
+            if active_status:
+                logger.info(f"Duplicate request for {source_path}. Job {job_id} is already in progress.")
+                return jsonify({
+                    'job_id': job_id,
+                    'status': 'already_running',
+                    'message': f'A graph generation job is already in progress for {source_path}',
+                    'output_path': output_path
+                }), 202
+
             initial_status = {
                 "job_id": job_id,
                 "status": "pending",
@@ -83,26 +100,7 @@ async def extract_quotes():
             }
             update_job_status(job_id, initial_status)
 
-            async def run_extraction():
-                try:
-                    logger.info(f'Starting background graph generation for {input_path} (Job: {job_id})...')
-                    initial_status["status"] = "running"
-                    update_job_status(job_id, initial_status)
-                    
-                    await generate_graph(input_path, output_path)
-                    
-                    initial_status["status"] = "completed"
-                    initial_status["output_path"] = output_path
-                    initial_status["completed_at"] = time.time()
-                    update_job_status(job_id, initial_status)
-                    logger.info(f'Graph generation completed successfully for {output_path}')
-                except Exception as e:
-                    logger.error(f"Background graph generation failed for job {job_id}: {str(e)}")
-                    initial_status["status"] = "failed"
-                    initial_status["error"] = str(e)
-                    update_job_status(job_id, initial_status)
-
-            asyncio.create_task(run_extraction())
+            asyncio.create_task(background_run_extraction(job_id, input_path, output_path, initial_status))
 
             return jsonify({
                 'job_id': job_id,
@@ -129,8 +127,29 @@ def handle_bad_request(e):
 
     return app
 
+class LifespanMiddleware:
+    """ASGI middleware to handle startup and shutdown events."""
+    def __init__(self, app):
+        self.app = app
+
+    async def __call__(self, scope, receive, send):
+        if scope["type"] == "lifespan":
+            while True:
+                message = await receive()
+                if message["type"] == "lifespan.startup":
+                    # Trigger resumption when the event loop is officially running
+                    logger.info("ASGI startup: triggering job resumption scan...")
+                    asyncio.create_task(resume_interrupted_jobs())
+                    await send({"type": "lifespan.startup.complete"})
+                elif message["type"] == "lifespan.shutdown":
+                    await send({"type": "lifespan.shutdown.complete"})
+                    return
+        return await self.app(scope, receive, send)
+
 def create_asgi_app():
-    return WsgiToAsgi(create_app())
+    flask_app = create_app()
+    asgi_app = WsgiToAsgi(flask_app)
+    return LifespanMiddleware(asgi_app)
 
 if __name__ == "__main__":
     asgi_app = create_asgi_app()

pyproject.toml

@@ -14,7 +14,6 @@ dependencies = [
     "opensearch-py>=2.8.0",
     "requests-aws4auth>=1.3.1",
     "flask[async]==3.1.2",
-    "waitress>=3.0.0",
     "uvicorn>=0.30.0",
 ]
 

src/utils/__init__.py

@@ -1 +1,9 @@
-from .job_tracker import update_job_status, read_job_status, get_status_path
+from .job_tracker import (
+    update_job_status, 
+    read_job_status, 
+    get_status_path,
+    get_job_id_for_path,
+    get_active_job_status,
+    background_run_extraction,
+    resume_interrupted_jobs
+)

src/utils/job_tracker.py

@@ -1,6 +1,11 @@
 import json
 import logging
 import fsspec
+import hashlib
+import time
+import asyncio
+from typing import Optional
+from src.visualiser_graph_generator import generate_graph, generate_output_path
 
 logger = logging.getLogger(__name__)
 
@@ -21,7 +26,7 @@ def update_job_status(job_id: str, status_data: dict):
     except Exception as e:
         logger.error(f"Failed to update job status in S3 for job {job_id}: {str(e)}")
 
-def read_job_status(job_id: str) -> dict:
+def read_job_status(job_id: str) -> dict| None:
     """Reads the job status from S3."""
     path = get_status_path(job_id)
     fs = fsspec.filesystem("s3")
@@ -37,3 +42,76 @@ def read_job_status(job_id: str) -> dict:
     except Exception as e:
         logger.error(f"Failed to read job status from S3 for job {job_id}: {str(e)}")
         return None
+
+def get_job_id_for_path(source_path: str) -> str:
+    """Generates a predictable job_id based on a hash of the source_path."""
+    return hashlib.sha256(source_path.encode()).hexdigest()
+
+def get_active_job_status(job_id: str, timeout_hours: int = 24) -> Optional[dict]:
+    """
+    Returns the job status ONLY if it is currently active (pending/running) 
+    and has not exceeded the timeout. Returns None if the job is stale or not active.
+    """
+    status = read_job_status(job_id)
+    if not status:
+        return None
+        
+    if status.get("status") in ["pending", "running"]:
+        created_at = status.get("created_at", 0)
+        is_stale = (time.time() - created_at) > (timeout_hours * 3600)
+        
+        if not is_stale:
+            return status
+            
+    return None
+
+async def background_run_extraction(job_id: str, input_path: str, output_path: str, status: dict):
+    """Background task for graph generation and status tracking."""
+    try:
+        logger.info(f'Starting background graph generation for {input_path} (Job: {job_id})...')
+        status["status"] = "running"
+        update_job_status(job_id, status)
+        
+        await generate_graph(input_path, output_path)
+        
+        status["status"] = "completed"
+        status["output_path"] = output_path
+        status["completed_at"] = time.time()
+        update_job_status(job_id, status)
+        logger.info(f'Graph generation completed successfully for {output_path}')
+    except Exception as e:
+        logger.error(f"Background graph generation failed for job {job_id}: {str(e)}")
+        status["status"] = "failed"
+        status["error"] = str(e)
+        update_job_status(job_id, status)
+
+async def resume_interrupted_jobs():
+    """Scans for jobs stuck in 'running' state and restarts them if they are fresh (<24h)."""
+    logger.info("Scanning for interrupted jobs to resume...")
+    fs = fsspec.filesystem("s3")
+    
+    try:
+        # List all status files
+        status_files = fs.glob(f"{STATUS_ROOT}/*.json")
+        
+        for file_path in status_files:
+            job_id = file_path.split("/")[-1].replace(".json", "")
+            status = read_job_status(job_id)
+            
+            if status and status.get("status") == "running":
+                created_at = status.get("created_at", 0)
+                is_fresh = (time.time() - created_at) < (24 * 3600)
+                
+                if is_fresh:
+                    source_path = status.get("source_path")
+                    if source_path:
+                        try:
+                            input_path, output_path = generate_output_path(source_path)
+                            logger.info(f"Resuming interrupted job {job_id} for {source_path}")
+                            asyncio.create_task(background_run_extraction(job_id, input_path, output_path, status))
+                        except Exception as e:
+                            logger.error(f"Failed to prepare resumption for job {job_id}: {str(e)}")
+                else:
+                    logger.info(f"Skipping stale interrupted job {job_id} (over 24h old)")
+    except Exception as e:
+        logger.error(f"Error during job resumption scan: {str(e)}")