From 863be2f739d300e590537223dc5056aea3036275 Mon Sep 17 00:00:00 2001
From: Eric Lunderberg <elunderberg@octoml.ai>
Date: Thu, 9 May 2024 21:02:52 +0000
Subject: [PATCH] [Disco] Treat hangup of disco worker process as kShutdown

Prior to this commit, each disco worker needed to receive
`DiscoAction::kShutdown` in order to close cleanly.  While this is
sent from the destructor of `ProcessSessionObj`, which owns the worker
processes, this does not guarantee that the disco workers will receive
the shutdown command.  For example, the controller process holding the
`ProcessSessionObj` may reach a timeout and be terminated, preventing
it from sending the `DiscoAction::kShutdown` command.

This commit updates the disco worker to check for a closed pipe that
occurs between two packets, and to treat this as if the
`DiscoAction::kShutdown` command were received.  A closed pipe that
occurs at any other location is still treated as an error and
reported.
---
 src/runtime/disco/process_session.cc | 38 +++++++++++++++++++++++++---
 1 file changed, 34 insertions(+), 4 deletions(-)

diff --git a/src/runtime/disco/process_session.cc b/src/runtime/disco/process_session.cc
index 6474db479e94..6687a64e7f85 100644
--- a/src/runtime/disco/process_session.cc
+++ b/src/runtime/disco/process_session.cc
@@ -48,11 +48,21 @@ class DiscoPipeMessageQueue : private dmlc::Stream, private DiscoProtocol<DiscoP
   }
 
   TVMArgs Recv() {
-    DequeueNextPacket();
+    bool is_implicit_shutdown = DequeueNextPacket();
     TVMValue* values = nullptr;
     int* type_codes = nullptr;
     int num_args = 0;
-    RPCReference::RecvPackedSeq(&values, &type_codes, &num_args, this);
+
+    if (is_implicit_shutdown) {
+      num_args = 2;
+      values = ArenaAlloc<TVMValue>(num_args);
+      type_codes = ArenaAlloc<int>(num_args);
+      TVMArgsSetter setter(values, type_codes);
+      setter(0, static_cast<int>(DiscoAction::kShutDown));
+      setter(1, 0);
+    } else {
+      RPCReference::RecvPackedSeq(&values, &type_codes, &num_args, this);
+    }
     return TVMArgs(values, type_codes, num_args);
   }
 
@@ -62,18 +72,38 @@ class DiscoPipeMessageQueue : private dmlc::Stream, private DiscoProtocol<DiscoP
     write_buffer_.clear();
   }
 
-  void DequeueNextPacket() {
+  /* \brief Read next packet and reset unpacker
+   *
+   * Read the next packet into `read_buffer_`, releasing all arena
+   * allocations performed by the unpacker and resetting the unpacker
+   * to its initial state.
+   *
+   * \return A boolean value.  If true, this packet should be treated
+   *    equivalently to a `DiscoAction::kShutdown` event.  If false,
+   *    this packet should be unpacked.
+   */
+  bool DequeueNextPacket() {
     uint64_t packet_nbytes = 0;
     int read_size = pipe_.Read(&packet_nbytes, sizeof(packet_nbytes));
+    if (read_size == 0) {
+      // Special case, connection dropped between packets.  Treat as a
+      // request to shutdown.
+      return true;
+    }
+
     ICHECK_EQ(read_size, sizeof(packet_nbytes))
         << "Pipe closed without proper shutdown. Please make sure to explicitly call "
            "`Session::Shutdown`";
     read_buffer_.resize(packet_nbytes);
-    pipe_.Read(read_buffer_.data(), packet_nbytes);
+    read_size = pipe_.Read(read_buffer_.data(), packet_nbytes);
+    ICHECK_EQ(read_size, packet_nbytes)
+        << "Pipe closed without proper shutdown. Please make sure to explicitly call "
+           "`Session::Shutdown`";
     read_offset_ = 0;
     this->RecycleAll();
     RPCCode code = RPCCode::kReturn;
     this->Read(&code);
+    return false;
   }
 
   size_t Read(void* data, size_t size) final {