Merge pull request #3387 from PolicyEngine/codex/budget-window-batch

Add budget-window batch economy endpoint

Author: anth-volk

.env.example

@@ -18,3 +18,9 @@ OPENAI_API_KEY=policyengine_openai_api_key
 
 # Token for Hugging Face models 
 HUGGING_FACE_TOKEN=policyengine_huggingface_token
+
+# Redis is required for budget-window economy requests and other API cache paths.
+# Local development and App Engine use an in-container/local Redis by default.
+CACHE_REDIS_HOST=127.0.0.1
+CACHE_REDIS_PORT=6379
+CACHE_REDIS_DB=0

.github/workflows/push.yml

@@ -226,7 +226,7 @@ jobs:
       - name: Install staging test dependencies
         run: pip install pytest httpx
       - name: Run staging smoke test
-        run: python -m pytest tests/integration/test_live_calculate.py tests/integration/test_live_economy.py -v
+        run: python -m pytest tests/integration/test_live_calculate.py tests/integration/test_live_economy.py tests/integration/test_live_budget_window_cache.py -v
         env:
           API_BASE_URL: ${{ needs.deploy-staging.outputs.url }}
           STAGING_API_TEST_PROBE_ID: ${{ needs.deploy-staging.outputs.version }}

Makefile

@@ -11,7 +11,7 @@ test-env-vars:
 	pytest tests/env_variables
 
 test:
-	MAX_HOUSEHOLDS=1000 coverage run -a --branch -m pytest tests/to_refactor tests/unit --disable-pytest-warnings
+	MAX_HOUSEHOLDS=1000 coverage run -a --branch -m pytest tests/to_refactor tests/unit tests/integration/test_budget_window_in_flight_dedupe.py --disable-pytest-warnings
 	coverage xml -i
 
 debug-test:

README.md

@@ -120,6 +120,8 @@ NOTE: Any output that needs to be calculated will not work. Therefore, only hous
 
 ### 6. Testing calculations
 
+Redis is required for API cache paths, including budget-window economy requests. The budget-window endpoint uses Redis for completed-result caching and in-flight batch deduplication; if Redis is unavailable, those requests fail instead of falling back to the database or an in-process cache.
+
 To test anything that utilizes Redis or the API's service workers (e.g. anything that requires society-wide calculations with the policy calculator), you'll also need to complete the following steps:
 
 1. Start Redis
@@ -136,6 +138,8 @@ brew install redis
 redis-server
 ```
 
+By default the API connects to Redis at `127.0.0.1:6379`, database `0`. Override this with `CACHE_REDIS_HOST`, `CACHE_REDIS_PORT`, and `CACHE_REDIS_DB` if your local Redis uses different connection settings.
+
 2. Start the API
 
 Run the below
@@ -144,6 +148,8 @@ Run the below
 FLASK_DEBUG=1 python -m flask --app policyengine_api.api run
 ```
 
+App Engine staging and production deployments install and start Redis in the API container before Gunicorn starts.
+
 NOTE: Calculations are not possible in the uk app without access to a specific dataset. Expect an error: "ValueError: Invalid response code 404 for url https://api.github.com/repos/policyengine/non-public-microdata/releases/tags/uk-2024-march-efo."
 
 ## Testing, Formatting, Changelogging

changelog.d/budget-window-batch.fixed.md

@@ -0,0 +1 @@
+Added a budget-window economy endpoint that batches yearly impact calculations with bounded server-side concurrency and returns aggregated progress plus totals.

gcp/README.md

@@ -2,6 +2,8 @@
 
 The deployment actions build Docker images and deploy them to Google App Engine. The docker images themselves are based off a starter image (to save each API docker image having to spend 5 minutes installing the same dependencies). The starter image is the `Dockerfile` in this directory.
 
+The App Engine API image installs `redis-server` and starts it through `gcp/policyengine_api/start.sh`. Redis is required at runtime for budget-window economy request caching and in-flight batch deduplication. The API reads `CACHE_REDIS_HOST`, `CACHE_REDIS_PORT`, and `CACHE_REDIS_DB`, defaulting to `127.0.0.1`, `6379`, and `0`.
+
 To update the starter image:
 * `python setup.py sdist` to build the python package
 * `twine upload dist/*` to upload the package to pypi as `policyengine-api`

gcp/policyengine_api/start.sh

@@ -1,18 +1,25 @@
 #!/bin/sh
 # Environment variables
 PORT="${PORT:-8080}"
-REDIS_PORT="${REDIS_PORT:-6379}"
+CACHE_REDIS_HOST="${CACHE_REDIS_HOST:-127.0.0.1}"
+CACHE_REDIS_PORT="${CACHE_REDIS_PORT:-6379}"
+CACHE_REDIS_DB="${CACHE_REDIS_DB:-0}"
+export CACHE_REDIS_HOST CACHE_REDIS_PORT CACHE_REDIS_DB
 
-# Start the API
-gunicorn -b :"$PORT" policyengine_api.api --timeout 300 --workers 5 --preload &
-
-# Start Redis with configuration for multiple clients
-redis-server --protected-mode no \
+# Start Redis with configuration for multiple clients.
+redis-server --bind "$CACHE_REDIS_HOST" \
+             --port "$CACHE_REDIS_PORT" \
+             --protected-mode yes \
              --maxclients 10000 \
              --timeout 0 &
 
 # Wait for Redis to be ready
