Merge pull request #1526 from PolicyEngine/codex/main-modal-worker-autoscaler

Keep production Modal workers warm

Author: anth-volk

.github/scripts/create_pr.sh

@@ -36,7 +36,7 @@ Related to #1178
 modal_release:
   new_app_target: frontier
   promote_existing_frontier: true
-  cleanup_target: none
+  cleanup_target: retired
 \`\`\`
 
 ## Version Updates

.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

.github/scripts/test_resolve_modal_release_config.py

@@ -69,6 +69,7 @@ def test_resolve_release_uses_weekly_default_for_empty_workflow_dispatch():
     assert resolved.should_deploy is True
     assert resolved.source == "workflow-dispatch-inputs"
     assert resolved.config is not None
+    assert resolved.config.cleanup_target == "retired"
 
 
 def test_resolve_release_skips_regular_push_even_with_pr_config():
@@ -130,3 +131,5 @@ def test_resolve_release_uses_weekly_default_when_no_pr_body_exists():
 
     assert resolved.should_deploy is True
     assert resolved.source == "weekly-default"
+    assert resolved.config is not None
+    assert resolved.config.cleanup_target == "retired"

.github/workflows/deploy-staged.yml

@@ -23,7 +23,7 @@ on:
       cleanup_target:
         description: "Modal app group to stop after updating the manifest"
         required: true
-        default: none
+        default: retired
         type: choice
         options:
           - none

changelog.d/1525.changed.md

@@ -0,0 +1 @@
+Keep a small warm Modal worker pool for the production environment and clean up retired Modal worker apps by default.

docs/engineering/skills/modal-release-prs.md

@@ -19,7 +19,7 @@ block:
 modal_release:
   new_app_target: frontier
   promote_existing_frontier: true
-  cleanup_target: none
+  cleanup_target: retired
 ```
 
 Allowed values:
@@ -40,13 +40,15 @@ different deployment behavior:
 modal_release:
   new_app_target: frontier
   promote_existing_frontier: true
-  cleanup_target: none
+  cleanup_target: retired
 ```
 
 This deploys the new worker to `frontier`, promotes the previous `frontier` to
 `current`, and moves the previous `current` into the manifest's `retired`
-history. Retired apps are not deleted unless `cleanup_target: retired` is
-explicitly configured.
+history. The default weekly release shape uses `cleanup_target: retired`, so
+apps in the retired history are stopped after the manifest is updated. Use
+`cleanup_target: none` only when the user explicitly asks to preserve retired
+worker apps after release.
 
 Do not use PR labels, branch names, model-specific tags, or title prefixes to
 control Modal release behavior. The PR body YAML block is the source of truth.

policyengine_household_api/modal_release/release_config.py

@@ -116,7 +116,7 @@ def default_weekly_config() -> ModalReleaseConfig:
     return ModalReleaseConfig(
         new_app_target=NewAppTarget.FRONTIER,
         promote_existing_frontier=True,
-        cleanup_target=CleanupTarget.NONE,
+        cleanup_target=CleanupTarget.RETIRED,
     )
 
 

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)