Merge pull request #19 from PolicyEngine/feat/modal-compute

feat: Move compute to Modal.com for sub-1s cold starts

Author: nikhilwoodruff

.env.example

@@ -20,5 +20,7 @@ DEBUG=true
 LOGFIRE_TOKEN=your-logfire-token-here
 LOGFIRE_ENVIRONMENT=local
 
-# Worker
-WORKER_POLL_INTERVAL=60
+# Modal.com (compute backend)
+# Authentication is via `modal token set` command, not env vars
+# Modal functions need a secret called "policyengine-db" with:
+#   DATABASE_URL, SUPABASE_URL, SUPABASE_KEY, STORAGE_BUCKET

Makefile

@@ -1,4 +1,4 @@
-.PHONY: install dev format lint test integration-test clean seed up down logs start-supabase stop-supabase reset rebuild create-state-bucket deploy-local init db-reset-prod
+.PHONY: install dev format lint test integration-test clean seed up down logs start-supabase stop-supabase reset rebuild create-state-bucket deploy-local init db-reset-prod modal-deploy modal-serve
 
 # AWS Configuration
 AWS_REGION ?= us-east-1
@@ -112,3 +112,17 @@ db-reset-prod:
 		echo "Aborted."; \
 		exit 1; \
 	fi
+
+modal-deploy:
+	@echo "Deploying Modal functions..."
+	@set -a && . .env.prod && set +a && \
+	uv run modal secret create policyengine-db \
+		"DATABASE_URL=$$SUPABASE_POOLER_URL" \
+		"SUPABASE_URL=$$SUPABASE_URL" \
+		"SUPABASE_KEY=$$SUPABASE_KEY" \
+		--force
+	uv run modal deploy src/policyengine_api/modal_app.py
+
+modal-serve:
+	@echo "Running Modal functions locally..."
+	uv run modal serve src/policyengine_api/modal_app.py

docker-compose.yml

@@ -23,31 +23,6 @@ services:
       retries: 10
       start_period: 10s
 
-  worker:
-    build: .
-    command: python -m policyengine_api.tasks.worker
-    ports:
-      - "${WORKER_PORT:-8080}:${WORKER_PORT:-8080}"
-    environment:
-      SUPABASE_URL: http://supabase_kong_policyengine-api-v2:8000
-      SUPABASE_KEY: ${SUPABASE_KEY}
-      SUPABASE_SERVICE_KEY: ${SUPABASE_SERVICE_KEY}
-      SUPABASE_DB_URL: postgresql://postgres:postgres@supabase_db_policyengine-api-v2:5432/postgres
-      LOGFIRE_TOKEN: ${LOGFIRE_TOKEN}
-      WORKER_POLL_INTERVAL: ${WORKER_POLL_INTERVAL:-5}
-      WORKER_PORT: ${WORKER_PORT:-8080}
-    volumes:
-      - ./src:/app/src
-      - dataset_cache:/tmp/policyengine_dataset_cache
-    networks:
-      - supabase_network_policyengine-api-v2
-    healthcheck:
-      test: ["CMD", "python", "-c", "import httpx; exit(0 if httpx.get('http://localhost:${WORKER_PORT:-8080}/health', timeout=2).status_code == 200 else 1)"]
-      interval: 10s
-      timeout: 5s
-      retries: 3
-      start_period: 30s
-
   test:
     build: .
     command: pytest tests/ -v
@@ -65,14 +40,9 @@ services:
     depends_on:
       api:
         condition: service_healthy
-      worker:
-        condition: service_started
     profiles:
       - test
 
 networks:
   supabase_network_policyengine-api-v2:
     external: true
-
-volumes:
-  dataset_cache:

pyproject.toml

@@ -22,6 +22,7 @@ dependencies = [
     "fastapi-cache2>=0.2.1",
     "boto3>=1.41.1",
     "fastapi-mcp>=0.4.0",
+    "modal>=0.68.0",
 ]
 
 [project.optional-dependencies]

src/policyengine_api/api/analysis.py

@@ -19,6 +19,7 @@
 from typing import Literal
 from uuid import UUID, uuid5
 
+import logfire
 from fastapi import APIRouter, Depends, HTTPException
 from pydantic import BaseModel, Field
 from sqlmodel import Session, select
@@ -285,6 +286,18 @@ def _build_response(
     )
 
 
+def _trigger_modal_report(report_id: str, tax_benefit_model_name: str) -> None:
+    """Trigger Modal function for economic impact report."""
+    import modal
+
+    if tax_benefit_model_name == "policyengine_uk":
+        fn = modal.Function.from_name("policyengine", "run_report_uk")
+    else:
+        fn = modal.Function.from_name("policyengine", "run_report_us")
+
+    fn.spawn(report_id=report_id)
+
+
 @router.post("/economic-impact", response_model=EconomicImpactResponse)
 def economic_impact(
     request: EconomicImpactRequest,
@@ -333,6 +346,11 @@ def economic_impact(
 
     report = _get_or_create_report(baseline_sim.id, reform_sim.id, label, session)
 
+    # Trigger Modal if report is pending
+    if report.status == ReportStatus.PENDING:
+        with logfire.span("trigger_modal_report", report_id=str(report.id)):
+            _trigger_modal_report(str(report.id), request.tax_benefit_model_name)
+
     return _build_response(report, baseline_sim, reform_sim, session)