Merge remote-tracking branch 'upstream/branch-0.36' into tag-mask

rapidsai · Dec 7, 2023 · 43acf5b · 43acf5b
2 parents c8c656f + 13a663b
commit 43acf5b
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 19 deletions.
diff --git a/cpp/include/ucxx/request_tag_multi.h b/cpp/include/ucxx/request_tag_multi.h
@@ -43,6 +43,8 @@ class RequestTagMulti : public Request {
   std::mutex
     _completedRequestsMutex{};   ///< Mutex to control access to completed requests container
   size_t _completedRequests{0};  ///< Count requests that already completed
+  ucs_status_t _finalStatus{
+    UCS_OK};  ///< Shortcut to the final status, a.k.a. the first error to occur
 
  public:
   std::vector<BufferRequestPtr> _bufferRequests{};  ///< Container of all requests posted

diff --git a/cpp/src/request_tag_multi.cpp b/cpp/src/request_tag_multi.cpp
@@ -179,28 +179,27 @@ void RequestTagMulti::markCompleted(ucs_status_t status, RequestCallbackUserData
                  _delayedSubmission->_data.getTag()._tag);
   std::lock_guard<std::mutex> lock(_completedRequestsMutex);
 
+  if (_finalStatus == UCS_OK && status != UCS_OK) _finalStatus = status;
+
   if (++_completedRequests == _totalFrames) {
-    auto s = UCS_OK;
-
-    // Get the first non-UCS_OK status and set that as complete status
-    for (const auto& br : _bufferRequests) {
-      if (br->request) {
-        s = br->request->getStatus();
-        if (s != UCS_OK) break;
-      }
-    }
-    // Check the status of the current (and final) message as it may have completed before
-    // `_bufferRequests->request` was populated.
-    if (s == UCS_OK) s = status;
+    setStatus(_finalStatus);
 
-    setStatus(s);
+    ucxx_trace_req(
+      "RequestTagMulti::markCompleted request: %p, tag: %lx, completed: %lu/%lu, final status: %d "
+      "(%s)",
+      this,
+      _delayedSubmission->_data.getTag()._tag,
+      _completedRequests,
+      _totalFrames,
+      _finalStatus,
+      ucs_status_string(_finalStatus));
+  } else {
+    ucxx_trace_req("RequestTagMulti::markCompleted request: %p, tag: %lx, completed: %lu/%lu",
+                   this,
+                   _delayedSubmission->_data.getTag()._tag,
+                   _completedRequests,
+                   _totalFrames);
   }
-
-  ucxx_trace_req("RequestTagMulti::markCompleted request: %p, tag: %lx, completed: %lu/%lu",
-                 this,
-                 _delayedSubmission->_data.getTag()._tag,
-                 _completedRequests,
-                 _totalFrames);
 }
 
 void RequestTagMulti::recvHeader()