Skip to content

UCT/API: Purge outstanding work without an RX token #12006

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
gleon99 merged 3 commits into from
Oct 4, 2026
Merged

UCT/API: Purge outstanding work without an RX token #12006

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
2799ee8
UCT/IB/MLX5: Purge outstanding work without an RX token
evgeny-leksikov Sep 30, 2026
baa63ab
UCT/IB/MLX5: Reject a NULL RX token next to its dereference
evgeny-leksikov Oct 2, 2026
4e6b54e
TEST/UCT: Keep the RX token alive until outstanding purge returns
evgeny-leksikov Oct 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 8 additions & 2 deletions src/uct/api/v2/uct_v2.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1945,12 +1945,14 @@ typedef enum {
typedef struct {
/** Mask of valid fields, using bits from @ref
* uct_ep_outstanding_purge_field_t. @ref
* UCT_EP_OUTSTANDING_FIELD_RX_TOKEN and @ref
* UCT_EP_OUTSTANDING_FIELD_CB must be set. */
* UCT_EP_OUTSTANDING_FIELD_CB must be set. If @ref
* UCT_EP_OUTSTANDING_FIELD_RX_TOKEN is omitted, every outstanding
* operation is treated as undelivered. */

@svc-nvidia-pr-review svc-nvidia-pr-review Sep 30, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The new "treated as undelivered" rule is only documented on field_mask; the function/callback doc still describes the callback as reporting undelivered operations, so a caller that reposts from op_info can duplicate data that the peer already received. Can we add this note to the @ref uct_ep_outstanding_purge doc as well, and mention that without the token the callback may also report operations the peer already received, so reposting them can duplicate data?

@svc-nvidia-pr-review svc-nvidia-pr-review Oct 1, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

minor: pls document the optional case on the rx_token field as well ("Valid when @ref UCT_EP_OUTSTANDING_FIELD_RX_TOKEN is set"), like arg does — the field_mask comment alone is easy to miss.

uint64_t field_mask;

/**
* Opaque RX token received from the remote peer.
* Valid when @ref UCT_EP_OUTSTANDING_FIELD_RX_TOKEN is set.
*/
const void *rx_token;

Expand All @@ -1975,6 +1977,10 @@ typedef struct {
* @note This routine should be called only after the error handler for @a ep
* returns @ref UCS_INPROGRESS. It can be used only if the interface
* supports @c UCT_IFACE_FLAG_V2_QUERY_TOKEN.
*
* @note If @ref UCT_EP_OUTSTANDING_FIELD_RX_TOKEN is omitted, @a cb may also
* be invoked for operations the peer already received. Reposting those
* operations can duplicate data.
*/
ucs_status_t
uct_ep_outstanding_purge(uct_ep_h ep,
Expand Down
38 changes: 22 additions & 16 deletions src/uct/ib/mlx5/rc/rc_mlx5_ep.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1347,13 +1347,6 @@ static ucs_status_t uct_rc_mlx5_ep_outstanding_purge_check_params(
return UCS_ERR_INVALID_PARAM;
}

if (!(params->field_mask & UCT_EP_OUTSTANDING_FIELD_RX_TOKEN)) {
ucs_error("rc mlx5: rx token is not set");
return UCS_ERR_INVALID_PARAM;
}

ucs_assert(params->rx_token != NULL);

if (!(params->field_mask & UCT_EP_OUTSTANDING_FIELD_CB) ||
(params->cb == NULL)) {
ucs_error("rc mlx5: callback is not set or is NULL");
Expand Down Expand Up @@ -1452,18 +1445,31 @@ ucs_status_t uct_rc_mlx5_ep_outstanding_purge(
num_outstanding_packets = uct_rc_mlx5_txwq_outstanding_num_packets(
&iface->super.super, txwq, start_pi, end_pi);

rx_token = params->rx_token;
receiver_next_psn = ntohl(*rx_token) & UCT_IB_MLX5_PSN_MASK;
first_failed_psn = (uct_ib_mlx5_txwq_get_next_wqe_psn(txwq) -
num_outstanding_packets) & UCT_IB_MLX5_PSN_MASK;
psn_diff = (receiver_next_psn - first_failed_psn) &
UCT_IB_MLX5_PSN_MASK;
if (psn_diff > num_outstanding_packets) {
return UCS_ERR_INVALID_PARAM;
first_failed_psn = (uct_ib_mlx5_txwq_get_next_wqe_psn(txwq) -
num_outstanding_packets) & UCT_IB_MLX5_PSN_MASK;
if (params->field_mask & UCT_EP_OUTSTANDING_FIELD_RX_TOKEN) {
rx_token = params->rx_token;
if (rx_token == NULL) {

@svc-nvidia-pr-review svc-nvidia-pr-review Oct 2, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

minor: the rx_token == NULL validation moved out of uct_rc_mlx5_ep_outstanding_purge_check_params(), which exists for exactly this. Can we keep this NULL check there together with the other parameter validation, so all parameter checks stay in one place and happen before any WQ walking?

@svc-nvidia-pr-review svc-nvidia-pr-review Oct 2, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Parameter validation is now split between uct_rc_mlx5_ep_outstanding_purge_check_params() (NULL params, cb) and the main purge function (rx_token). A reader/maintainer has to look in two places for the purge param contract, and the rx_token check runs after the ucs_assertv_always() and packet-count work instead of with the other early checks.

pls keep this check in uct_rc_mlx5_ep_outstanding_purge_check_params() with the other param validations, so all the purge param contract stays in one place.

ucs_error("rc mlx5: rx token is NULL");
return UCS_ERR_INVALID_PARAM;
}

receiver_next_psn = ntohl(*rx_token) & UCT_IB_MLX5_PSN_MASK;
psn_diff = (receiver_next_psn - first_failed_psn) &
UCT_IB_MLX5_PSN_MASK;
if (psn_diff > num_outstanding_packets) {
return UCS_ERR_INVALID_PARAM;
}

delivered = 1;
} else {

@svc-nvidia-pr-review svc-nvidia-pr-review Sep 30, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe set receiver_next_psn = first_failed_psn here and keep delivered = 1 before the loop? Then uct_ib_mlx5_wqe_is_delivered() returns false on the first WQE and the existing single code path marks everything undelivered, without the dead receiver_next_psn = 0.

@evgeny-leksikov evgeny-leksikov Sep 30, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

overcomplicating, receiver_next_psn is not used if delivered == 1, it's initialized to suppress compiler warns

/* No peer receive position: every outstanding operation is
* undelivered. */
receiver_next_psn = 0;
delivered = 0;
}

wqe_first_psn = first_failed_psn;
delivered = 1;
for (pi = start_pi; pi != end_pi;
pi = uct_ib_mlx5_txwq_next_wqe_index(pi, wqe_size)) {
ctrl = uct_ib_mlx5_txwq_get_wqe(txwq, pi);
Expand Down
117 changes: 69 additions & 48 deletions test/gtest/uct/test_peer_failure.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -562,44 +562,52 @@ class test_uct_purge_outstanding : public uct_test {
return UCS_OK;
}

void purge_outstanding(purge_ctx *ctx)
void purge_outstanding(purge_ctx *ctx, bool with_rx_token = true)

@svc-nvidia-pr-review svc-nvidia-pr-review Oct 2, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

minor: the = true default value on purge_outstanding() is unused, the only caller passes it explicitly.

Suggested change
void purge_outstanding(purge_ctx *ctx, bool with_rx_token = true)
void purge_outstanding(purge_ctx *ctx, bool with_rx_token)

{
uct_ep_outstanding_purge_params_t purge_params = {};
uct_iface_attr_v2_t tx_attr = {};
uct_iface_attr_v2_t rx_attr = {};
uct_ep_attr_t ep_attr = {};
/* Lives until purge returns. The params keep a pointer into it. */
std::vector<uint8_t> rx_token;

tx_attr.field_mask = UCT_IFACE_ATTR_FIELD_TX_TOKEN_LENGTH;
ASSERT_UCS_OK(uct_iface_query_v2(m_sender->iface(), &tx_attr));
purge_params.field_mask = UCT_EP_OUTSTANDING_FIELD_CB |
UCT_EP_OUTSTANDING_FIELD_ARG;
purge_params.cb = purge_cb;
purge_params.arg = ctx;

rx_attr.field_mask = UCT_IFACE_ATTR_FIELD_RX_TOKEN_LENGTH;
ASSERT_UCS_OK(uct_iface_query_v2(m_receiver->iface(), &rx_attr));
if (with_rx_token) {
uct_iface_attr_v2_t tx_attr = {};
uct_iface_attr_v2_t rx_attr = {};
uct_ep_attr_t ep_attr = {};
std::vector<uint8_t> tx_token;

std::vector<uint8_t> tx_token(tx_attr.tx_token_length);
std::vector<uint8_t> rx_token(rx_attr.rx_token_length);
tx_attr.field_mask = UCT_IFACE_ATTR_FIELD_TX_TOKEN_LENGTH;
ASSERT_UCS_OK(uct_iface_query_v2(m_sender->iface(), &tx_attr));

ep_attr.field_mask = UCT_EP_ATTR_FIELD_TX_TOKEN;
ep_attr.tx_token = tx_token.data();
ASSERT_UCS_OK(uct_ep_query(m_sender->ep(0), &ep_attr));
rx_attr.field_mask = UCT_IFACE_ATTR_FIELD_RX_TOKEN_LENGTH;
ASSERT_UCS_OK(uct_iface_query_v2(m_receiver->iface(), &rx_attr));

rx_attr.field_mask = UCT_IFACE_ATTR_FIELD_TX_TOKEN |
UCT_IFACE_ATTR_FIELD_RX_TOKEN;
rx_attr.tx_token = tx_token.data();
rx_attr.rx_token = rx_token.data();
ASSERT_UCS_OK(uct_iface_query_v2(m_receiver->iface(), &rx_attr));
tx_token.resize(tx_attr.tx_token_length);
rx_token.resize(rx_attr.rx_token_length);

ep_attr.field_mask = UCT_EP_ATTR_FIELD_TX_TOKEN;
ep_attr.tx_token = tx_token.data();
ASSERT_UCS_OK(uct_ep_query(m_sender->ep(0), &ep_attr));

rx_attr.field_mask = UCT_IFACE_ATTR_FIELD_TX_TOKEN |
UCT_IFACE_ATTR_FIELD_RX_TOKEN;
rx_attr.tx_token = tx_token.data();
rx_attr.rx_token = rx_token.data();
ASSERT_UCS_OK(uct_iface_query_v2(m_receiver->iface(), &rx_attr));

purge_params.field_mask |= UCT_EP_OUTSTANDING_FIELD_RX_TOKEN;
purge_params.rx_token = rx_token.data();
}

purge_params.field_mask = UCT_EP_OUTSTANDING_FIELD_RX_TOKEN |
UCT_EP_OUTSTANDING_FIELD_CB |
UCT_EP_OUTSTANDING_FIELD_ARG;
purge_params.rx_token = rx_token.data();
purge_params.cb = purge_cb;
purge_params.arg = ctx;
ASSERT_UCS_OK(uct_ep_outstanding_purge(m_sender->ep(0),
&purge_params));
}

void test_purge_outstanding(const send_func_t &send_func,
purge_ctx &ctx)
void test_purge_outstanding(const send_func_t &send_func, purge_ctx &ctx,
bool with_rx_token = true)
{
static constexpr uint32_t NUM_MSG_BEFORE_INVALIDATE = 2;
uct_ep_invalidate_params_t invalidate_params = {};
Expand Down Expand Up @@ -628,7 +636,7 @@ class test_uct_purge_outstanding : public uct_test {

num_outstanding = ctx.comp.count;
num_completions = ctx.num_completions;
purge_outstanding(&ctx);
purge_outstanding(&ctx, with_rx_token);

EXPECT_GT(ctx.num_ops_purged, 0u);
EXPECT_LE(ctx.num_ops_purged, num_posted);
Expand All @@ -645,6 +653,34 @@ class test_uct_purge_outstanding : public uct_test {
EXPECT_EQ(0, ctx.comp.count);
}

void test_am_short(bool with_rx_token)
{
const uct_iface_attr_t &attr = m_sender->iface_attr();
const size_t size = ucs_min((size_t)64,
attr.cap.am.max_short);
std::vector<uint8_t> payload(size);

mem_buffer::pattern_fill(payload.data(), payload.size(), SEND_SEED);

purge_ctx ctx = {this, UCT_EP_OP_AM_SHORT, {completion_cb, 0, UCS_OK}};
ctx.send_buf = payload.data();
ctx.send_len = payload.size();

ASSERT_UCS_OK(uct_iface_set_am_handler(m_receiver->iface(), AM_SHORT_ID,
am_handler, NULL, 0));

send_func_t am_short = [&](uct_ep_h ep, uct_completion_t*) {
return uct_ep_am_short(ep, AM_SHORT_ID, AM_SHORT_HEADER,
ctx.send_buf, ctx.send_len);
};
test_purge_outstanding(am_short, ctx, with_rx_token);

EXPECT_GT(ctx.num_ops_purged_at_completion, 0u);
EXPECT_GT(ctx.num_ops_posted_after_flush, 0u);
EXPECT_EQ(ctx.num_ops_purged, ctx.num_ops_purged_at_completion +
ctx.num_ops_posted_after_flush);
}

entity *m_sender;
entity *m_receiver;
unsigned m_err_count = 0;
Expand All @@ -656,28 +692,13 @@ const uint64_t test_uct_purge_outstanding::AM_SHORT_HEADER;
UCS_TEST_SKIP_COND_P(test_uct_purge_outstanding, am_short,
!check_caps(UCT_IFACE_FLAG_AM_SHORT))
{
const uct_iface_attr_t &attr = m_sender->iface_attr();
const size_t size = ucs_min((size_t)64, attr.cap.am.max_short);
std::vector<uint8_t> payload(size);
mem_buffer::pattern_fill(payload.data(), payload.size(), SEND_SEED);

purge_ctx ctx = {this, UCT_EP_OP_AM_SHORT, {completion_cb, 0, UCS_OK}};
ctx.send_buf = payload.data();
ctx.send_len = payload.size();

ASSERT_UCS_OK(uct_iface_set_am_handler(m_receiver->iface(), AM_SHORT_ID,
am_handler, NULL, 0));

send_func_t am_short = [&](uct_ep_h ep, uct_completion_t*) {
return uct_ep_am_short(ep, AM_SHORT_ID, AM_SHORT_HEADER, ctx.send_buf,
ctx.send_len);
};
test_purge_outstanding(am_short, ctx);
test_am_short(true);
}

EXPECT_GT(ctx.num_ops_purged_at_completion, 0u);
EXPECT_GT(ctx.num_ops_posted_after_flush, 0u);
EXPECT_EQ(ctx.num_ops_purged, ctx.num_ops_purged_at_completion +
ctx.num_ops_posted_after_flush);
UCS_TEST_SKIP_COND_P(test_uct_purge_outstanding, am_short_no_rx_token,

@svc-nvidia-pr-review svc-nvidia-pr-review Sep 30, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This test body is a copy of am_short; can we extract a helper (e.g. do_test_am_short(bool with_rx_token)) and call it from both tests?

@svc-nvidia-pr-review svc-nvidia-pr-review Sep 30, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

minor: this is an exact copy of am_short, pls move the body to a helper (e.g. test_am_short(bool with_rx_token)) and call it from both tests.

@svc-nvidia-pr-review svc-nvidia-pr-review Oct 1, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

the new test is a line-for-line copy of am_short except for the last argument; pls extract the body into a helper, e.g. test_am_short(bool with_rx_token), and keep both test cases as one-liners.

!check_caps(UCT_IFACE_FLAG_AM_SHORT))
{
test_am_short(false);
}

UCS_TEST_SKIP_COND_P(test_uct_purge_outstanding, put_short,
Expand Down