-sleep 2
+until redis-cli -h "$CACHE_REDIS_HOST" -p "$CACHE_REDIS_PORT" ping >/dev/null 2>&1; do
+  sleep 1
+done
+
+# Start the API
+gunicorn -b :"$PORT" policyengine_api.api --timeout 300 --workers 5 --preload &
 
 # Keep the script running and handle shutdown gracefully
 trap "pkill -P $$; exit 1" INT TERM

policyengine_api/api.py

@@ -4,6 +4,7 @@
 
 import time
 import sys
+import os
 
 start_time = time.time()
 
@@ -89,8 +90,9 @@ def log_timing(message):
     {
         "CACHE_TYPE": "RedisCache",
         "CACHE_KEY_PREFIX": "policyengine",
-        "CACHE_REDIS_HOST": "127.0.0.1",
-        "CACHE_REDIS_PORT": 6379,
+        "CACHE_REDIS_HOST": os.environ.get("CACHE_REDIS_HOST", "127.0.0.1"),
+        "CACHE_REDIS_PORT": int(os.environ.get("CACHE_REDIS_PORT", "6379")),
+        "CACHE_REDIS_DB": int(os.environ.get("CACHE_REDIS_DB", "0")),
         "CACHE_DEFAULT_TIMEOUT": 300,
     }
 )

policyengine_api/libs/simulation_api_modal.py

@@ -7,7 +7,7 @@
 
 import os
 import sys
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Optional
 
 import httpx
@@ -42,6 +42,28 @@ def name(self) -> str:
         return self.job_id
 
 
+@dataclass
+class ModalBudgetWindowBatchExecution:
+    """
+    Represents a budget-window batch execution in the Modal simulation API.
+    """
+
+    batch_job_id: str
+    status: str
+    progress: Optional[int] = None
+    completed_years: list[str] = field(default_factory=list)
+    running_years: list[str] = field(default_factory=list)
+    queued_years: list[str] = field(default_factory=list)
+    failed_years: list[str] = field(default_factory=list)
+    result: Optional[dict] = None
+    error: Optional[str] = None
+
+    @property
+    def name(self) -> str:
+        """Alias for batch_job_id."""
+        return self.batch_job_id
+
+
 class SimulationAPIModal:
     """
     HTTP client for the Modal Simulation API.
@@ -154,6 +176,57 @@ def run(self, payload: dict) -> ModalSimulationExecution:
             )
             raise
 
+    def run_budget_window_batch(self, payload: dict) -> ModalBudgetWindowBatchExecution:
+        """
+        Submit a budget-window batch job to the Modal API.
+        """
+        try:
+            modal_payload = dict(payload)
+            if "model_version" in modal_payload:
+                modal_payload["version"] = modal_payload.pop("model_version")
+            modal_payload.pop("data_version", None)
+
+            response = self.client.post(
+                f"{self.base_url}/simulate/economy/budget-window",
+                json=modal_payload,
+            )
+            response.raise_for_status()
+            data = response.json()
+
+            logger.log_struct(
+                {
+                    "message": "Modal budget-window batch submitted",
+                    "batch_job_id": data.get("batch_job_id"),
+                    "status": data.get("status"),
+                },
+                severity="INFO",
+            )
+
+            return ModalBudgetWindowBatchExecution(
+                batch_job_id=data["batch_job_id"],
+                status=data["status"],
+            )
+
+        except httpx.HTTPStatusError as e:
+            logger.log_struct(
+                {
+                    "message": f"Modal batch API HTTP error: {e.response.status_code}",
+                    "response_text": e.response.text[:500],
+                },
+                severity="ERROR",
+            )
+            raise
+
+        except httpx.RequestError as e:
+            logger.log_struct(
+                {
+                    "message": f"Modal batch API request error: {str(e)}",
+                    "run_id": (payload.get("_telemetry") or {}).get("run_id"),
+                },
+                severity="ERROR",
+            )
+            raise
+
     def resolve_app_name(
         self, country: str, version: Optional[str] = None
     ) -> tuple[str, str]:
@@ -235,6 +308,51 @@ def get_execution_by_id(self, job_id: str) -> ModalSimulationExecution:
             )
             raise
 
+    def get_budget_window_batch_by_id(
+        self, batch_job_id: str
+    ) -> ModalBudgetWindowBatchExecution:
+        """
+        Poll the Modal API for the current status of a budget-window batch.
+        """
+        try:
+            response = self.client.get(
+                f"{self.base_url}/budget-window-jobs/{batch_job_id}"
+            )
+            if response.status_code not in (200, 202, 500):
+                response.raise_for_status()
+            data = response.json()
+
+            return ModalBudgetWindowBatchExecution(
+                batch_job_id=batch_job_id,
+                status=data["status"],
+                progress=data.get("progress"),
+                completed_years=data.get("completed_years", []),
+                running_years=data.get("running_years", []),
+                queued_years=data.get("queued_years", []),
+                failed_years=data.get("failed_years", []),
+                result=data.get("result"),
+                error=data.get("error"),
+            )
+
+        except httpx.HTTPStatusError as e:
+            logger.log_struct(
+                {
+                    "message": f"Modal batch API HTTP error polling job {batch_job_id}: {e.response.status_code}",
+                    "response_text": e.response.text[:500],
+                },
+                severity="ERROR",
+            )
+            raise
+
+        except httpx.RequestError as e:
+            logger.log_struct(
+                {
+                    "message": f"Modal batch API request error polling job {batch_job_id}: {str(e)}",
+                },
+                severity="ERROR",
+            )
+            raise
+
     def get_execution_status(self, execution: ModalSimulationExecution) -> str:
         """
         Get the status string from an execution.