Refactor callback logic to use shared_ptr rather than tokens and add …

…broadcasting test

cpp/mrc/include/mrc/runtime/remote_descriptor.hpp

@@ -426,9 +426,6 @@ class Descriptor2 : public std::enable_shared_from_this<Descriptor2>
     template <typename T>
     memory::buffer serialize(std::shared_ptr<memory::memory_resource> mr);
 
-    template <typename T>
-    memory::buffer_view serialize(memory::buffer_view buffer);
-
     template <typename T>
     [[nodiscard]] const T deserialize() const;
 
@@ -442,6 +439,7 @@ class Descriptor2 : public std::enable_shared_from_this<Descriptor2>
         m_encoded_object(std::move(encoded_object)), m_data_plane_resources(data_plane_resources) {}
 
     void setup_remote_payloads();
+    void register_remote_descriptor();
 
     std::any m_value;
 
@@ -453,27 +451,16 @@ class Descriptor2 : public std::enable_shared_from_this<Descriptor2>
 template <typename T>
 memory::buffer Descriptor2::serialize(std::shared_ptr<memory::memory_resource> mr)
 {
-    m_encoded_object = std::move(mrc::codable::encode2<T>(std::any_cast<T>(m_value)));
-
-    setup_remote_payloads();
-
-    // Allocate enough bytes to hold the encoded object
-    auto buffer = memory::buffer(m_encoded_object->proto().ByteSizeLong(), mr);
-
-    if (!m_encoded_object->proto().SerializeToArray(buffer.data(), buffer.bytes()))
+    if (!m_encoded_object)
     {
-        LOG(FATAL) << "Failed to serialize EncodedObjectProto to bytes";
+        m_encoded_object = std::move(mrc::codable::encode2<T>(std::any_cast<T>(m_value)));
+        this->setup_remote_payloads();
     }
 
-    return buffer;
-}
-
-template <typename T>
-memory::buffer_view Descriptor2::serialize(memory::buffer_view buffer)
-{
-    m_encoded_object = std::move(mrc::codable::encode2<T>(std::any_cast<T>(m_value)));
+    this->register_remote_descriptor();
 
-    setup_remote_payloads();
+    // Allocate enough bytes to hold the encoded object
+    auto buffer = memory::buffer(m_encoded_object->proto().ByteSizeLong(), mr);
 
     if (!m_encoded_object->proto().SerializeToArray(buffer.data(), buffer.bytes()))
     {

cpp/mrc/src/internal/data_plane/data_plane_resources.cpp

@@ -161,22 +161,22 @@ DataPlaneResources2::DataPlaneResources2()
     DVLOG(10) << "initialize the registration cache for this context";
     m_registration_cache = std::make_shared<ucx::RegistrationCache2>(m_context);
 
-    auto decrement_callback = ucxx::AmReceiverCallbackType([this](std::shared_ptr<ucxx::Request> req) {
+    auto pull_complete_callback = ucxx::AmReceiverCallbackType([this](std::shared_ptr<ucxx::Request> req) {
         auto status = req->getStatus();
         if (status != UCS_OK)
         {
-            LOG(ERROR) << "Error calling decrement_callback, request failed with status " << status << "("
+            LOG(ERROR) << "Error calling pull_complete_callback, request failed with status " << status << "("
                        << ucs_status_string(status) << ")";
         }
 
-        auto* dec_message = reinterpret_cast<remote_descriptor::RemoteDescriptorDecrementMessage*>(
+        auto* message = reinterpret_cast<remote_descriptor::DescriptorPullCompletionMessage*>(
             req->getRecvBuffer()->data());
 
-        decrement_tokens(dec_message);
+        complete_remote_pull(message);
     });
     m_worker->registerAmReceiverCallback(
         ucxx::AmReceiverCallbackInfo(ucxx::AmReceiverCallbackOwnerType("MRC"), ucxx::AmReceiverCallbackIdType(0)),
-        decrement_callback);
+        pull_complete_callback);
 
     // flush any work that needs to be done by the workers
     this->flush();
@@ -386,44 +386,54 @@ uint64_t DataPlaneResources2::get_next_object_id()
     return m_next_object_id++;
 }
 
-uint64_t DataPlaneResources2::register_remote_decriptor(
-    std::shared_ptr<runtime::Descriptor2> remote_descriptor)
+uint64_t DataPlaneResources2::register_remote_decriptor(std::shared_ptr<runtime::Descriptor2> descriptor)
 {
-    auto object_id = get_next_object_id();
-    remote_descriptor->encoded_object().set_object_id(object_id);
+    // If the descriptor has an object_id > 0, the descriptor has already been registered and should not be re-registered
+    auto object_id = descriptor->encoded_object().object_id();
+    if (object_id > 0)
+    {
+        m_descriptor_by_id[object_id].push_back(descriptor);
+        return object_id;
+    }
+
+    object_id = get_next_object_id();
+    descriptor->encoded_object().set_object_id(object_id);
     {
         std::unique_lock lock(m_remote_descriptors_mutex);
         m_remote_descriptors_cv.wait(lock, [this] {
-            return m_remote_descriptor_by_id.size() < m_max_remote_descriptors;
+            return m_descriptor_by_id.size() < m_max_remote_descriptors;
         });
-        m_remote_descriptor_by_id[object_id] = remote_descriptor;
+        m_descriptor_by_id[object_id].push_back(descriptor);
     }
     return object_id;
 }
 
 uint64_t DataPlaneResources2::registered_remote_descriptor_count()
 {
-    return m_remote_descriptor_by_id.size();
+    return m_descriptor_by_id.size();
 }
 
-uint64_t DataPlaneResources2::registered_remote_descriptor_token_count(uint64_t object_id)
+uint64_t DataPlaneResources2::registered_remote_descriptor_ptr_count(uint64_t object_id)
 {
-    return m_remote_descriptor_by_id.at(object_id)->encoded_object().tokens();
+    return m_descriptor_by_id.at(object_id).size();
 }
 
-void DataPlaneResources2::decrement_tokens(remote_descriptor::RemoteDescriptorDecrementMessage* dec_message)
+void DataPlaneResources2::complete_remote_pull(remote_descriptor::DescriptorPullCompletionMessage* message)
 {
-    if (dec_message->tokens > 0)
+    // If the mapping between object_id to descriptor shared ptrs exists, then there exists >= 1 shared ptrs
+    if (m_descriptor_by_id.find(message->object_id) != m_descriptor_by_id.end())
     {
-        auto remote_descriptor = m_remote_descriptor_by_id[dec_message->object_id];
-        auto tokens            = remote_descriptor->encoded_object().tokens();
-        tokens -= dec_message->tokens;
-        remote_descriptor->encoded_object().set_tokens(tokens);
-        if (tokens == 0)
+        // Once we've completed pulling of a descriptor, we remove a descriptor shared ptr from the vector
+        // When the vector becomes empty, there will be no more shared ptrs pointing to the descriptor object,
+        // it will be destructed accordingly.
+        // We should also remove that mapping as the object_id corresponding to that mapping will not be reused.
+        auto& descriptors = m_descriptor_by_id[message->object_id];
+        descriptors.pop_back();
+        if (descriptors.size() == 0)
         {
             {
                 std::unique_lock lock(m_remote_descriptors_mutex);
-                m_remote_descriptor_by_id.erase(dec_message->object_id);
+                m_descriptor_by_id.erase(message->object_id);
             }
             m_remote_descriptors_cv.notify_one();
         }

cpp/mrc/src/internal/data_plane/data_plane_resources.hpp

@@ -196,9 +196,9 @@ class DataPlaneResources2
                                                  ucs_memory_type_t mem_type);
     std::shared_ptr<ucxx::Request> am_recv_async(std::shared_ptr<ucxx::Endpoint> endpoint);
 
-    uint64_t register_remote_decriptor(std::shared_ptr<runtime::Descriptor2> remote_descriptor);
+    uint64_t register_remote_decriptor(std::shared_ptr<runtime::Descriptor2> descriptor);
     uint64_t registered_remote_descriptor_count();
-    uint64_t registered_remote_descriptor_token_count(uint64_t object_id);
+    uint64_t registered_remote_descriptor_ptr_count(uint64_t object_id);
 
   private:
     std::optional<uint64_t> m_instance_id;  // Global ID used to identify this instance
@@ -212,22 +212,23 @@ class DataPlaneResources2
     std::map<std::string, std::shared_ptr<ucxx::Endpoint>> m_endpoints_by_address;
     std::map<uint64_t, std::shared_ptr<ucxx::Endpoint>> m_endpoints_by_id;
 
-    std::atomic<uint64_t> m_next_object_id{0};
+    // An object_id of 0 (default protobuf int field value) signifies an unregistered descriptor
+    std::atomic<uint64_t> m_next_object_id{1};
 
     // std::shared_ptr<node::Queue<std::unique_ptr<runtime::ValueDescriptor>>> m_outbound_descriptors;
     // std::map<InstanceID, std::weak_ptr<node::Queue<std::unique_ptr<runtime::ValueDescriptor>>>>
     // m_inbound_port_channels;
 
     uint64_t get_next_object_id();
 
-    void decrement_tokens(remote_descriptor::RemoteDescriptorDecrementMessage* dec_message);
+    void complete_remote_pull(remote_descriptor::DescriptorPullCompletionMessage* message);
 
     uint64_t m_max_remote_descriptors{std::numeric_limits<uint64_t>::max()};
     boost::fibers::mutex m_remote_descriptors_mutex{};
     boost::fibers::condition_variable m_remote_descriptors_cv{};
 
   protected:
-    std::map<uint64_t, std::shared_ptr<runtime::Descriptor2>> m_remote_descriptor_by_id;
+    std::map<uint64_t, std::vector<std::shared_ptr<runtime::Descriptor2>>> m_descriptor_by_id;
 };
 
 }  // namespace mrc::data_plane

cpp/mrc/src/internal/remote_descriptor/messages.hpp

@@ -21,6 +21,11 @@
 
 namespace mrc::remote_descriptor {
 
+struct DescriptorPullCompletionMessage
+{
+    std::uint64_t object_id;
+};
+
 struct RemoteDescriptorDecrementMessage
 {
     std::uint64_t object_id;

cpp/mrc/src/public/runtime/remote_descriptor.cpp

@@ -385,14 +385,13 @@ std::shared_ptr<Descriptor2> Descriptor2::create(memory::buffer_view view, data_
     data_plane_resources.wait_requests(requests);
 
     // For the remote descriptor message, send decrement to the remote resources
-    remote_descriptor::RemoteDescriptorDecrementMessage dec_message;
+    remote_descriptor::DescriptorPullCompletionMessage dec_message;
     dec_message.object_id = descriptor->proto().object_id();
-    dec_message.tokens    = descriptor->proto().tokens();
 
     // TODO(Peter): Define `ucxx::AmReceiverCallbackInfo` at central place, must be known by all MRC processes.
     // Send a decrement message using custom AM receiver callback
     auto decrement_request = ep->amSend(&dec_message,
-                                        sizeof(remote_descriptor::RemoteDescriptorDecrementMessage),
+                                        sizeof(remote_descriptor::DescriptorPullCompletionMessage),
                                         UCS_MEMORY_TYPE_HOST,
                                         ucxx::AmReceiverCallbackInfo("MRC", 0));
 
@@ -405,7 +404,6 @@ void Descriptor2::setup_remote_payloads()
 
     // Transfer the info object
     remote_object.set_instance_id(m_data_plane_resources.get_instance_id());
-    remote_object.set_tokens(std::numeric_limits<uint64_t>::max());
 
     // Loop over all local payloads and convert them to remote payloads
 
@@ -432,7 +430,10 @@ void Descriptor2::setup_remote_payloads()
         deferred_msg->set_memory_block_size(ucx_block->bytes());
         deferred_msg->set_remote_key(ucx_block->packed_remote_keys());
     }
+}
 
+void Descriptor2::register_remote_descriptor()
+{
     m_data_plane_resources.register_remote_decriptor(shared_from_this());
 }
 }  // namespace mrc::runtime