Merge branch 'master' of github.com:cmu-db/noisepage into replication

src/include/main/db_main.h

@@ -17,6 +17,7 @@
 #include "network/postgres/postgres_protocol_interpreter.h"
 #include "optimizer/statistics/stats_storage.h"
 #include "replication/replication_manager.h"
+#include "self_driving/model_server/model_server_manager.h"
 #include "settings/settings_manager.h"
 #include "settings/settings_param.h"
 #include "storage/garbage_collector_thread.h"
@@ -433,6 +434,13 @@ class DBMain {
                                            network_port_, connection_thread_count_, uds_file_directory_);
       }
 
+      std::unique_ptr<modelserver::ModelServerManager> model_server_manager = DISABLED;
+      if (model_server_enable_) {
+        NOISEPAGE_ASSERT(use_messenger_, "Pilot requires messenger layer.");
+        model_server_manager =
+            std::make_unique<modelserver::ModelServerManager>(model_server_path_, messenger_layer->GetMessenger());
+      }
+
       db_main->settings_manager_ = std::move(settings_manager);
       db_main->metrics_manager_ = std::move(metrics_manager);
       db_main->metrics_thread_ = std::move(metrics_thread);
@@ -783,6 +791,13 @@ class DBMain {
     uint16_t messenger_port_ = 9022;
     uint16_t replication_port_ = 15445;
     std::string replication_hosts_path_ = "replication.conf";
+    bool model_server_enable_ = false;
+    /**
+     * The ModelServer script is located at PROJECT_ROOT/script/model by default, and also assume
+     * the build binary at PROJECT_ROOT/build/bin/noisepage. This should be override or set explicitly
+     * in use cases where such assumptions are no longer true.
+     */
+    std::string model_server_path_ = "../../script/model/model_server.py";
 
     /**
      * Instantiates the SettingsManager and reads all of the settings to override the Builder's settings.
@@ -844,6 +859,8 @@ class DBMain {
       use_replication_ = settings_manager->GetBool(settings::Param::replication_enable);
       replication_port_ = settings_manager->GetInt(settings::Param::replication_port);
       replication_hosts_path_ = settings_manager->GetString(settings::Param::replication_hosts_path);
+      model_server_enable_ = settings_manager->GetBool(settings::Param::model_server_enable);
+      model_server_path_ = settings_manager->GetString(settings::Param::model_server_path);
 
       return settings_manager;
     }

src/include/replication/replication_manager.h

@@ -28,7 +28,7 @@ class Replica {
 
   messenger::ConnectionDestination replica_info_;
   messenger::ConnectionId connection_;
-  uint64_t last_heartbeat_;  //< Time (unix epoch) that the replica heartbeat was last successful.
+  uint64_t last_heartbeat_;  ///< Time (unix epoch) that the replica heartbeat was last successful.
 };
 
 /**
@@ -91,9 +91,9 @@ class ReplicationManager {
   common::ManagedPointer<messenger::ConnectionId> GetReplicaConnection(const std::string &replica_name);
 
   common::ManagedPointer<messenger::Messenger> messenger_;
-  std::string identity_;                               //< The identity of this replica.
-  uint16_t port_;                                      //< The port that replication runs on.
-  std::unordered_map<std::string, Replica> replicas_;  //< Replica Name -> Connection ID.
+  std::string identity_;                               ///< The identity of this replica.
+  uint16_t port_;                                      ///< The port that replication runs on.
+  std::unordered_map<std::string, Replica> replicas_;  ///< Replica Name -> Connection ID.
   std::mutex mutex_;
   std::condition_variable cvar_;
 };

src/include/self_driving/model_server/model_server_manager.h

@@ -12,14 +12,19 @@
 
 #pragma once
 
+#include <atomic>
+#include <condition_variable>
 #include <string>
+#include <thread>
 #include <utility>
 #include <vector>
 
 #include "common/json.h"
 #include "common/managed_pointer.h"
+#include "messenger/messenger_defs.h"
 
 namespace noisepage::messenger {
+class ConnectionRouter;
 class Messenger;
 }  // namespace noisepage::messenger
 

src/include/settings/settings_defs.h

@@ -368,4 +368,21 @@ SETTING_string(
     false,
     noisepage::settings::Callbacks::NoOp
 )
+
+SETTING_bool(
+    model_server_enable,
+    "Whether to enable the ModelServerManager (default: false)",
+    false,
+    false,
+    noisepage::settings::Callbacks::NoOp
+)
+
+// Relative path assuming binary locate at PROJECT_ROOT/build/bin/, and model_server.py at PROJECT_ROOT/script/model
+SETTING_string(
+    model_server_path,
+    "The python model server script to invoke (default: ../../script/model/model_server.py)",
+    "../../script/model/model_server.py",
+    false,
+    noisepage::settings::Callbacks::NoOp
+)
     // clang-format on

src/messenger/messenger.cpp

@@ -93,6 +93,81 @@
 
 namespace noisepage::messenger {
 
+///** An abstraction around ZeroMQ messages which explicitly have the sender specified. */
+// class ZmqMessage {
+// public:
+//  /**
+//   * Build a new ZmqMessage from the supplied information.
+//   * @param send_msg_id     The ID of the message on the sender side.
+//   * @param recv_cb_id      The ID of the message on the receiver side.
+//   * @param sender_id       The identity of the sender.
+//   * @param message         The contents of the message.
+//   * @return A ZmqMessage encapsulating the given message.
+//   */
+//  static ZmqMessage Build(uint64_t send_msg_id, uint64_t recv_cb_id, const std::string &sender_id,
+//                          std::string_view message) {
+//    return ZmqMessage{sender_id, fmt::format("{}-{}-{}", send_msg_id, recv_cb_id, message)};
+//  }
+//
+//  /**
+//   * Parse the given payload into a ZmqMessage.
+//   * @param routing_id      The message's routing ID.
+//   * @param message         The message received.
+//   * @return A ZmqMessage encapsulating the given message.
+//   */
+//  static ZmqMessage Parse(const std::string &routing_id, const std::string &message) {
+//    return ZmqMessage{routing_id, message};
+//  }
+//
+//  /** @return The ID of this message (sender side). */
+//  uint64_t GetMessageIdSender() const { return send_msg_id_; }
+//
+//  /** @return The ID of this message (receiver side). */
+//  uint64_t GetCallbackIdReceiver() const { return recv_cb_id_; }
+//
+//  /** @return The routing ID of this message. */
+//  std::string_view GetRoutingId() const { return std::string_view(routing_id_); }
+//
+//  /** @return The message itself. */
+//  std::string_view GetMessage() const { return message_; }
+//
+//  /** @return The raw payload of the message. */
+//  std::string_view GetRawPayload() const { return std::string_view(payload_); }
+//
+// private:
+//  friend Messenger;
+//
+//  /** Construct a new ZmqMessage with the given routing ID and payload. Payload of form ID-MESSAGE. */
+//  ZmqMessage(std::string routing_id, std::string payload)
+//      : routing_id_(std::move(routing_id)), payload_(std::move(payload)), message_(payload_) {
+//    if (payload_.empty()) {
+//      send_msg_id_ = 0;
+//      recv_cb_id_ = 0;
+//    } else {
+//      // TODO(WAN): atoi, stoull, from_chars, etc? Error checking in general.
+//      UNUSED_ATTRIBUTE int check =
+//          std::sscanf(payload_.c_str(), "%" SCNu64 "-%" SCNu64 "-", &send_msg_id_, &recv_cb_id_);
+//      NOISEPAGE_ASSERT(2 == check, "Couldn't parse the message header.");
+//
+//      // Remove the prefix up to the second '-'
+//      message_.remove_prefix(message_.find_first_of('-') + 1);
+//      message_.remove_prefix(message_.find_first_of('-') + 1);
+//    }
+//  }
+//
+//  /** The routing ID of the message. */
+//  std::string routing_id_;
+//  /** The payload in the message, of form ID-MESSAGE.  */
+//  std::string payload_;
+//
+//  /** The cached id of the message (sender side). */
+//  uint64_t send_msg_id_;
+//  /** The cached id of the message (receiver side). */
+//  uint64_t recv_cb_id_;
+//  /** The cached actual message. */
+//  std::string_view message_;
+//};
+
 ZmqMessage ZmqMessage::Build(uint64_t source_cb_id, uint64_t dest_cb_id, const std::string &routing_id,
                              std::string_view message) {
   return ZmqMessage{routing_id, fmt::format("{}-{}-{}", source_cb_id, dest_cb_id, message)};
@@ -112,7 +187,10 @@ ZmqMessage::ZmqMessage(std::string routing_id, std::string payload)
     UNUSED_ATTRIBUTE int check =
         std::sscanf(payload_.c_str(), "%" SCNu64 "-%" SCNu64 "-", &source_cb_id_, &dest_cb_id_);
     NOISEPAGE_ASSERT(2 == check, "Couldn't parse the message header.");
-    message_.remove_prefix(message_.find_last_of('-') + 1);
+
+    // Remove the prefix up to the second '-'
+    message_.remove_prefix(message_.find_first_of('-') + 1);
+    message_.remove_prefix(message_.find_first_of('-') + 1);
   }
 }
 
