Improve tagProbe and AM receiver callback

cpp/include/ucxx/typedefs.h

@@ -72,6 +72,20 @@ enum TagMask : ucp_tag_t {};
  */
 static constexpr TagMask TagMaskFull{std::numeric_limits<std::underlying_type_t<TagMask>>::max()};
 
+/**
+ * @brief Information about probed tag message.
+ *
+ * Contains information returned when probing by a tag message received by the worker but
+ * not yet consumed.
+ */
+class TagRecvInfo {
+ public:
+  Tag senderTag;  ///< Sender tag
+  size_t length;  ///< The size of the received data
+
+  explicit TagRecvInfo(const ucp_tag_recv_info_t&);
+};
+
 /**
  * @brief A UCP configuration map.
  *
@@ -124,9 +138,10 @@ typedef std::function<std::shared_ptr<Buffer>(size_t)> AmAllocatorType;
  * @brief Active Message receiver callback.
  *
  * Type for a custom Active Message receiver callback, executed by the remote worker upon
- * Active Message request completion.
+ * Active Message request completion. The first parameter is the request that completed,
+ * the second is the handle of the UCX endpoint of the sender.
  */
-typedef std::function<void(std::shared_ptr<Request>)> AmReceiverCallbackType;
+typedef std::function<void(std::shared_ptr<Request>, ucp_ep_h)> AmReceiverCallbackType;
 
 /**
  * @brief Active Message receiver callback owner name.

cpp/include/ucxx/worker.h

@@ -10,6 +10,7 @@
 #include <queue>
 #include <string>
 #include <thread>
+#include <utility>
 
 #include <ucp/api/ucp.h>
 
@@ -684,21 +685,29 @@ class Worker : public Component {
    *
    * Checks the worker for any uncaught tag messages. An uncaught tag message is any
    * tag message that has been fully or partially received by the worker, but not matched
-   * by a corresponding `ucp_tag_recv_*` call.
+   * by a corresponding `ucp_tag_recv_*` call. Additionally, returns information about the
+   * tag message.
    *
    * @code{.cpp}
    * // `worker` is `std::shared_ptr<ucxx::Worker>`
-   * assert(!worker->tagProbe(0));
+   * auto probe = worker->tagProbe(0);
+   * assert(!probe.first)
    *
    * // `ep` is a remote `std::shared_ptr<ucxx::Endpoint` to the local `worker`
    * ep->tagSend(buffer, length, 0);
    *
-   * assert(worker->tagProbe(0));
+   * probe = worker->tagProbe(0);
+   * assert(probe.first);
+   * assert(probe.second.tag == 0);
+   * assert(probe.second.length == length);
    * @endcode
    *
-   * @returns `true` if any uncaught messages were received, `false` otherwise.
+   * @returns pair where first elements is `true` if any uncaught messages were received,
+   *          `false` otherwise, and second element contain the information from the tag
+   *          receive.
    */
