NIXL/UCX: Optimize request handling (#982)

Author: iyastreb

src/plugins/ucx/ucx_backend.cpp

@@ -24,6 +24,7 @@
 #include <optional>
 #include <limits>
 #include <future>
+#include <set>
 #include <string.h>
 #include <unistd.h>
 #include "absl/strings/numbers.h"
@@ -271,20 +272,12 @@ void nixlUcxEngine::vramFiniCtx()
 
 class nixlUcxIntReq {
 public:
-    std::unique_ptr<std::string> amBuffer;
-
-    bool
-    is_complete() const {
-        return completed_;
+    operator nixlUcxReq() noexcept {
+        return static_cast<nixlUcxReq>(this);
     }
 
     void
-    completed() {
-        completed_ = true;
-    }
-
-    void
-    setConnection(ucx_connection_ptr_t conn) {
+    setConnection(nixlUcxConnection *conn) {
         conn_ = conn;
     }
 
@@ -295,51 +288,28 @@ class nixlUcxIntReq {
     }
 
 private:
-    bool completed_ = false;
-    ucx_connection_ptr_t conn_;
+    nixlUcxConnection *conn_;
 };
 
-static void
-nixlUcxReqSetConnection(nixlUcxReq req, ucx_connection_ptr_t conn) {
-    nixlUcxIntReq *req_int = reinterpret_cast<nixlUcxIntReq *>(req);
-    req_int->setConnection(conn);
-}
-
-static void _internalRequestInit(void *request)
-{
-    /* Initialize request in-place (aka "placement new")*/
-    new(request) nixlUcxIntReq;
-}
-
-static void _internalRequestFini(void *request)
-{
-    /* Finalize request */
-    nixlUcxIntReq *req = (nixlUcxIntReq*)request;
-    req->~nixlUcxIntReq();
-}
-
-
-static void _internalRequestReset(nixlUcxIntReq *req) {
-    _internalRequestFini((void *)req);
-    _internalRequestInit((void *)req);
-}
-
 /****************************************
  * Backend request management
 *****************************************/
 
 class nixlUcxBackendH : public nixlBackendReqH {
 private:
+    std::set<ucx_connection_ptr_t> connections_;
     std::vector<nixlUcxIntReq *> requests_;
     nixlUcxWorker *worker;
     size_t worker_id;
 
     // Notification to be sent after completion of all requests
     struct Notif {
-	    std::string agent;
-	    nixl_blob_t payload;
-	    Notif(const std::string& remote_agent, const nixl_blob_t& msg)
-		    : agent(remote_agent), payload(msg) {}
+        std::string agent;
+        nixl_blob_t payload;
+
+        Notif(const std::string &remote_agent, const nixl_blob_t &msg)
+            : agent(remote_agent),
+              payload(msg) {}
     };
     std::optional<Notif> notif;
 
@@ -358,8 +328,11 @@ class nixlUcxBackendH : public nixlBackendReqH {
     }
 
     void
-    append(nixlUcxIntReq *req) {
-        requests_.push_back(req);
+    append(nixlUcxReq req, ucx_connection_ptr_t conn) {
+        auto req_int = static_cast<nixlUcxIntReq *>(req);
+        req_int->setConnection(conn.get());
+        requests_.push_back(req_int);
+        connections_.insert(conn);
     }
 
     virtual bool
@@ -371,15 +344,16 @@ class nixlUcxBackendH : public nixlBackendReqH {
     release() {
         // TODO: Error log: uncompleted requests found! Cancelling ...
         for (nixlUcxIntReq *req : requests_) {
-            if (!req->is_complete()) {
+            nixl_status_t ret = ucx_status_to_nixl(ucp_request_check_status(req));
+            if (ret == NIXL_IN_PROG) {
                 // TODO: Need process this properly.
                 // it may not be enough to cancel UCX request
-                worker->reqCancel((nixlUcxReq)req);
+                worker->reqCancel(req);
             }
-            _internalRequestReset(req);
-            worker->reqRelease((nixlUcxReq)req);
+            worker->reqRelease(req);
         }
         requests_.clear();
+        connections_.clear();
         return NIXL_SUCCESS;
     }
 
@@ -394,37 +368,37 @@ class nixlUcxBackendH : public nixlBackendReqH {
         while (worker->progress())
             ;
 
-        /* Go over all request updating their status */
-        nixl_status_t out_ret = NIXL_SUCCESS;
-        for (nixlUcxIntReq *req : requests_) {
-            nixl_status_t ret;
-            if (!req->is_complete()) {
-                ret = ucx_status_to_nixl(ucp_request_check_status((nixlUcxReq)req));
-                switch (ret) {
-                case NIXL_SUCCESS:
-                    /* Mark as completed */
-                    req->completed();
-                    break;
-                case NIXL_IN_PROG:
-                    out_ret = NIXL_IN_PROG;
-                    break;
-                default:
-                    // Any other ret value is ERR and will be returned
-                    nixl_status_t conn_status = req->checkConnection(worker_id);
-                    return (conn_status == NIXL_SUCCESS) ? ret : conn_status;
-                }
-            }
+        /* If last request is incomplete, return NIXL_IN_PROG early without
+         * checking other requests */
+        nixlUcxIntReq *req = requests_.back();
+        nixl_status_t ret = ucx_status_to_nixl(ucp_request_check_status(req));
+        if (ret == NIXL_IN_PROG) {
+            return NIXL_IN_PROG;
+        } else if (ret != NIXL_SUCCESS) {
+            nixl_status_t conn_status = req->checkConnection(worker_id);
+            return (conn_status == NIXL_SUCCESS) ? ret : conn_status;
         }
 
+        /* Last request completed successfully, all the others must be in the
+         * same state. TODO: remove extra checks? */
         size_t incomplete_reqs = 0;
+        nixl_status_t out_ret = NIXL_SUCCESS;
         for (nixlUcxIntReq *req : requests_) {
-            if (req->is_complete()) {
-                _internalRequestReset(req);
-                worker->reqRelease((nixlUcxReq)req);
-            } else {
+            nixl_status_t ret = ucx_status_to_nixl(ucp_request_check_status(req));
+            if (__builtin_expect(ret == NIXL_SUCCESS, 0)) {
+                worker->reqRelease(req);
+            } else if (ret == NIXL_IN_PROG) {
+                if (out_ret == NIXL_SUCCESS) {
+                    out_ret = NIXL_IN_PROG;
+                }
                 requests_[incomplete_reqs++] = req;
+            } else {
+                // Any other ret value is ERR and will be returned
+                nixl_status_t conn_status = req->checkConnection(worker_id);
+                out_ret = (conn_status == NIXL_SUCCESS) ? ret : conn_status;
             }
         }
+
         requests_.resize(incomplete_reqs);
         return out_ret;
     }
@@ -1127,13 +1101,8 @@ nixlUcxEngine::nixlUcxEngine(const nixlBackendInitParams &init_params)
         err_handling_mode = ucx_err_mode_from_string(err_handling_mode_it->second);
     }
 
-    uc = std::make_unique<nixlUcxContext>(devs,
-                                          sizeof(nixlUcxIntReq),
-                                          _internalRequestInit,
-                                          _internalRequestFini,
-                                          init_params.enableProgTh,
-                                          num_workers,
-                                          init_params.syncMode);
+    uc = std::make_unique<nixlUcxContext>(
+        devs, sizeof(nixlUcxIntReq), init_params.enableProgTh, num_workers, init_params.syncMode);
 
     for (size_t i = 0; i < num_workers; i++) {
         uws.emplace_back(std::make_unique<nixlUcxWorker>(*uc, err_handling_mode));
@@ -1360,10 +1329,7 @@ _retHelper(nixl_status_t ret, nixlUcxBackendH *hndl, nixlUcxReq &req, ucx_connec
     /* if transfer wasn't immediately completed */
     switch(ret) {
     case NIXL_IN_PROG:
-        // TODO: this cast does not look safe
-        // We need to allocate a vector of nixlUcxIntReq and set nixlUcxReqt
-        hndl->append((nixlUcxIntReq *)req);
-        nixlUcxReqSetConnection(req, conn);
+        hndl->append(req, conn);
     case NIXL_SUCCESS:
         // Nothing to do
         break;
@@ -1587,8 +1553,10 @@ nixlUcxEngine::postXfer(const nixl_xfer_op_t &operation,
         if (ret == NIXL_SUCCESS) {
             nixlUcxReq req;
             auto rmd = (nixlUcxPublicMetadata *)remote[0].metadataP;
-            ret = notifSendPriv(
-                remote_agent, opt_args->notifMsg, req, rmd->conn->getEp(int_handle->getWorkerId()));
+            ret = notifSendPriv(remote_agent,
+                                opt_args->notifMsg,
+                                rmd->conn->getEp(int_handle->getWorkerId()),
+                                &req);
             if (_retHelper(ret, int_handle, req, rmd->conn)) {
                 return ret;
             }
@@ -1624,7 +1592,7 @@ nixl_status_t nixlUcxEngine::checkXfer (nixlBackendReqH* handle) const
 
     nixlUcxReq req;
     nixl_status_t status =
-        notifSendPriv(notif->agent, notif->payload, req, conn->getEp(intHandle->getWorkerId()));
+        notifSendPriv(notif->agent, notif->payload, conn->getEp(intHandle->getWorkerId()), &req);
     notif.reset();
     status = _retHelper(status, intHandle, req, conn);
     if (status != NIXL_SUCCESS) {
@@ -1759,23 +1727,31 @@ int nixlUcxEngine::progress() {
 nixl_status_t
 nixlUcxEngine::notifSendPriv(const std::string &remote_agent,
                              const std::string &msg,
-                             nixlUcxReq &req,
-                             const std::unique_ptr<nixlUcxEp> &ep) const {
+                             const std::unique_ptr<nixlUcxEp> &ep,
+                             nixlUcxReq *req) const {
     nixlSerDes ser_des;
-    nixl_status_t ret;
 
     ser_des.addStr("name", localAgent);
     ser_des.addStr("msg", msg);
     // TODO: replace with mpool for performance
 
-    auto buffer = std::make_unique<std::string>(ser_des.exportStr());
-    ret = ep->sendAm(
-        NOTIF_STR, NULL, 0, (void *)buffer->data(), buffer->size(), UCP_AM_SEND_FLAG_EAGER, req);
-    if (ret == NIXL_IN_PROG) {
-        nixlUcxIntReq* nReq = (nixlUcxIntReq*)req;
-        nReq->amBuffer = std::move(buffer);
-    }
-    return ret;
+    std::string *buffer = new std::string(ser_des.exportStr());
+    auto deleter = [buffer, req](void *completed_request, void *ptr) {
+        delete buffer;
+        if ((req == nullptr) && (completed_request != nullptr)) {
+            /* Caller is not interested in the request, free it */
+            ucp_request_free(completed_request);
+        }
+    };
+
+    return ep->sendAm(NOTIF_STR,
+                      nullptr,
+                      0,
+                      (void *)buffer->data(),
+                      buffer->size(),
+                      UCP_AM_SEND_FLAG_EAGER,
+                      req,
+                      deleter);
 }
 
 ucx_connection_ptr_t
@@ -1827,26 +1803,16 @@ nixl_status_t nixlUcxEngine::getNotifs(notif_list_t &notif_list)
     return NIXL_SUCCESS;
 }
 
-nixl_status_t nixlUcxEngine::genNotif(const std::string &remote_agent, const std::string &msg) const
-{
-    nixl_status_t ret;
-    nixlUcxReq req;
-
+nixl_status_t
+nixlUcxEngine::genNotif(const std::string &remote_agent, const std::string &msg) const {
     auto conn = getConnection(remote_agent);
     if (!conn) {
         return NIXL_ERR_NOT_FOUND;
     }
 
-    ret = notifSendPriv(remote_agent, msg, req, conn->getEp(getWorkerId()));
-    switch(ret) {
-    case NIXL_IN_PROG:
-        /* do not track the request */
-        getWorker(getWorkerId())->reqRelease(req);
-    case NIXL_SUCCESS:
-        break;
-    default:
-        /* error case */
-        return ret;
+    nixl_status_t ret = notifSendPriv(remote_agent, msg, conn->getEp(getWorkerId()));
+    if (ret == NIXL_IN_PROG) {
+        ret = NIXL_SUCCESS;
     }
-    return NIXL_SUCCESS;
+    return ret;
 }

src/plugins/ucx/ucx_backend.h

@@ -288,8 +288,8 @@ class nixlUcxEngine : public nixlBackendEngine {
     nixl_status_t
     notifSendPriv(const std::string &remote_agent,
                   const std::string &msg,
-                  nixlUcxReq &req,
-                  const std::unique_ptr<nixlUcxEp> &ep) const;
+                  const std::unique_ptr<nixlUcxEp> &ep,
+                  nixlUcxReq *req = nullptr) const;
 
     ucx_connection_ptr_t
     getConnection(const std::string &remote_agent) const;