@@ -276,20 +354,14 @@ ConnectionId::ConnectionId(common::ManagedPointer<Messenger> messenger, const Co
   socket_->set(zmq::sockopt::linger, 0);             // Discard all pending messages immediately on socket close.
   socket_->connect(target.GetDestination());
   routing_id_ = ZmqUtil::GetRoutingId(common::ManagedPointer(socket_));
-  MESSENGER_LOG_INFO(fmt::format("[PID={}] Registered (but haven't opened!) a new connection to {} ({}) as {}.",
+  MESSENGER_LOG_TRACE(fmt::format("[PID={}] Registered (but haven't opened!) a new connection to {} ({}) as {}.",
                                  ::getpid(), target_name_, target.GetDestination(), routing_id_.c_str()));
   // Add the new socket to the list of sockets that will be polled by the server loop.
   messenger->polled_sockets_->AddPollItem(socket_.get(), nullptr);
 }
 
 ConnectionId::~ConnectionId() = default;
 
-ConnectionId::ConnectionId(ConnectionId &&other) {
-  this->socket_ = std::move(other.socket_);
-  this->routing_id_ = std::move(other.routing_id_);
-  this->target_name_ = std::move(other.target_name_);
-}
-
 ConnectionRouter::ConnectionRouter(common::ManagedPointer<Messenger> messenger, const ConnectionDestination &target,
                                    const std::string &identity, CallbackFn callback)
     : callback_(std::move(callback)), identity_(identity) {
@@ -392,7 +464,7 @@ ConnectionId Messenger::MakeConnection(const ConnectionDestination &target) {
 }
 
 void Messenger::SendMessage(common::ManagedPointer<ConnectionId> connection_id, const std::string &message,
-                            CallbackFn callback, uint64_t remote_cb_id) {
+                            CallbackFn callback, uint64_t recv_cb_id) {
   // Note that a ConnectionId cannot be used from multiple threads. This is an inherent ZeroMQ limitation.
   // If you need another thread to connect to the same ConnectionDestination, just get a new ConnectionId.
   // This allows us to avoid taking latches around the socket that the ConnectionId wraps.
@@ -404,14 +476,14 @@ void Messenger::SendMessage(common::ManagedPointer<ConnectionId> connection_id,
   callbacks_mutex_.unlock();
 
   // Build and send the message.
-  ZmqMessage msg = ZmqMessage::Build(send_msg_id, remote_cb_id, connection_id->routing_id_, message);
+  ZmqMessage msg = ZmqMessage::Build(send_msg_id, recv_cb_id, connection_id->routing_id_, message);
   ZmqUtil::SendMsgPayload(common::ManagedPointer(connection_id->socket_), msg);
   MESSENGER_LOG_TRACE(
       fmt::format("[PID={}] Messenger SENT-TO {}: {} ", ::getpid(), connection_id->target_name_, msg.GetRawPayload()));
 }
 
 void Messenger::SendMessage(common::ManagedPointer<ConnectionRouter> router_id, const std::string &recv_id,
-                            const std::string &message, CallbackFn callback, uint64_t remote_cb_id) {
+                            const std::string &message, CallbackFn callback, uint64_t recv_cb_id) {
   uint64_t send_msg_id = GetNextSendMessageId();
 
   // Register the callback that will be invoked when a response to this message is received.
@@ -422,7 +494,7 @@ void Messenger::SendMessage(common::ManagedPointer<ConnectionRouter> router_id,
   // Build and send the message. Note that ConnectionRouter is a ROUTER socket.
   zmq::message_t router_data(recv_id.data(), recv_id.size());
   if (router_id->socket_->send(router_data, zmq::send_flags::sndmore).has_value()) {
-    ZmqMessage reply = ZmqMessage::Build(send_msg_id, remote_cb_id, router_id->identity_, message);
+    ZmqMessage reply = ZmqMessage::Build(send_msg_id, recv_cb_id, router_id->identity_, message);
     ZmqUtil::SendMsgIdentity(common::ManagedPointer(router_id->socket_.get()), router_id->identity_);
     ZmqUtil::SendMsgPayload(common::ManagedPointer(router_id->socket_.get()), reply);
     MESSENGER_LOG_TRACE(fmt::format("[PID={}] Messenger ({}) SENT-TO {}: {} ", ::getpid(), router_id->identity_,
@@ -476,8 +548,9 @@ void Messenger::ServerLoop() {
         bool has_custom_serverloop = poll_items.server_callbacks_[i] != nullptr;
         MESSENGER_LOG_TRACE("[PID={}] Messenger RECV-FR {} (custom serverloop: {}): {}", ::getpid(), msg.GetRoutingId(),
                             has_custom_serverloop, msg.GetRawPayload());
-        ProcessMessage(msg);
-        if (has_custom_serverloop) {
+        if (!has_custom_serverloop) {
+          ProcessMessage(msg);
+        } else {
           auto &server_callback = poll_items.server_callbacks_[i];
           (*server_callback)(common::ManagedPointer(this), msg);
         }

src/replication/replication_manager.cpp

@@ -24,6 +24,10 @@ ReplicationManager::ReplicationManager(common::ManagedPointer<noisepage::messeng
                                   [this](common::ManagedPointer<messenger::Messenger> messenger,
                                          const messenger::ZmqMessage &msg) { EventLoop(messenger, msg); });
   BuildReplicaList(replication_hosts_path);
+
+  // TODO(WAN): development purposes
+  replication_logger->set_level(spdlog::level::trace);
+
   for (const auto &replica : replicas_) {
     ReplicaHeartbeat(replica.first);
   }
@@ -105,7 +109,7 @@ void ReplicationManager::EventLoop(common::ManagedPointer<noisepage::messenger::
                                    const noisepage::messenger::ZmqMessage &msg) {
   switch (static_cast<MessageType>(msg.GetDestinationCallbackId())) {
     case MessageType::HEARTBEAT:
-      REPLICATION_LOG_INFO(fmt::format("Heartbeat from: {}", msg.GetRoutingId()));
+      REPLICATION_LOG_TRACE(fmt::format("Heartbeat from: {}", msg.GetRoutingId()));
       break;
     default:
       break;
@@ -121,10 +125,12 @@ void ReplicationManager::ReplicaHeartbeat(const std::string &replica_name) {
     auto epoch_now_ms = std::chrono::duration_cast<std::chrono::milliseconds>(epoch_now);
     if (0 == replica.last_heartbeat_) {
       replica.last_heartbeat_ = epoch_now_ms.count();
+      REPLICATION_LOG_TRACE(
+          fmt::format("Replica {}: heartbeat initialized at {}.", replica_name, replica.last_heartbeat_));
     }
   }
 
-  REPLICATION_LOG_INFO(fmt::format("Replica {}: heartbeat start.", replica_name));
+  REPLICATION_LOG_TRACE(fmt::format("Replica {}: heartbeat start.", replica_name));
   try {
     messenger_->SendMessage(
         GetReplicaConnection(replica_name), "",
@@ -133,21 +139,21 @@ void ReplicationManager::ReplicaHeartbeat(const std::string &replica_name) {
           auto epoch_now = std::chrono::system_clock::now().time_since_epoch();
           auto epoch_now_ms = std::chrono::duration_cast<std::chrono::milliseconds>(epoch_now);
           replica.last_heartbeat_ = epoch_now_ms.count();
-          REPLICATION_LOG_INFO(fmt::format("Replica {}: last heartbeat {}, heartbeat {} OK.", replica_name,
-                                           replica.last_heartbeat_, epoch_now_ms.count()));
+          REPLICATION_LOG_TRACE(fmt::format("Replica {}: last heartbeat {}, heartbeat {} OK.", replica_name,
+                                            replica.last_heartbeat_, epoch_now_ms.count()));
         },
         static_cast<uint64_t>(MessageType::HEARTBEAT));
   } catch (const MessengerException &e) {
-    REPLICATION_LOG_INFO(
+    REPLICATION_LOG_TRACE(
         fmt::format("Replica {}: last heartbeat {}, heartbeat failed.", replica_name, replica.last_heartbeat_));
     auto epoch_now = std::chrono::system_clock::now().time_since_epoch();
     auto epoch_now_ms = std::chrono::duration_cast<std::chrono::milliseconds>(epoch_now);
     if (epoch_now_ms.count() - replica.last_heartbeat_ >= REPLICATION_CARDIAC_ARREST_MS) {
-      REPLICATION_LOG_INFO(fmt::format("Replica {}: last heartbeat {}, declared dead {}.", replica_name,
+      REPLICATION_LOG_WARN(fmt::format("Replica {}: last heartbeat {}, declared dead {}.", replica_name,
                                        replica.last_heartbeat_, epoch_now_ms.count()));
     }
   }
-  REPLICATION_LOG_INFO(fmt::format("Replica {}: heartbeat end.", replica_name));
+  REPLICATION_LOG_TRACE(fmt::format("Replica {}: heartbeat end.", replica_name));
 }
 
 common::ManagedPointer<noisepage::messenger::ConnectionId> ReplicationManager::GetReplicaConnection(

src/self_driving/model_server/model_server_manager.cpp

@@ -1,12 +1,14 @@
+#include "self_driving/model_server/model_server_manager.h"
+
 #include <sys/wait.h>
+
 #include <csignal>
 #include <thread>  // NOLINT
 
 #include "common/json.h"
 #include "loggers/model_server_logger.h"
 #include "messenger/connection_destination.h"
 #include "messenger/messenger.h"
-#include "self_driving/model_server/model_server_manager.h"
 
 namespace noisepage::modelserver {
 static constexpr const char *MODEL_CONN_ID_NAME = "model-server-conn";
@@ -50,8 +52,11 @@ ModelServerManager::ModelServerManager(const std::string &model_bin,
         }
       })) {
   // Model Initialization handling logic
-  auto msm_handler = [&](common::ManagedPointer<messenger::Messenger> messenger, std::string_view sender_id,
-                         std::string_view message, uint64_t recv_cb_id) {
+  auto msm_handler = [&](common::ManagedPointer<messenger::Messenger> messenger, const messenger::ZmqMessage &msg) {
+    uint64_t sender_id = msg.GetSourceCallbackId();
+    uint64_t recv_cb_id = msg.GetDestinationCallbackId();
+    std::string_view message = msg.GetMessage();
+
     // ModelServer connected
     MODEL_SERVER_LOG_TRACE("[PID={},SENDER_ID={}] Messenger RECV: {}, {}", ::getpid(), sender_id, message, recv_cb_id);
 
@@ -168,10 +173,10 @@ bool ModelServerManager::TrainWith(const std::vector<std::string> &methods, cons
   j["data"]["save_path"] = save_path;
 
   // Callback to notify the waiter for result, or failure to parse the result.
-  auto callback = [&, future](common::ManagedPointer<messenger::Messenger> messenger, std::string_view sender_id,
-                              std::string_view message, uint64_t recv_cb_id) {
-    MODEL_SERVER_LOG_INFO("Callback :recv_cb_id={}, message={}", recv_cb_id, message);
-    future->Done(message);
+  auto callback = [&, future](common::ManagedPointer<messenger::Messenger> messenger,
+                              const messenger::ZmqMessage &msg) {
+    MODEL_SERVER_LOG_INFO("Callback :recv_cb_id={}, message={}", msg.GetDestinationCallbackId(), msg.GetMessage());
+    future->Done(msg.GetMessage());
   };
 
   return SendMessage(j.dump(), callback);
@@ -189,10 +194,9 @@ std::pair<std::vector<std::vector<double>>, bool> ModelServerManager::DoInferenc
   ModelServerFuture<std::vector<std::vector<double>>> future;
 
   // Callback to notify waiter with result
-  auto callback = [&](common::ManagedPointer<messenger::Messenger> messenger, std::string_view sender_id,
-                      std::string_view message, uint64_t recv_cb_id) {
-    MODEL_SERVER_LOG_INFO("Callback :recv_cb_id={}, message={}", recv_cb_id, message);
-    future.Done(message);
+  auto callback = [&](common::ManagedPointer<messenger::Messenger> messenger, const messenger::ZmqMessage &msg) {
+    MODEL_SERVER_LOG_INFO("Callback :recv_cb_id={}, message={}", msg.GetDestinationCallbackId(), msg.GetMessage());
+    future.Done(msg.GetMessage());
   };
 
   // Fail to send the message