Address P1 review feedback on Modal snapshot rollout

Two fixes flagged by review on #1528:

1. Gateway backward compatibility during the release transition.
   The default Modal release flow promotes the existing frontier worker
   to current without redeploying it. With this PR shipped, the new
   worker exposes the `HouseholdWorker` class but the promoted current
   worker still only exposes the pre-PR top-level
   `handle_household_request` function. `call_worker_function` now tries
   `modal.Cls.from_name(..., "HouseholdWorker")` first and falls back to
   `modal.Function.from_name(..., "handle_household_request")` on
   `modal.exception.NotFoundError`, so both shapes route correctly
   during the one release cycle where both coexist. Added two gateway
   unit tests covering the class path and the function-fallback path.

2. Reset live network state captured by the memory snapshot.
   Memory snapshots preserve Python object state but not live TCP
   sockets; the SQLAlchemy pool and the Cloud SQL Connector created at
   snapshot time hold sockets that closed when Modal froze the
   container. Add `@modal.enter(snap=False) reset_post_snapshot_state`
   on `HouseholdWorker` to run on every snapshot-restored container
   start: re-establish the credentials file in case `/tmp` was not
   preserved, then call `analytics_setup.cleanup()` (drops the Cloud SQL
   Connector singleton) and `db.engine.dispose()` (drops the SQLAlchemy
   connection pool). Subsequent queries open fresh connections. Added a
   unit test asserting the hook is declared on the class.

References:
- https://modal.com/docs/guide/memory-snapshot
- https://modal.com/blog/mem-snapshots

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: hua7450

policyengine_household_api/modal_release/gateway.py

@@ -163,13 +163,22 @@ def resolve_app_for_request(
 def call_worker_function(app_name: str, payload: dict[str, Any]) -> Response:
     import modal
 
-    worker_cls = modal.Cls.from_name(
-        app_name,
-        "HouseholdWorker",
-    )
-    return _response_from_dispatch_result(
-        worker_cls().handle_household_request.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

@@ -78,13 +78,41 @@ def load_flask_app(self) -> None:
 
         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 re-runs `configure_google_credentials()` in case Modal did
+        # not preserve `/tmp` across snapshots and the credentials file is
+        # missing on the restored filesystem.
+        # See: https://modal.com/docs/guide/memory-snapshot
+        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]:
-        # Idempotent: if env var is already set (snapshot-restored or fresh
-        # container that ran the snap hook), this is a no-op. Kept here so
-        # filesystem-only credential state is re-established on restored
-        # containers if Modal does not preserve /tmp across snapshots.
-        configure_google_credentials()
         return dispatch_to_flask_app(self.flask_app, payload)

tests/unit/modal_release/test_gateway.py

@@ -231,3 +231,101 @@ def test_country_versions_rejects_unsupported_country():
     assert response.status_code == 404
     assert not worker_requests
     assert response.get_json()["message"] == "Unsupported country `zz`"
+
+
+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``."""
+    import modal
+
+    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."""
+    import modal
+
+    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

@@ -67,3 +67,12 @@ def test_household_worker_exposes_snapshot_entrypoint(worker_app):
     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")