Add dataset seeding with Supabase storage and worker automation

- Seed script now creates datasets from ensure_datasets and uploads to Supabase storage
- Worker automatically scans for pending simulations every 60 seconds
- Downloads datasets from storage, runs simulations, marks complete
- Comprehensive logfire instrumentation in worker for step tracking
- Storage policies migration for datasets bucket access

Author: nikhilwoodruff

pyproject.toml

@@ -3,7 +3,7 @@ name = "policyengine-api-v2"
 version = "0.1.0"
 description = "FastAPI service for PolicyEngine microsimulations"
 readme = "README.md"
-requires-python = ">=3.11"
+requires-python = ">=3.13"
 dependencies = [
     "fastapi>=0.115.0",
     "uvicorn[standard]>=0.32.0",
@@ -13,7 +13,7 @@ dependencies = [
     "storage3>=0.8.1",
     "celery[redis]>=5.4.0",
     "redis>=5.1.1",
-    "policyengine>=0.1.0",
+    "policyengine>=3.1.5",
     "policyengine-uk>=2.0.0",
     "policyengine-us>=1.0.0",
     "pydantic>=2.9.2",
@@ -22,6 +22,7 @@ dependencies = [
     "tables>=3.10.2",
     "logfire[fastapi,httpx,sqlalchemy]>=0.60.0",
     "fastapi-cache2>=0.2.1",
+    "boto3>=1.41.1",
 ]
 
 [project.optional-dependencies]

scripts/seed.py

@@ -14,14 +14,18 @@
 sys.path.insert(0, str(Path(__file__).parent.parent / "src"))
 
 from policyengine.tax_benefit_models.uk import uk_latest
+from policyengine.tax_benefit_models.uk.datasets import ensure_datasets as ensure_uk_datasets
 from policyengine.tax_benefit_models.us import us_latest
+from policyengine.tax_benefit_models.us.datasets import ensure_datasets as ensure_us_datasets
 from policyengine_api.models import (
     TaxBenefitModel,
     TaxBenefitModelVersion,
     Variable,
     Parameter,
     ParameterValue,
+    Dataset,
 )
+from policyengine_api.services.storage import upload_dataset
 from sqlmodel import Session, create_engine, select
 from policyengine_api.config.settings import settings
 from rich.console import Console
@@ -140,6 +144,73 @@ def seed_model(model_version, session) -> TaxBenefitModelVersion:
     return db_version
 
 
+def seed_datasets(session):
+    """Seed datasets and upload to S3."""
+    console.print("[bold blue]Seeding datasets...")
+
+    # UK datasets
+    console.print("  Creating UK datasets...")
+    uk_datasets = ensure_uk_datasets()
+
+    for _, pe_dataset in track(list(uk_datasets.items()), description="UK datasets"):
+        # Check if dataset already exists
+        existing = session.exec(
+            select(Dataset).where(Dataset.name == pe_dataset.name)
+        ).first()
+
+        if existing:
+            console.print(f"  Dataset {pe_dataset.name} already exists, skipping")
+            continue
+
+        # Upload to S3
+        object_name = upload_dataset(pe_dataset.filepath)
+        console.print(f"  Uploaded {pe_dataset.filepath} to S3 as {object_name}")
+
+        # Create database record
+        db_dataset = Dataset(
+            name=pe_dataset.name,
+            description=pe_dataset.description,
+            filepath=object_name,  # Store S3 key, not local path
+            year=pe_dataset.year,
+            tax_benefit_model="uk_latest",
+        )
+        session.add(db_dataset)
+        session.commit()
+        console.print(f"  [green]✓[/green] Created dataset: {db_dataset.name}")
+
+    # US datasets
+    console.print("  Creating US datasets...")
+    us_datasets = ensure_us_datasets()
+
+    for _, pe_dataset in track(list(us_datasets.items()), description="US datasets"):
+        # Check if dataset already exists
+        existing = session.exec(
+            select(Dataset).where(Dataset.name == pe_dataset.name)
+        ).first()
+
+        if existing:
+            console.print(f"  Dataset {pe_dataset.name} already exists, skipping")
+            continue
+
+        # Upload to S3
+        object_name = upload_dataset(pe_dataset.filepath)
+        console.print(f"  Uploaded {pe_dataset.filepath} to S3 as {object_name}")
+
+        # Create database record
+        db_dataset = Dataset(
+            name=pe_dataset.name,
+            description=pe_dataset.description,
+            filepath=object_name,  # Store S3 key, not local path
+            year=pe_dataset.year,
+            tax_benefit_model="us_latest",
+        )
+        session.add(db_dataset)
+        session.commit()
+        console.print(f"  [green]✓[/green] Created dataset: {db_dataset.name}")
+
+    console.print(f"[green]✓[/green] Seeded {len(uk_datasets) + len(us_datasets)} datasets\n")
+
+
 def main():
     """Main seed function."""
     console.print("[bold green]PolicyEngine database seeding[/bold green]\n")
@@ -153,8 +224,10 @@ def main():
         us_version = seed_model(us_latest, session)
         console.print(f"[green]✓[/green] US model seeded: {us_version.id}\n")
 
+        # Seed datasets
+        seed_datasets(session)
+
     console.print("\n[bold green]✓ Database seeding complete![/bold green]")
-    console.print("\n[yellow]Note:[/yellow] Dataset creation skipped. To add datasets, upload H5 files via the API.")
 
 
 if __name__ == "__main__":

src/policyengine_api/services/storage.py

@@ -20,22 +20,22 @@ def upload_dataset(file_path: str, object_name: str | None = None) -> str:
         object_name: Name to store in bucket (defaults to filename)
 
     Returns:
-        Public URL of uploaded file
+        Object name (key) in storage
     """
     supabase = get_supabase_client()
 
     if object_name is None:
         object_name = Path(file_path).name
 
-    # Upload file
+    # Upload file using Supabase storage client
     with open(file_path, "rb") as f:
         supabase.storage.from_(settings.storage_bucket).upload(
-            object_name, f, {"content-type": "application/octet-stream"}
+            object_name,
+            f,
+            {"content-type": "application/octet-stream", "upsert": "true"}
         )
 
-    # Get public URL
-    url = supabase.storage.from_(settings.storage_bucket).get_public_url(object_name)
-    return url
+    return object_name
 
 
 def download_dataset(object_name: str, local_path: str) -> str:
@@ -50,7 +50,10 @@ def download_dataset(object_name: str, local_path: str) -> str:
     """
     supabase = get_supabase_client()
 
-    # Download file
+    # Ensure parent directory exists
+    Path(local_path).parent.mkdir(parents=True, exist_ok=True)
+
+    # Download file using Supabase storage client
     data = supabase.storage.from_(settings.storage_bucket).download(object_name)
 
     # Save locally

src/policyengine_api/tasks/celery_app.py

@@ -1,4 +1,5 @@
 from celery import Celery
+from celery.schedules import crontab
 import logfire
 
 from policyengine_api.config.settings import settings
@@ -23,4 +24,10 @@
     result_serializer="json",
     timezone="UTC",
     enable_utc=True,
+    beat_schedule={
+        "scan-pending-simulations": {
+            "task": "scan_pending_simulations",
+            "schedule": 60.0,  # Run every 60 seconds
+        },
+    },
 )