UCP/AM: CR fixes openucx#2

amastbaum · Jun 17, 2024 · 0da08c2 · 0da08c2
1 parent cc90dea
commit 0da08c2
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 43 deletions.
diff --git a/src/ucp/core/ucp_am.c b/src/ucp/core/ucp_am.c
@@ -329,8 +329,8 @@ ucs_status_t ucp_worker_set_am_recv_handler(ucp_worker_h worker,
     }
 
     if (param->id > UINT16_MAX) {
-        ucs_error("AM handler id is too large (%u). "
-                  "Should not be greater than uint16_t", param->id);
+        ucs_error("AM handler id %u is outside of allowed range [0, %u]",
+                  param->id, UINT16_MAX);
         return UCS_ERR_INVALID_PARAM;
     }
 
@@ -959,8 +959,8 @@ UCS_PROFILE_FUNC(ucs_status_ptr_t, ucp_am_send_nbx,
 
     if (id > UINT16_MAX) {
         ret = UCS_STATUS_PTR(UCS_ERR_INVALID_PARAM);
-        ucs_error("AM handler id is too large (%u). "
-                  "Should not be greater than uint16_t", id);
+        ucs_error("AM handler id %u is outside of allowed range [0, %u]",
+                  id, UINT16_MAX);
         goto out;
     }
 

diff --git a/test/gtest/ucp/test_ucp_am.cc b/test/gtest/ucp/test_ucp_am.cc
@@ -709,8 +709,7 @@ class test_ucp_am_id : public test_ucp_am_nbx {
     }
 
     ucs_status_t set_am_id_overflow_data_handler(entity &e, unsigned am_id,
-                               ucp_am_recv_callback_t cb, void *arg,
-                               unsigned flags = 0)
+                               ucp_am_recv_callback_t cb)
     {
         ucp_am_handler_param_t param;
 
@@ -720,66 +719,37 @@ class test_ucp_am_id : public test_ucp_am_nbx {
                            UCP_AM_HANDLER_PARAM_FIELD_ARG;
         param.id         = am_id;
         param.cb         = cb;
-        param.arg        = arg;
-
-        if (flags != 0) {
-            param.field_mask |= UCP_AM_HANDLER_PARAM_FIELD_FLAGS;
-            param.flags       = flags;
-        }
+        param.arg        = this;
 
         return ucp_worker_set_am_recv_handler(e.worker(), &param);
     }
 
     void test_am_id_handler(size_t size, size_t header_size = 0ul)
     {
-        mem_buffer sbuf(size, tx_memtype());
-        sbuf.pattern_fill(SEED);
-        m_hdr.resize(header_size);
-        ucs::fill_random(m_hdr);
         reset_counters();
-        ucp_mem_h memh = NULL;
 
         EXPECT_EQ(UCS_OK, set_am_id_overflow_data_handler(receiver(),
-                                        0x1, am_id_overflow_data_cb_1,
-                                        this, 0));
+                                        0x1, am_id_overflow_data_cb_1));
         EXPECT_NE(UCS_OK, set_am_id_overflow_data_handler(receiver(),
-                                        0xffff0001, am_id_overflow_data_cb_2,
-                                        this, 0));
-
-        ucp::data_type_desc_t sdt_desc(m_dt, sbuf.ptr(), size);
-
-        if (prereg()) {
-            memh = sender().mem_map(sbuf.ptr(), size);
-        }
+                                        0xffff0001, am_id_overflow_data_cb_2));
 
         ucp_request_param_t param;
-        param.op_attr_mask = UCP_OP_ATTR_FIELD_DATATYPE;
-        param.datatype     = sdt_desc.dt();
-
-        if (memh != NULL) {
-            param.op_attr_mask |= UCP_OP_ATTR_FIELD_MEMH;
-            param.memh          = memh;
-        }
+        param.op_attr_mask = 0;
 
         m_send_counter++;
-        ucs_status_ptr_t sptr = ucp_am_send_nbx(sender().ep(), 0x1, m_hdr.data(),
-                                                m_hdr.size(), sbuf.ptr(), size,
-                                                &param);
+        ucs_status_ptr_t sptr = ucp_am_send_nbx(sender().ep(), 0x1, NULL,
+                                                0ul, NULL, 0, &param);
         wait_receives();
         EXPECT_EQ(UCS_OK, request_wait(sptr));
         EXPECT_EQ(m_recv_counter_cb_1, 1);
         EXPECT_EQ(m_recv_counter_cb_2, 0);
 
         m_send_counter++;
-        sptr = ucp_am_send_nbx(sender().ep(), 0xffff0001, m_hdr.data(),
-                               m_hdr.size(), sbuf.ptr(), size, &param);
+        sptr = ucp_am_send_nbx(sender().ep(), 0xffff0001,
+                               NULL, 0ul, NULL, 0, &param);
         EXPECT_EQ(UCS_PTR_STATUS(sptr), UCS_ERR_INVALID_PARAM);
         EXPECT_EQ(m_recv_counter_cb_1, 1);
         EXPECT_EQ(m_recv_counter_cb_2, 0);
-
-        if (prereg()) {
-            sender().mem_unmap(memh);
-        }
     }
 
     virtual ucs_status_t am_id_overflow_data_handler(