Add Active Messages support

cpp/include/ucxx/delayed_submission.h

@@ -7,20 +7,18 @@
 #include <functional>
 #include <memory>
 #include <mutex>
+#include <utility>
 #include <vector>
 
 #include <ucp/api/ucp.h>
 
-#include <ucxx/component.h>
 #include <ucxx/log.h>
 
 namespace ucxx {
 
 typedef std::function<void()> DelayedSubmissionCallbackType;
 
-typedef std::shared_ptr<DelayedSubmissionCallbackType> DelayedSubmissionCallbackPtrType;
-
-class DelayedSubmission : public Component {
+class DelayedSubmission {
  public:
   bool _send{false};       ///< Whether this is a send (`true`) operation or recv (`false`)
   void* _buffer{nullptr};  ///< Raw pointer to data buffer
@@ -60,7 +58,7 @@ class DelayedSubmission : public Component {
 
 class DelayedSubmissionCollection {
  private:
-  std::vector<DelayedSubmissionCallbackPtrType>
+  std::vector<std::pair<std::shared_ptr<Request>, DelayedSubmissionCallbackType>>
     _collection{};      ///< The collection of all known delayed submission operations.
   std::mutex _mutex{};  ///< Mutex to provide access to the collection.
 
@@ -95,10 +93,12 @@ class DelayedSubmissionCollection {
    * Register a request for delayed submission with a callback that will be executed when
    * the request is in fact submitted when `process()` is called.
    *
+   * @param[in] request   the request to which the callback belongs, ensuring it remains
+   *                      alive until the callback is invoked.
    * @param[in] callback  the callback that will be executed by `process()` when the
    *                      operation is submitted.
    */
-  void registerRequest(DelayedSubmissionCallbackType callback);
+  void registerRequest(std::shared_ptr<Request> request, DelayedSubmissionCallbackType callback);
 };
 
 }  // namespace ucxx

cpp/include/ucxx/worker.h

@@ -386,10 +386,13 @@ class Worker : public Component {
    * thread, thus decreasing computation on the caller thread, but potentially increasing
    * transfer latency.
    *
+   * @param[in] request  the request to which the callback belongs, ensuring it remains
+   *                     alive until the callback is invoked.
    * @param[in] callback the callback set to execute the UCP transfer routine during the
    *                     worker thread loop.
    */
-  void registerDelayedSubmission(DelayedSubmissionCallbackType callback);
+  void registerDelayedSubmission(std::shared_ptr<Request> request,
+                                 DelayedSubmissionCallbackType callback);
 
   /**
    * @brief Inquire if worker has been created with delayed submission enabled.

cpp/src/delayed_submission.cpp

@@ -35,26 +35,25 @@ void DelayedSubmissionCollection::process()
       toProcess = std::move(_collection);
     }
 
-    for (auto& callbackPtr : toProcess) {
-      auto& callback = *callbackPtr;
+    for (auto& pair : toProcess) {
+      auto& req      = pair.first;
+      auto& callback = pair.second;
 
-      ucxx_trace_req("Submitting request: %p", callback.target<void (*)(std::shared_ptr<void>)>());
+      ucxx_trace_req("Submitting request: %p", req.get());
 
       if (callback) callback();
     }
   }
 }
 
-void DelayedSubmissionCollection::registerRequest(DelayedSubmissionCallbackType callback)
+void DelayedSubmissionCollection::registerRequest(std::shared_ptr<Request> request,
+                                                  DelayedSubmissionCallbackType callback)
 {
-  auto r = std::make_shared<DelayedSubmissionCallbackType>(callback);
-
   {
     std::lock_guard<std::mutex> lock(_mutex);
-    _collection.push_back(r);
+    _collection.push_back(std::make_pair(request, callback));
   }
-  ucxx_trace_req("Registered submit request: %p",
-                 callback.target<void (*)(std::shared_ptr<void>)>());
+  ucxx_trace_req("Registered submit request: %p", request.get());
 }
 
 }  // namespace ucxx

cpp/src/internal/request_am.cpp

@@ -19,7 +19,6 @@ RecvAmMessage::RecvAmMessage(internal::AmData* amData,
 {
   _request->_delayedSubmission =
     std::make_shared<DelayedSubmission>(false, _buffer->data(), _buffer->getSize());
-  _request->_delayedSubmission->setParent(_request);
 }
 
 void RecvAmMessage::setUcpRequest(void* request) { _request->_request = request; }

cpp/src/request_am.cpp

@@ -28,7 +28,12 @@ std::shared_ptr<RequestAm> createRequestAmSend(
   auto req = std::shared_ptr<RequestAm>(new RequestAm(
     endpoint, buffer, length, memoryType, enablePythonFuture, callbackFunction, callbackData));
 
-  req->_delayedSubmission->setParent(req);
+  // A delayed notification request is not populated immediately, instead it is
+  // delayed to allow the worker progress thread to set its status, and more
+  // importantly the Python future later on, so that we don't need the GIL here.
+  req->_worker->registerDelayedSubmission(
+    req, std::bind(std::mem_fn(&Request::populateDelayedSubmission), req.get()));
+
   return req;
 }
 
@@ -61,12 +66,6 @@ RequestAm::RequestAm(std::shared_ptr<Endpoint> endpoint,
 {
   _callback     = callbackFunction;
   _callbackData = callbackData;
-
-  // A delayed notification request is not populated immediately, instead it is
-  // delayed to allow the worker progress thread to set its status, and more
-  // importantly the Python future later on, so that we don't need the GIL here.
-  _worker->registerDelayedSubmission(
-    std::bind(std::mem_fn(&Request::populateDelayedSubmission), this));
 }
 
 RequestAm::RequestAm(std::shared_ptr<Component> endpointOrWorker,
@@ -165,16 +164,24 @@ ucs_status_t RequestAm::recvCallback(void* arg,
     ucs_status_ptr_t status =
       ucp_am_recv_data_nbx(worker->getHandle(), data, buf->data(), length, &request_param);
 
-    ucxx_trace_req_f(
-      ownerString.c_str(),
-      status,
-      "amRecv rndv",
-      "ep %p, buffer %p, size %lu, future %p, future handle %p, populateDelayedSubmission",
-      ep,
-      buf->data(),
-      length,
-      req->_future.get(),
-      req->_future->getHandle());
+    if (req->_enablePythonFuture)
+      ucxx_trace_req_f(ownerString.c_str(),
+                       status,
+                       "amRecv rndv",
+                       "ep %p, buffer %p, size %lu, future %p, future handle %p, recvCallback",
+                       ep,
+                       buf->data(),
+                       length,
+                       req->_future.get(),
+                       req->_future->getHandle());
+    else
+      ucxx_trace_req_f(ownerString.c_str(),
+                       status,
+                       "amRecv rndv",
+                       "ep %p, buffer %p, size %lu, recvCallback",
+                       ep,
+                       buf->data(),
+                       length);
 
     if (req->isCompleted()) {
       // The request completed/errored immediately
@@ -198,16 +205,24 @@ ucs_status_t RequestAm::recvCallback(void* arg,
     std::shared_ptr<Buffer> buf = amData->_allocators.at(UCS_MEMORY_TYPE_HOST)(length);
     if (length > 0) memcpy(buf->data(), data, length);
 
-    ucxx_trace_req_f(
-      ownerString.c_str(),
-      nullptr,
-      "amRecv eager",
-      "ep: %p, buffer %p, size %lu, future %p, future handle %p, populateDelayedSubmission",
-      ep,
-      buf->data(),
-      length,
-      req->_future.get(),
-      req->_future->getHandle());
+    if (req->_enablePythonFuture)
+      ucxx_trace_req_f(ownerString.c_str(),
+                       nullptr,
+                       "amRecv eager",
+                       "ep: %p, buffer %p, size %lu, future %p, future handle %p, recvCallback",
+                       ep,
+                       buf->data(),
+                       length,
+                       req->_future.get(),
+                       req->_future->getHandle());
+    else
+      ucxx_trace_req_f(ownerString.c_str(),
+                       nullptr,
+                       "amRecv eager",
+                       "ep: %p, buffer %p, size %lu, recvCallback",
+                       ep,
+                       buf->data(),
+                       length);
 
     internal::RecvAmMessage recvAmMessage(amData, ep, req, buf);
     recvAmMessage.callback(nullptr, UCS_OK);

cpp/src/request_stream.cpp

@@ -23,13 +23,6 @@ RequestStream::RequestStream(std::shared_ptr<Endpoint> endpoint,
             enablePythonFuture),
     _length(length)
 {
-  auto worker = endpoint->getWorker();
-
-  // A delayed notification request is not populated immediately, instead it is
-  // delayed to allow the worker progress thread to set its status, and more
-  // importantly the Python future later on, so that we don't need the GIL here.
-  worker->registerDelayedSubmission(
-    std::bind(std::mem_fn(&Request::populateDelayedSubmission), this));
 }
 
 std::shared_ptr<RequestStream> createRequestStream(std::shared_ptr<Endpoint> endpoint,
@@ -41,7 +34,12 @@ std::shared_ptr<RequestStream> createRequestStream(std::shared_ptr<Endpoint> end
   auto req = std::shared_ptr<RequestStream>(
     new RequestStream(endpoint, send, buffer, length, enablePythonFuture));
 
-  req->_delayedSubmission->setParent(req);
+  // A delayed notification request is not populated immediately, instead it is
+  // delayed to allow the worker progress thread to set its status, and more
+  // importantly the Python future later on, so that we don't need the GIL here.
+  req->_worker->registerDelayedSubmission(
+    req, std::bind(std::mem_fn(&Request::populateDelayedSubmission), req.get()));
+
   return req;
 }
 

cpp/src/request_tag.cpp

@@ -31,7 +31,12 @@ std::shared_ptr<RequestTag> createRequestTag(std::shared_ptr<Component> endpoint
                                                         callbackFunction,
                                                         callbackData));
 
-  req->_delayedSubmission->setParent(req);
+  // A delayed notification request is not populated immediately, instead it is
+  // delayed to allow the worker progress thread to set its status, and more
+  // importantly the Python future later on, so that we don't need the GIL here.
+  req->_worker->registerDelayedSubmission(
+    req, std::bind(std::mem_fn(&Request::populateDelayedSubmission), req.get()));
+
   return req;
 }
 
@@ -53,12 +58,6 @@ RequestTag::RequestTag(std::shared_ptr<Component> endpointOrWorker,
     throw ucxx::Error("An endpoint is required to send tag messages");
   _callback     = callbackFunction;
   _callbackData = callbackData;
-
-  // A delayed notification request is not populated immediately, instead it is
-  // delayed to allow the worker progress thread to set its status, and more
-  // importantly the Python future later on, so that we don't need the GIL here.
-  _worker->registerDelayedSubmission(
-    std::bind(std::mem_fn(&Request::populateDelayedSubmission), this));
 }
 
 void RequestTag::callback(void* request, ucs_status_t status, const ucp_tag_recv_info_t* info)

cpp/src/worker.cpp

@@ -261,12 +261,13 @@ bool Worker::progress()
   return ret;
 }
 
-void Worker::registerDelayedSubmission(DelayedSubmissionCallbackType callback)
+void Worker::registerDelayedSubmission(std::shared_ptr<Request> request,
+                                       DelayedSubmissionCallbackType callback)
 {
   if (_delayedSubmissionCollection == nullptr) {
     callback();
   } else {
-    _delayedSubmissionCollection->registerRequest(callback);
+    _delayedSubmissionCollection->registerRequest(request, callback);
 
     /* Waking the progress event is needed here because the UCX request is
      * not dispatched immediately. Thus we must signal the progress task so