[MOD-3251] Stop heartbeats before sending the container snapshot RPC

modal/_container_entrypoint.py

@@ -699,7 +699,7 @@ def main(container_args: api_pb2.ContainerArguments, client: Client):
 
     _client: _Client = synchronizer._translate_in(client)  # TODO(erikbern): ugly
 
-    with container_io_manager.heartbeats(), UserCodeEventLoop() as event_loop:
+    with container_io_manager.heartbeats(function_def.is_checkpointing_function), UserCodeEventLoop() as event_loop:
         # If this is a serialized function, fetch the definition from the server
         if function_def.definition_type == api_pb2.Function.DEFINITION_TYPE_SERIALIZED:
             ser_cls, ser_fun = container_io_manager.get_serialized_function()
@@ -779,6 +779,7 @@ def main(container_args: api_pb2.ContainerArguments, client: Client):
         if function_def.is_checkpointing_function:
             container_io_manager.memory_snapshot()
 
+
         # Install hooks for interactive functions.
         if function_def.pty_info.pty_type != api_pb2.PTYInfo.PTY_TYPE_UNSPECIFIED:
 

modal/_container_io_manager.py

@@ -73,8 +73,9 @@ class _ContainerIOManager:
     _input_concurrency: Optional[int]
     _semaphore: Optional[asyncio.Semaphore]
     _environment_name: str
-    _waiting_for_memory_snapshot: bool
     _heartbeat_loop: Optional[asyncio.Task]
+    _pause_heartbeats: Optional[asyncio.Condition]
+    _waiting_for_memory_snapshot: bool
 
     _is_interactivity_enabled: bool
     _fetching_inputs: bool
@@ -101,8 +102,9 @@ def _init(self, container_args: api_pb2.ContainerArguments, client: _Client):
 
         self._semaphore = None
         self._environment_name = container_args.environment_name
-        self._waiting_for_memory_snapshot = False
         self._heartbeat_loop = None
+        self._pause_heartbeats = asyncio.Condition()
+        self._waiting_for_memory_snapshot = False
 
         self._is_interactivity_enabled = False
         self._fetching_inputs = True
@@ -146,22 +148,20 @@ async def _heartbeat_handle_cancellations(self) -> bool:
         # Return True if a cancellation event was received, in that case
         # we shouldn't wait too long for another heartbeat
 
-        # Don't send heartbeats for tasks waiting to be checkpointed.
-        # Calling gRPC methods open new connections which block the
-        # checkpointing process.
-        if self._waiting_for_memory_snapshot:
-            return False
-
         request = api_pb2.ContainerHeartbeatRequest(supports_graceful_input_cancellation=True)
         if self.current_input_id is not None:
             request.current_input_id = self.current_input_id
         if self.current_input_started_at is not None:
             request.current_input_started_at = self.current_input_started_at
 
-        # TODO(erikbern): capture exceptions?
-        response = await retry_transient_errors(
-            self._client.stub.ContainerHeartbeat, request, attempt_timeout=HEARTBEAT_TIMEOUT
-        )
+        async with self._pause_heartbeats:
+            while self._waiting_for_memory_snapshot:
+                await self._pause_heartbeats.wait()
+
+            # TODO(erikbern): capture exceptions?
+            response = await retry_transient_errors(
+                self._client.stub.ContainerHeartbeat, request, attempt_timeout=HEARTBEAT_TIMEOUT
+            )
 
         if response.HasField("cancel_input_event"):
             # Pause processing of the current input by signaling self a SIGUSR1.
@@ -196,10 +196,11 @@ async def _heartbeat_handle_cancellations(self) -> bool:
         return False
 
     @asynccontextmanager
-    async def heartbeats(self) -> AsyncGenerator[None, None]:
+    async def heartbeats(self, disable_init: bool) -> AsyncGenerator[None, None]:
         async with TaskContext() as tc:
             self._heartbeat_loop = t = tc.create_task(self._run_heartbeat_loop())
             t.set_name("heartbeat loop")
+            self._waiting_for_memory_snapshot = disable_init
             try:
                 yield
             finally:
@@ -574,6 +575,11 @@ async def memory_restore(self) -> None:
             await asyncio.sleep(0.01)
             continue
 
+        # Turn heartbeats back on
+        async with self._pause_heartbeats:
+            self._waiting_for_memory_snapshot = False
+            self._pause_heartbeats.notify_all()
+
         logger.debug("Container: restored")
 
         # Look for state file and create new client with updated credentials.
