[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

@@ -106,7 +106,6 @@ def _init(self, container_args: api_pb2.ContainerArguments, client: _Client):
         self._waiting_for_memory_snapshot = False
         self._heartbeat_loop = None
         self._pause_heartbeats = asyncio.Condition()
-        self._snapshot_running = False
 
         self._is_interactivity_enabled = False
         self._fetching_inputs = True
@@ -163,15 +162,13 @@ async def _heartbeat_handle_cancellations(self) -> bool:
             request.current_input_started_at = self.current_input_started_at
 
         async with self._pause_heartbeats:
-            print("heartbeat acquired lock")
             while self._snapshot_running:
                 await self._pause_heartbeats.wait()
 
             # TODO(erikbern): capture exceptions?
             response = await retry_transient_errors(
                 self._client.stub.ContainerHeartbeat, request, attempt_timeout=HEARTBEAT_TIMEOUT
             )
-        print("heartbeat released lock")
 
         if response.HasField("cancel_input_event"):
             # Pause processing of the current input by signaling self a SIGUSR1.
@@ -206,11 +203,11 @@ async def _heartbeat_handle_cancellations(self) -> bool:
         return False
 
     @asynccontextmanager
-    async def heartbeats(self) -> AsyncGenerator[None, None]:
+    async def heartbeats(self, enable: bool) -> AsyncGenerator[None, None]:
         async with TaskContext() as tc:
             self._heartbeat_loop = t = tc.create_task(self._run_heartbeat_loop())
-            self._snapshot_running = False
             t.set_name("heartbeat loop")
+            self._snapshot_running = enable
             try:
                 yield
             finally:
@@ -587,10 +584,8 @@ async def memory_restore(self) -> None:
 
         # Turn heartbeats back on
         async with self._pause_heartbeats:
-            print("restore acquired lock")
             self._snapshot_running = False
             self._pause_heartbeats.notify_all()
-        print("restore released lock")
 
         logger.debug("Container: restored")
 
@@ -644,14 +639,12 @@ async def memory_snapshot(self) -> None:
 
         # Pause heartbeats since they keep the client connection open which causes the snapshotter to crash
         async with self._pause_heartbeats:
-            print("snapshot acquired lock")
             self._snapshot_running = True
             self._pause_heartbeats.notify_all()
 
         await self._client.stub.ContainerCheckpoint(
             api_pb2.ContainerCheckpointRequest(checkpoint_id=self.checkpoint_id)
         )
-        print("snapshot sent request")
 
         self._waiting_for_memory_snapshot = True
         await self._client._close(forget_credentials=True)