diff --git a/src/ucp/core/ucp_worker.c b/src/ucp/core/ucp_worker.c index 9f3591bb593..45ced6b3524 100644 --- a/src/ucp/core/ucp_worker.c +++ b/src/ucp/core/ucp_worker.c @@ -3111,10 +3111,10 @@ void ucp_worker_destroy(ucp_worker_h worker) ucs_free(worker); } -static ucs_status_t ucp_worker_address_pack(ucp_worker_h worker, - uint32_t address_flags, - size_t *address_length_p, - void **address_p) +static ucs_status_t +ucp_worker_address_pack(ucp_worker_h worker, uint32_t address_flags, + const char *address_device_name, + size_t *address_length_p, void **address_p) { ucp_context_h context = worker->context; unsigned flags = ucp_worker_default_address_pack_flags(worker); @@ -3128,16 +3128,29 @@ static ucs_status_t ucp_worker_address_pack(ucp_worker_h worker, */ ucs_assert(flags & UCP_ADDRESS_PACK_FLAG_WORKER_UUID); + if (address_device_name != NULL) { + ucp_context_dev_tl_bitmap(context, address_device_name, &tl_bitmap); + } else { + tl_bitmap = context->tl_bitmap; + } + if (address_flags & UCP_WORKER_ADDRESS_FLAG_NET_ONLY) { - UCS_STATIC_BITMAP_RESET_ALL(&tl_bitmap); - UCS_STATIC_BITMAP_FOR_EACH_BIT(tl_id, &worker->context->tl_bitmap) { + UCS_STATIC_BITMAP_FOR_EACH_BIT(tl_id, &tl_bitmap) { iface_attr = ucp_worker_iface_get_attr(worker, tl_id); - if (iface_attr->cap.flags & UCT_IFACE_FLAG_INTER_NODE) { - UCS_STATIC_BITMAP_SET(&tl_bitmap, tl_id); + if (!(iface_attr->cap.flags & UCT_IFACE_FLAG_INTER_NODE)) { + UCS_STATIC_BITMAP_RESET(&tl_bitmap, tl_id); } } - } else { - UCS_STATIC_BITMAP_SET_ALL(&tl_bitmap); + } + + if ((address_device_name != NULL) && UCS_STATIC_BITMAP_IS_ZERO(tl_bitmap)) { + ucs_error("worker %p: no addressable resources for device %s with " + "network-only flag %c", + worker, address_device_name, + ((address_flags & UCP_WORKER_ADDRESS_FLAG_NET_ONLY) != 0) ? + 'y' : + 'n'); + return UCS_ERR_NO_DEVICE; } return ucp_address_pack(worker, NULL, &tl_bitmap, flags, @@ -3149,6 +3162,7 @@ ucs_status_t ucp_worker_query(ucp_worker_h worker, ucp_worker_attr_t *attr) { ucs_status_t status = UCS_OK; + const char *address_device_name; uint32_t address_flags; if (attr->field_mask & UCP_WORKER_ATTR_FIELD_THREAD_MODE) { @@ -3158,9 +3172,12 @@ ucs_status_t ucp_worker_query(ucp_worker_h worker, if (attr->field_mask & UCP_WORKER_ATTR_FIELD_ADDRESS) { address_flags = UCP_ATTR_VALUE(WORKER, attr, address_flags, ADDRESS_FLAGS, 0); - status = ucp_worker_address_pack(worker, address_flags, - &attr->address_length, - (void**)&attr->address); + address_device_name = UCP_ATTR_VALUE(WORKER, attr, address_device_name, + ADDRESS_DEVICE_NAME, NULL); + status = ucp_worker_address_pack(worker, address_flags, + address_device_name, + &attr->address_length, + (void**)&attr->address); } if (attr->field_mask & UCP_WORKER_ATTR_FIELD_MAX_AM_HEADER) { @@ -3434,7 +3451,7 @@ ucs_status_t ucp_worker_get_address(ucp_worker_h worker, UCP_WORKER_THREAD_CS_ENTER_CONDITIONAL(worker); - status = ucp_worker_address_pack(worker, 0, address_length_p, + status = ucp_worker_address_pack(worker, 0, NULL, address_length_p, (void**)address_p); UCP_WORKER_THREAD_CS_EXIT_CONDITIONAL(worker); diff --git a/test/gtest/ucp/test_ucp_worker.cc b/test/gtest/ucp/test_ucp_worker.cc index 9a94d58d85b..cbaa2ec4cc7 100644 --- a/test/gtest/ucp/test_ucp_worker.cc +++ b/test/gtest/ucp/test_ucp_worker.cc @@ -8,6 +8,7 @@ #include #include #include +#include extern "C" { #include @@ -839,6 +840,117 @@ class test_ucp_worker_address_query : public ucp_test { add_variant_with_value(variants, UCP_FEATURE_TAG, 0, ""); add_variant_with_value(variants, UCP_FEATURE_TAG, 1, "unified"); } + +protected: + enum class address_device_type { + any, + net, + non_net + }; + + std::string find_address_device(const address_device_type device_type) + { + ucp_worker_h worker = sender().worker(); + ucp_context_h context = worker->context; + ucp_rsc_index_t tl_id; + + UCS_STATIC_BITMAP_FOR_EACH_BIT(tl_id, &context->tl_bitmap) { + const auto &resource = context->tl_rscs[tl_id].tl_rsc; + const auto &iface_attr = *ucp_worker_iface_get_attr(worker, tl_id); + const auto is_network = (iface_attr.cap.flags & + UCT_IFACE_FLAG_INTER_NODE) != 0; + + if (iface_can_connect(iface_attr) && + ((device_type == address_device_type::any) || + ((device_type == address_device_type::net) == is_network))) { + return resource.dev_name; + } + } + + return {}; + } + + void check_address_by_device(const std::string &address_device_name, + const uint32_t address_flags) + { + ucp_worker_h worker = sender().worker(); + ucp_context_h context = worker->context; + const ucp_address_entry_t *address_entry = nullptr; + ucp_worker_attr_t worker_attr{}; + ucp_unpacked_address_t unpacked_address{}; + ucp_tl_bitmap_t tl_bitmap; + ucp_rsc_index_t tl_id; + ucs_status_t status; + + ucp_context_dev_tl_bitmap(context, address_device_name.c_str(), + &tl_bitmap); + UCS_STATIC_BITMAP_FOR_EACH_BIT(tl_id, &tl_bitmap) { + const auto &iface_attr = *ucp_worker_iface_get_attr(worker, tl_id); + + if (!iface_matches_address_flags(iface_attr, address_flags)) { + UCS_STATIC_BITMAP_RESET(&tl_bitmap, tl_id); + } + } + + ASSERT_FALSE(UCS_STATIC_BITMAP_IS_ZERO(tl_bitmap)); + + worker_attr.field_mask = UCP_WORKER_ATTR_FIELD_ADDRESS | + UCP_WORKER_ATTR_FIELD_ADDRESS_DEVICE_NAME; + if (address_flags != 0) { + worker_attr.field_mask |= UCP_WORKER_ATTR_FIELD_ADDRESS_FLAGS; + worker_attr.address_flags = address_flags; + } + + worker_attr.address_device_name = address_device_name.c_str(); + + status = ucp_worker_query(worker, &worker_attr); + ASSERT_UCS_OK(status); + if (worker_attr.address == nullptr) { + UCS_TEST_ABORT("ucp_worker_query() returned a null address"); + } + + status = ucp_address_unpack(worker, worker_attr.address, + ucp_worker_default_address_pack_flags( + worker), + &unpacked_address); + ASSERT_UCS_OK(status); + EXPECT_EQ(UCS_STATIC_BITMAP_POPCOUNT(tl_bitmap), + unpacked_address.address_count); + + ucp_unpacked_address_for_each(address_entry, &unpacked_address) { + auto found = false; + + UCS_STATIC_BITMAP_FOR_EACH_BIT(tl_id, &tl_bitmap) { + const auto &resource = context->tl_rscs[tl_id]; + + if ((resource.md_index == address_entry->md_index) && + (resource.tl_name_csum == address_entry->tl_name_csum)) { + found = true; + break; + } + } + + EXPECT_TRUE(found); + } + + ucs_free(unpacked_address.address_list); + ucp_worker_release_address(worker, worker_attr.address); + } + +private: + static bool iface_can_connect(const uct_iface_attr_t &attrs) + { + return (attrs.cap.flags & (UCT_IFACE_FLAG_CONNECT_TO_IFACE | + UCT_IFACE_FLAG_CONNECT_TO_EP)) != 0; + } + + static bool iface_matches_address_flags(const uct_iface_attr_t &iface_attr, + const uint32_t address_flags) + { + return iface_can_connect(iface_attr) && + (((address_flags & UCP_WORKER_ADDRESS_FLAG_NET_ONLY) == 0) || + ((iface_attr.cap.flags & UCT_IFACE_FLAG_INTER_NODE) != 0)); + } }; UCS_TEST_P(test_ucp_worker_address_query, query) @@ -860,6 +972,62 @@ UCS_TEST_P(test_ucp_worker_address_query, query) ucp_worker_release_address(sender().worker(), worker_attr.address); } +UCS_TEST_P(test_ucp_worker_address_query, query_address_by_device) +{ + const auto address_device_name = find_address_device( + address_device_type::any); + + ASSERT_FALSE(address_device_name.empty()); + check_address_by_device(address_device_name, 0); +} + +UCS_TEST_P(test_ucp_worker_address_query, query_address_by_device_net_only) +{ + const auto address_device_name = find_address_device( + address_device_type::net); + + if (address_device_name.empty()) { + UCS_TEST_SKIP_R("no network device"); + } + + check_address_by_device(address_device_name, + UCP_WORKER_ADDRESS_FLAG_NET_ONLY); +} + +UCS_TEST_P(test_ucp_worker_address_query, empty_intersection) +{ + const auto address_device_name = find_address_device( + address_device_type::non_net); + scoped_log_handler wrap_err(wrap_errors_logger); + ucp_worker_attr_t worker_attr; + + if (address_device_name.empty()) { + UCS_TEST_SKIP_R("no non-network device"); + } + + worker_attr.field_mask = UCP_WORKER_ATTR_FIELD_ADDRESS | + UCP_WORKER_ATTR_FIELD_ADDRESS_FLAGS | + UCP_WORKER_ATTR_FIELD_ADDRESS_DEVICE_NAME; + worker_attr.address_flags = UCP_WORKER_ADDRESS_FLAG_NET_ONLY; + worker_attr.address_device_name = address_device_name.c_str(); + + EXPECT_EQ(UCS_ERR_NO_DEVICE, + ucp_worker_query(sender().worker(), &worker_attr)); +} + +UCS_TEST_P(test_ucp_worker_address_query, query_address_unknown_device) +{ + scoped_log_handler wrap_err(wrap_errors_logger); + ucp_worker_attr_t worker_attr; + + worker_attr.field_mask = UCP_WORKER_ATTR_FIELD_ADDRESS | + UCP_WORKER_ATTR_FIELD_ADDRESS_DEVICE_NAME; + worker_attr.address_device_name = "no-such-ucx-device"; + const auto status = ucp_worker_query(sender().worker(), &worker_attr); + + EXPECT_EQ(UCS_ERR_NO_DEVICE, status); +} + UCP_INSTANTIATE_TEST_CASE(test_ucp_worker_address_query) class test_ucp_worker_address_version : public ucp_test {