Keep production Modal workers warm

Author: anth-volk

.github/scripts/modal-deploy-release.sh

@@ -2,7 +2,6 @@
 set -euo pipefail
 
 config_json="${1:?Usage: modal-deploy-release.sh CONFIG_JSON}"
-modal_environment="${MODAL_ENVIRONMENT:-main}"
 output_file="${GITHUB_OUTPUT:-}"
 
 require_env() {
@@ -36,10 +35,13 @@ github_output() {
 }
 
 require_env \
+  MODAL_ENVIRONMENT \
   USER_ANALYTICS_DB_USERNAME \
   USER_ANALYTICS_DB_PASSWORD \
   USER_ANALYTICS_DB_CONNECTION_NAME
 
+modal_environment="${MODAL_ENVIRONMENT}"
+
 uv run alembic upgrade head
 analytics_database_revision="$(
   uv run python -m policyengine_household_api.modal_release.analytics_revision
@@ -67,6 +69,7 @@ bash .github/scripts/modal-sync-secrets.sh
 new_app_target="$(config_value new_app_target)"
 if [ "${new_app_target}" != "none" ]; then
   HOUSEHOLD_MODAL_WORKER_APP_NAME="${worker_app_name}" \
+    MODAL_ENVIRONMENT="${modal_environment}" \
     uv run modal deploy \
       --env "${modal_environment}" \
       -m policyengine_household_api.modal_release.worker_app

.github/scripts/test_modal_deploy_release.py

@@ -0,0 +1,26 @@
+import os
+import subprocess
+
+
+def test_modal_deploy_release_requires_explicit_modal_environment():
+    env = {
+        **os.environ,
+        "USER_ANALYTICS_DB_USERNAME": "user",
+        "USER_ANALYTICS_DB_PASSWORD": "password",
+        "USER_ANALYTICS_DB_CONNECTION_NAME": "project:region:instance",
+    }
+    env.pop("MODAL_ENVIRONMENT", None)
+
+    result = subprocess.run(
+        [
+            "bash",
+            ".github/scripts/modal-deploy-release.sh",
+            '{"new_app_target":"none","promote_existing_frontier":false,"cleanup_target":"none"}',
+        ],
+        capture_output=True,
+        env=env,
+        text=True,
+    )
+
+    assert result.returncode == 1
+    assert "MODAL_ENVIRONMENT" in result.stdout

changelog.d/1525.changed.md

@@ -0,0 +1 @@
+Keep a small warm Modal worker pool for the production environment.

policyengine_household_api/modal_release/worker_app.py

@@ -23,12 +23,37 @@
 )
 
 
-@app.function(
-    image=household_api_worker_image(),
-    secrets=[household_api_secret()],
-    timeout=180,
-    scaledown_window=300,
-)
+def worker_modal_environment(
+    modal_environment: str | None = None,
+) -> str:
+    environment = (
+        modal_environment
+        if modal_environment is not None
+        else os.getenv("MODAL_ENVIRONMENT")
+    )
+    if not environment:
+        raise RuntimeError("MODAL_ENVIRONMENT must be set for Modal workers")
+    return environment
+
+
+def worker_function_options(
+    modal_environment: str | None = None,
+) -> dict[str, Any]:
+    environment = worker_modal_environment(modal_environment)
+    options: dict[str, Any] = {
+        "image": household_api_worker_image(),
+        "secrets": [household_api_secret()],
+        "timeout": 180,
+        "scaledown_window": 300,
+    }
+    if environment == "main":
+        options["min_containers"] = 3
+        options["buffer_containers"] = 2
+        options["scaledown_window"] = 600
+    return options
+
+
+@app.function(**worker_function_options())
 def handle_household_request(payload: dict[str, Any]) -> dict[str, Any]:
     configure_google_credentials()
     from policyengine_household_api.api import app as flask_app

tests/unit/modal_release/test_worker_app.py

@@ -0,0 +1,46 @@
+import importlib
+
+import pytest
+
+
+pytestmark = pytest.mark.usefixtures("worker_app")
+
+
+@pytest.fixture
+def worker_app(monkeypatch):
+    monkeypatch.setenv("MODAL_ENVIRONMENT", "testing")
+    from policyengine_household_api.modal_release import worker_app
+
+    return importlib.reload(worker_app)
+
+
+def test_worker_function_options_keep_main_workers_warm(worker_app):
+    options = worker_app.worker_function_options(modal_environment="main")
+
+    assert options["min_containers"] == 3
+    assert options["buffer_containers"] == 2
+    assert options["scaledown_window"] == 600
+
+
+def test_worker_function_options_do_not_keep_staging_workers_warm():
+    from policyengine_household_api.modal_release.worker_app import (
+        worker_function_options,
+    )
+
+    options = worker_function_options(modal_environment="staging")
+
+    assert "min_containers" not in options
+    assert "buffer_containers" not in options
+    assert options["scaledown_window"] == 300
+
+
+def test_worker_function_options_do_not_keep_workers_warm_without_env(
+    monkeypatch,
+):
+    monkeypatch.delenv("MODAL_ENVIRONMENT", raising=False)
+    from policyengine_household_api.modal_release.worker_app import (
+        worker_function_options,
+    )
+
+    with pytest.raises(RuntimeError, match="MODAL_ENVIRONMENT"):
+        worker_function_options(modal_environment=None)