Retain copy of header for AM send requests

Retain a copy of headers for AM send requests as workaround for possible UCX bug openucx/ucx#8639 .
rapidsai · Jan 16, 2025 · b5dd27b · b5dd27b
1 parent 93edd75
commit b5dd27b
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 8 deletions.
diff --git a/cpp/include/ucxx/request_am.h b/cpp/include/ucxx/request_am.h
@@ -31,6 +31,9 @@ class RequestAm : public Request {
  private:
   friend class internal::RecvAmMessage;
 
+  std::string _header{};  ///< Retain copy of header for send requests as workaround for
+                          ///< https://github.com/openucx/ucx/issues/8639
+
   /**
    * @brief Private constructor of `ucxx::RequestAm`.
    *

diff --git a/cpp/src/request_am.cpp b/cpp/src/request_am.cpp
@@ -389,18 +389,18 @@ void RequestAm::request()
         ucp_request_param_t param = {.op_attr_mask = UCP_OP_ATTR_FIELD_CALLBACK |
                                                      UCP_OP_ATTR_FIELD_FLAGS |
                                                      UCP_OP_ATTR_FIELD_USER_DATA,
-                                     .flags = UCP_AM_SEND_FLAG_REPLY | UCP_AM_SEND_FLAG_COPY_HEADER,
+                                     .flags     = UCP_AM_SEND_FLAG_REPLY,
                                      .datatype  = ucp_dt_make_contig(1),
                                      .user_data = this};
 
-        param.cb.send         = _amSendCallback;
-        AmHeader header       = {.memoryType           = amSend._memoryType,
-                                 .receiverCallbackInfo = amSend._receiverCallbackInfo};
-        auto headerSerialized = header.serialize();
-        void* request         = ucp_am_send_nbx(_endpoint->getHandle(),
+        param.cb.send   = _amSendCallback;
+        AmHeader header = {.memoryType           = amSend._memoryType,
+                           .receiverCallbackInfo = amSend._receiverCallbackInfo};
+        _header         = header.serialize();
+        void* request   = ucp_am_send_nbx(_endpoint->getHandle(),
                                         0,
-                                        headerSerialized.data(),
-                                        headerSerialized.size(),
+                                        _header.data(),
+                                        _header.size(),
                                         amSend._buffer,
                                         amSend._length,
                                         &param);