-  [[nodiscard]] bool tagProbe(const Tag tag);
+  [[nodiscard]] std::pair<bool, TagRecvInfo> tagProbe(const Tag tag,
+                                                      const TagMask tagMask = TagMaskFull);
 
   /**
    * @brief Enqueue a tag receive operation.

cpp/src/internal/request_am.cpp

@@ -29,7 +29,7 @@ RecvAmMessage::RecvAmMessage(internal::AmData* amData,
 
   if (receiverCallback) {
     _request->_callback = [this, receiverCallback](ucs_status_t, std::shared_ptr<void>) {
-      receiverCallback(_request);
+      receiverCallback(_request, _ep);
     };
   }
 }

cpp/src/worker.cpp

@@ -576,7 +576,12 @@ void Worker::removeInflightRequest(const Request* const request)
   }
 }
 
-bool Worker::tagProbe(const Tag tag)
+TagRecvInfo::TagRecvInfo(const ucp_tag_recv_info_t& info)
+  : senderTag(Tag(info.sender_tag)), length(info.length)
+{
+}
+
+std::pair<bool, TagRecvInfo> Worker::tagProbe(const Tag tag, const TagMask tagMask)
 {
   if (!isProgressThreadRunning()) {
     progress();
@@ -592,9 +597,9 @@ bool Worker::tagProbe(const Tag tag)
   }
 
   ucp_tag_recv_info_t info;
-  ucp_tag_message_h tag_message = ucp_tag_probe_nb(_handle, tag, TagMaskFull, 0, &info);
+  ucp_tag_message_h tag_message = ucp_tag_probe_nb(_handle, tag, tagMask, 0, &info);
 
-  return tag_message != NULL;
+  return {tag_message != NULL, TagRecvInfo(info)};
 }
 
 std::shared_ptr<Request> Worker::tagRecv(void* buffer,

cpp/tests/request.cpp

@@ -59,17 +59,18 @@ class RequestTest : public ::testing::TestWithParam<
 
   void SetUp()
   {
+    std::tie(_bufferType,
+             _registerCustomAmAllocator,
+             _enableDelayedSubmission,
+             _progressMode,
+             _messageLength) = GetParam();
+
     if (_bufferType == ucxx::BufferType::RMM) {
 #if !UCXX_ENABLE_RMM
       GTEST_SKIP() << "UCXX was not built with RMM support";
 #endif
     }
 
-    std::tie(_bufferType,
-             _registerCustomAmAllocator,
-             _enableDelayedSubmission,
-             _progressMode,
-             _messageLength) = GetParam();
     _memoryType =
       (_bufferType == ucxx::BufferType::RMM) ? UCS_MEMORY_TYPE_CUDA : UCS_MEMORY_TYPE_HOST;
     _messageSize = _messageLength * sizeof(int);
@@ -168,13 +169,14 @@ TEST_P(RequestTest, ProgressAm)
     GTEST_SKIP() << "Interrupting UCP worker progress operation in wait mode is not possible";
   }
 
+  if (_registerCustomAmAllocator && _memoryType == UCS_MEMORY_TYPE_CUDA) {
 #if !UCXX_ENABLE_RMM
-  GTEST_SKIP() << "UCXX was not built with RMM support";
+    GTEST_SKIP() << "UCXX was not built with RMM support";
 #else
-  if (_registerCustomAmAllocator && _memoryType == UCS_MEMORY_TYPE_CUDA) {
     _worker->registerAmAllocator(UCS_MEMORY_TYPE_CUDA, [](size_t length) {
       return std::make_shared<ucxx::RMMBuffer>(length);
     });
+#endif
   }
 
   allocate(1, false);
@@ -198,7 +200,6 @@ TEST_P(RequestTest, ProgressAm)
 
   // Assert data correctness
   ASSERT_THAT(_recv[0], ContainerEq(_send[0]));
-#endif
 }
 
 TEST_P(RequestTest, ProgressAmReceiverCallback)
@@ -207,13 +208,14 @@ TEST_P(RequestTest, ProgressAmReceiverCallback)
     GTEST_SKIP() << "Interrupting UCP worker progress operation in wait mode is not possible";
   }
 
+  if (_registerCustomAmAllocator && _memoryType == UCS_MEMORY_TYPE_CUDA) {
 #if !UCXX_ENABLE_RMM
-  GTEST_SKIP() << "UCXX was not built with RMM support";
+    GTEST_SKIP() << "UCXX was not built with RMM support";
 #else
-  if (_registerCustomAmAllocator && _memoryType == UCS_MEMORY_TYPE_CUDA) {
     _worker->registerAmAllocator(UCS_MEMORY_TYPE_CUDA, [](size_t length) {
       return std::make_shared<ucxx::RMMBuffer>(length);
     });
+#endif
   }
 
   // Define AM receiver callback's owner and id for callback
@@ -226,7 +228,7 @@ TEST_P(RequestTest, ProgressAmReceiverCallback)
   // Define AM receiver callback and register with worker
   std::vector<std::shared_ptr<ucxx::Request>> receivedRequests;
   auto callback = ucxx::AmReceiverCallbackType(
-    [this, &receivedRequests, &mutex](std::shared_ptr<ucxx::Request> req) {
+    [this, &receivedRequests, &mutex](std::shared_ptr<ucxx::Request> req, ucp_ep_h) {
       {
         std::lock_guard<std::mutex> lock(mutex);
         receivedRequests.push_back(req);
@@ -260,7 +262,6 @@ TEST_P(RequestTest, ProgressAmReceiverCallback)
 
   // Assert data correctness
   ASSERT_THAT(_recv[0], ContainerEq(_send[0]));
-#endif
 }
 
 TEST_P(RequestTest, ProgressStream)

cpp/tests/worker.cpp

@@ -110,7 +110,8 @@ TEST_F(WorkerTest, TagProbe)
   auto progressWorker = getProgressFunction(_worker, ProgressMode::Polling);
   auto ep             = _worker->createEndpointFromWorkerAddress(_worker->getAddress());
 
-  ASSERT_FALSE(_worker->tagProbe(ucxx::Tag{0}));
+  auto probed = _worker->tagProbe(ucxx::Tag{0});
+  ASSERT_FALSE(probed.first);
 
   std::vector<int> buf{123};
   std::vector<std::shared_ptr<ucxx::Request>> requests;
@@ -119,10 +120,14 @@ TEST_F(WorkerTest, TagProbe)
 
   loopWithTimeout(std::chrono::milliseconds(5000), [this, progressWorker]() {
     progressWorker();
-    return _worker->tagProbe(ucxx::Tag{0});
+    auto probed = _worker->tagProbe(ucxx::Tag{0});
+    return probed.first;
   });
 
-  ASSERT_TRUE(_worker->tagProbe(ucxx::Tag{0}));
+  probed = _worker->tagProbe(ucxx::Tag{0});
+  ASSERT_TRUE(probed.first);
+  ASSERT_EQ(probed.second.senderTag, ucxx::Tag{0});
+  ASSERT_EQ(probed.second.length, buf.size() * sizeof(int));
 }
 
 TEST_F(WorkerTest, AmProbe)
@@ -189,7 +194,7 @@ TEST_P(WorkerProgressTest, ProgressAmReceiverCallback)
   // Define AM receiver callback and register with worker
   std::vector<std::shared_ptr<ucxx::Request>> receivedRequests;
   auto callback = ucxx::AmReceiverCallbackType(
-    [this, &receivedRequests, &mutex](std::shared_ptr<ucxx::Request> req) {
+    [this, &receivedRequests, &mutex](std::shared_ptr<ucxx::Request> req, ucp_ep_h) {
       {
         std::lock_guard<std::mutex> lock(mutex);
         receivedRequests.push_back(req);

python/ucxx/ucxx/_lib/libucxx.pyx

@@ -25,6 +25,7 @@ from libcpp.memory cimport (
     unique_ptr,
 )
 from libcpp.optional cimport nullopt
+from libcpp.pair cimport pair
 from libcpp.string cimport string
 from libcpp.utility cimport move
 from libcpp.vector cimport vector
@@ -684,13 +685,25 @@ cdef class UCXWorker():
 
         return num_canceled
 
-    def tag_probe(self, UCXXTag tag) -> bool:
-        cdef bint tag_matched
+    def tag_probe(self, UCXXTag tag, UCXXTagMask tag_mask = UCXXTagMaskFull) -> bool:
         cdef Tag cpp_tag = <Tag><size_t>tag.value
+        cdef TagMask cpp_tag_mask = <TagMask><size_t>tag_mask.value
+        cdef ucp_tag_recv_info_t empty_tag_recv_info
+        cdef pair[bint, TagRecvInfo]* probed
+        cdef bint tag_matched = False
 
         with nogil:
-            tag_matched = self._worker.get().tagProbe(cpp_tag)
+            # TagRecvInfo is not default-construtible, therefore we need to use a
+            # pointer, allocating it using a temporary ucp_tag_recv_info_t object
+            probed = new pair[bint, TagRecvInfo](
+                False,
+                TagRecvInfo(empty_tag_recv_info)
+            )
+            probed[0] = self._worker.get().tagProbe(cpp_tag, cpp_tag_mask)
+            tag_matched = probed[0].first
+            del probed
 
+        # TODO: Come up with good interface to expose TagRecvInfo as well
         return tag_matched
 
     def set_progress_thread_start_callback(

python/ucxx/ucxx/_lib/ucxx_api.pxd

@@ -9,6 +9,7 @@ from libcpp cimport bool as cpp_bool
 from libcpp.functional cimport function
 from libcpp.memory cimport shared_ptr, unique_ptr
 from libcpp.optional cimport nullopt_t, optional
+from libcpp.pair cimport pair
 from libcpp.string cimport string
 from libcpp.unordered_map cimport unordered_map as cpp_unordered_map
 from libcpp.vector cimport vector
@@ -54,6 +55,9 @@ cdef extern from "ucp/api/ucp.h" nogil:
 
     ctypedef uint64_t ucp_tag_t
 
+    ctypedef struct ucp_tag_recv_info_t:
+        pass
+
     ctypedef enum ucs_status_t:
         pass
 
@@ -174,10 +178,12 @@ cdef extern from "<ucxx/api.h>" namespace "ucxx" nogil:
         pass
     cdef enum TagMask:
         pass
+    cdef cppclass TagRecvInfo:
+        TagRecvInfo(const ucp_tag_recv_info_t&)
+        Tag senderTag
+        size_t length
     cdef cppclass AmReceiverCallbackInfo:
         pass
-    # ctypedef Tag CppTag
-    # ctypedef TagMask CppTagMask
 
     # Using function[Buffer] here doesn't seem possible due to Cython bugs/limitations.
     # The workaround is to use a raw C function pointer and let it be parsed by the
@@ -241,7 +247,7 @@ cdef extern from "<ucxx/api.h>" namespace "ucxx" nogil:
         size_t cancelInflightRequests(
             uint64_t period, uint64_t maxAttempts
         ) except +raise_py_error
-        bint tagProbe(const Tag) const
+        pair[bint, TagRecvInfo] tagProbe(const Tag, const TagMask) const
         void setProgressThreadStartCallback(
             function[void(void*)] callback, void* callbackArg
         )