@@ -617,18 +623,21 @@ async def memory_restore(self) -> None:
                 )
 
         self._client = await _Client.from_env()
-        self._waiting_for_memory_snapshot = False
 
     async def memory_snapshot(self) -> None:
         """Message server indicating that function is ready to be checkpointed."""
         if self.checkpoint_id:
             logger.debug(f"Checkpoint ID: {self.checkpoint_id} (Memory Snapshot ID)")
 
+        # Pause heartbeats since they keep the client connection open which causes the snapshotter to crash
+        async with self._pause_heartbeats:
+            self._waiting_for_memory_snapshot = True
+            self._pause_heartbeats.notify_all()
+
         await self._client.stub.ContainerCheckpoint(
             api_pb2.ContainerCheckpointRequest(checkpoint_id=self.checkpoint_id)
         )
 
-        self._waiting_for_memory_snapshot = True
         await self._client._close(forget_credentials=True)
 
         logger.debug("Memory snapshot request sent. Connection closed.")

test/container_app_test.py

@@ -4,6 +4,7 @@
 import pytest
 from typing import Dict
 from unittest import mock
+import asyncio
 
 from google.protobuf.empty_pb2 import Empty
 from google.protobuf.message import Message
@@ -58,12 +59,43 @@ async def test_container_snapshot_restore(container_client, tmpdir, servicer):
         encoding="utf-8",
     )
     with mock.patch.dict(
-        os.environ, {"MODAL_RESTORE_STATE_PATH": str(restore_path), "MODAL_SERVER_URL": servicer.container_addr}
+        os.environ, {"MODAL_RESTORE_STATE_PATH": str(restore_path), "MODAL_SERVER_URL": servicer.container_addr},
     ):
         io_manager.memory_snapshot()
         # In-memory Client instance should have update credentials, not old credentials
         assert old_client.credentials == ("ta-i-am-restored", "ts-i-am-restored")
 
+@pytest.mark.asyncio
+async def test_container_snapshot_restore_heartbeats(container_client, tmpdir, servicer):
+    io_manager = ContainerIOManager(api_pb2.ContainerArguments(), container_client)
+    restore_path = tmpdir.join("fake-restore-state.json")
+    restore_path.write_text(
+        json.dumps(
+            {
+                "task_id": "ta-i-am-restored",
+                "task_secret": "ts-i-am-restored",
+                "function_id": "fu-i-am-restored",
+            }
+        ),
+        encoding="utf-8",
+    )
+
+    asyncio.set_event_loop(asyncio.new_event_loop())
+
+    # Ensure that heartbeats do not run before the snapshot
+    # Ensure that heartbeats do run after the snapshot
+    with io_manager.heartbeats(True):
+        with mock.patch.dict(
+            os.environ, {"MODAL_RESTORE_STATE_PATH": str(restore_path), "MODAL_SERVER_URL": servicer.container_addr},
+        ):
+            with mock.patch("modal.runner.HEARTBEAT_INTERVAL", 1):
+                with mock.patch.object(container_client.stub, 'ContainerHeartbeat') as mock_heartbeat:
+
+                    await asyncio.sleep(2)
+                    mock_heartbeat.assert_not_called()
+                    io_manager.memory_snapshot()
+                    mock_heartbeat.assert_called_once()
+
 
 @pytest.mark.asyncio
 async def test_container_debug_snapshot(container_client, tmpdir, servicer):

test/container_test.py

@@ -819,6 +819,9 @@ def test_container_heartbeats(servicer):
     _run_container(servicer, "test.supports.functions", "square")
     assert any(isinstance(request, api_pb2.ContainerHeartbeatRequest) for request in servicer.requests)
 
+    _run_container(servicer, "test.supports.functions", "snapshotting_square")
+    assert any(isinstance(request, api_pb2.ContainerHeartbeatRequest) for request in servicer.requests)
+
 
 @skip_github_non_linux
 def test_cli(servicer):

test/live_reload_test.py

@@ -78,3 +78,4 @@ async def test_heartbeats(app_ref, server_url_env, servicer):
     # Typically [0s, 1s, 2s, 3s], but asyncio.sleep may lag.
     actual_heartbeats = servicer.app_heartbeats[apps[0]]
     assert abs(actual_heartbeats - (total_secs + 1)) <= 1
+

test/supports/functions.py

@@ -414,6 +414,11 @@ def f(self, x):
         return "".join(self._vals) + x
 
 
+@app.function(enable_memory_snapshot=True)
+def snapshotting_square(x):
+    return x * x
+
+
 @app.cls()
 class EventLoopCls:
     @enter()