Merge pull request #1528 from PolicyEngine/enable-modal-memory-snapshot

Enable Modal memory snapshots on household API worker

Author: hua7450

changelog.d/1527.changed.md

@@ -0,0 +1 @@
+Enable Modal memory snapshots on the household API worker so cold starts restore the pre-loaded Flask app and policyengine country systems instead of re-running the ~45s import chain on every fresh container.

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

@@ -121,7 +121,8 @@ Modal deployment history for deploy provenance.
 
 For `/calculate` and `/calculate_demo`, the Modal gateway reads the top-level
 request field `version` and removes it before dispatching to the worker's
-`handle_household_request` Modal function. Accepted values are:
+`HouseholdWorker.handle_household_request` Modal class method. Accepted
+values are:
 
 - omitted or `current`: route to the current worker
 - `frontier`: route to the frontier worker

policyengine_household_api/modal_release/gateway.py

@@ -168,11 +168,22 @@ def resolve_app_for_request(
 def call_worker_function(app_name: str, payload: dict[str, Any]) -> Response:
     import modal
 
-    worker_function = modal.Function.from_name(
-        app_name,
-        "handle_household_request",
-    )
-    return _response_from_dispatch_result(worker_function.remote(payload))
+    # Prefer the class-based worker (post #1528). During the release
+    # transition the existing frontier is promoted to current without a
+    # redeploy, so for one release cycle the current worker may still expose
+    # the pre-#1528 top-level `handle_household_request` function. Fall back
+    # to that shape if the class is not present.
+    try:
+        worker_cls = modal.Cls.from_name(app_name, "HouseholdWorker")
+        return _response_from_dispatch_result(
+            worker_cls().handle_household_request.remote(payload)
+        )
+    except modal.exception.NotFoundError:
+        worker_function = modal.Function.from_name(
+            app_name,
+            "handle_household_request",
+        )
+        return _response_from_dispatch_result(worker_function.remote(payload))
 
 
 def _extract_requested_version(body: bytes) -> tuple[bytes, str]:

policyengine_household_api/modal_release/worker_app.py

@@ -45,6 +45,7 @@ def worker_function_options(
         "secrets": [household_api_secret()],
         "timeout": 180,
         "scaledown_window": 300,
+        "enable_memory_snapshot": True,
     }
     if environment == "main":
         options["min_containers"] = 3
@@ -53,9 +54,69 @@ def worker_function_options(
     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
+@app.cls(**worker_function_options())
+class HouseholdWorker:
+    """Worker class for handling household API requests.
 
-    return dispatch_to_flask_app(flask_app, payload)
+    Uses a Modal class with ``@modal.enter(snap=True)`` so the heavy Flask
+    app import runs at memory-snapshot creation time. Subsequent container
+    starts restore from the snapshot in seconds rather than re-running the
+    full policyengine country-package import chain on every cold start.
+    """
+
+    @modal.enter(snap=True)
+    def load_flask_app(self) -> None:
+        # Importing `policyengine_household_api.api` runs
+        # `initialize_analytics_db_if_enabled` at module level, which opens a
+        # Cloud SQL connection in environments where analytics is enabled.
+        # That connection needs GOOGLE_APPLICATION_CREDENTIALS, set by
+        # `configure_google_credentials()`. Configure credentials first so the
+        # snapshot-time import can succeed even before any request method runs.
+        configure_google_credentials()
+
+        from policyengine_household_api.api import app as flask_app
+
+        self.flask_app = flask_app
+
+    @modal.enter(snap=False)
+    def reset_post_snapshot_state(self) -> None:
+        # Runs on every container start AFTER snapshot restore. Memory
+        # snapshots preserve Python object state but not live network
+        # connections; the SQLAlchemy pool and the Cloud SQL Connector
+        # captured in the snapshot hold sockets that closed at snapshot
+        # time. Reset them so the first request opens fresh connections.
+        #
+        # Also force-recreate the Google credentials file: Modal preserves
+        # env vars across snapshot restore, but /tmp is not guaranteed to
+        # be preserved. Without popping the env var first,
+        # configure_google_credentials() would short-circuit on the
+        # surviving GOOGLE_APPLICATION_CREDENTIALS and leave it pointing
+        # at a missing file, breaking analytics DB reconnects.
+        # See: https://modal.com/docs/guide/memory-snapshot
+        os.environ.pop("GOOGLE_APPLICATION_CREDENTIALS", None)
+        configure_google_credentials()
+
+        from policyengine_household_api.data import analytics_setup
+
+        if not analytics_setup.is_analytics_enabled():
+            return
+
+        analytics_setup.cleanup()
+
+        try:
+            with self.flask_app.app_context():
+                analytics_setup.db.engine.dispose()
+        except Exception as exc:
+            import logging
+
+            logging.getLogger(__name__).warning(
+                "Failed to dispose analytics DB engine after snapshot "
+                "restore; subsequent queries may reconnect lazily: %s",
+                exc,
+            )
+
+    @modal.method()
+    def handle_household_request(
+        self, payload: dict[str, Any]
+    ) -> dict[str, Any]:
+        return dispatch_to_flask_app(self.flask_app, payload)

tests/unit/modal_release/test_gateway.py

@@ -267,3 +267,97 @@ def from_name(name, *, create_if_missing):
 
     with pytest.raises(modal.exception.AuthError):
         load_modal_manifest()
+
+
+def test_call_worker_function_uses_class_when_available(monkeypatch):
+    """When the worker exposes the new ``HouseholdWorker`` class, the
+    gateway should dispatch through ``modal.Cls.from_name``."""
+    from policyengine_household_api.modal_release import gateway
+
+    captured = {}
+
+    class _StubMethod:
+        def remote(self, payload):
+            captured["dispatched_via"] = "class"
+            captured["payload"] = payload
+            return {"status_code": 200, "body": b'{"status":"ok"}'}
+
+    class _StubInstance:
+        handle_household_request = _StubMethod()
+
+    class _StubCls:
+        def __call__(self):
+            return _StubInstance()
+
+    def fake_cls_from_name(app_name, class_name):
+        captured["cls_app_name"] = app_name
+        captured["class_name"] = class_name
+        return _StubCls()
+
+    def fake_function_from_name(app_name, function_name):
+        captured["fallback_invoked"] = True
+        raise AssertionError("Function fallback must not be invoked")
+
+    monkeypatch.setattr(
+        modal.Cls, "from_name", staticmethod(fake_cls_from_name)
+    )
+    monkeypatch.setattr(
+        modal.Function, "from_name", staticmethod(fake_function_from_name)
+    )
+
+    response = gateway.call_worker_function(
+        "frontier-app", {"household": {"foo": "bar"}}
+    )
+
+    assert response.status_code == 200
+    assert captured["dispatched_via"] == "class"
+    assert captured["cls_app_name"] == "frontier-app"
+    assert captured["class_name"] == "HouseholdWorker"
+    assert captured["payload"] == {"household": {"foo": "bar"}}
+    assert "fallback_invoked" not in captured
+
+
+def test_call_worker_function_falls_back_to_function_for_legacy_workers(
+    monkeypatch,
+):
+    """During a release transition, the existing frontier worker gets
+    promoted to current without a redeploy, so the current worker may
+    still expose the pre-class ``handle_household_request`` function.
+    The gateway must fall back to ``modal.Function.from_name`` when the
+    class cannot be found."""
+    from policyengine_household_api.modal_release import gateway
+
+    captured = {}
+
+    def fake_cls_from_name(app_name, class_name):
+        raise modal.exception.NotFoundError(
+            f"No class named `{class_name}` in app `{app_name}`"
+        )
+
+    class _StubFunction:
+        def remote(self, payload):
+            captured["dispatched_via"] = "function"
+            captured["payload"] = payload
+            return {"status_code": 200, "body": b'{"status":"ok"}'}
+
+    def fake_function_from_name(app_name, function_name):
+        captured["fn_app_name"] = app_name
+        captured["function_name"] = function_name
+        return _StubFunction()
+
+    monkeypatch.setattr(
+        modal.Cls, "from_name", staticmethod(fake_cls_from_name)
+    )
+    monkeypatch.setattr(
+        modal.Function, "from_name", staticmethod(fake_function_from_name)
+    )
+
+    response = gateway.call_worker_function(
+        "current-app", {"household": {"foo": "bar"}}
+    )
+
+    assert response.status_code == 200
+    assert captured["dispatched_via"] == "function"
+    assert captured["fn_app_name"] == "current-app"
+    assert captured["function_name"] == "handle_household_request"
+    assert captured["payload"] == {"household": {"foo": "bar"}}

tests/unit/modal_release/test_worker_app.py

@@ -52,6 +52,38 @@ def test_worker_function_options_do_not_keep_workers_warm_without_env(
         worker_function_options(modal_environment=None)
 
 
+def test_worker_function_options_enable_memory_snapshot_in_all_envs(
+    worker_app,
+):
+    for environment in ("main", "staging", "testing"):
+        options = worker_app.worker_function_options(
+            modal_environment=environment
+        )
+        assert options["enable_memory_snapshot"] is True, (
+            f"enable_memory_snapshot must be True in `{environment}` "
+            "so cold starts restore from a memory snapshot instead of "
+            "re-running the ~45s policyengine import chain"
+        )
+
+
+def test_household_worker_exposes_snapshot_entrypoint(worker_app):
+    """The class must declare its snapshot-time hook so heavy imports
+    are captured in the memory snapshot rather than running per cold
+    start."""
+    worker_cls = worker_app.HouseholdWorker
+    assert hasattr(worker_cls, "load_flask_app")
+    assert hasattr(worker_cls, "handle_household_request")
+
+
+def test_household_worker_exposes_post_snapshot_reset_hook(worker_app):
+    """The class must declare a post-restore hook so network state
+    captured in the memory snapshot (SQLAlchemy pool, Cloud SQL
+    Connector) gets reset on every container start. Modal preserves
+    Python object state but not live TCP sockets across snapshots."""
+    worker_cls = worker_app.HouseholdWorker
+    assert hasattr(worker_cls, "reset_post_snapshot_state")
+
+
 def test_country_package_install_specs_use_release_package_versions_only():
     assert country_package_install_specs(
         {