Optimizes the Margo DTL's memory usage by matching the behavior of the UCX DTL by ilumsden · Pull Request #178 · flux-framework/dyad

ilumsden · 2026-06-05T18:07:22Z

I was using DYAD on Tuolumne recently, and I found two issues with DYAD.

First, I found that DYAD would crash whenever trying to send/recv data with a Mercury NA protocol string of ofi+cxi (i.e., what should be used for Slingshot). This was due to an oversight in the dyad_dtl_margo_rpc_pack function. That function assumes that the network address will fit into a statically sized 128 byte buffer. That assumption is invalid on Tuo because the Slingshot network normally uses network addresses larger than 128 bytes.

This PR fixes this first issue by updating dyad_dtl_margo_rpc_pack to call margo_addr_to_string twice. The first call is used to get the actual length of the network address. Then, after a malloc, the second call is used to actually obtain the network address in full.

Second, I found that DYAD was performing much worse than expected. When reviewing the code, I found that the main issue was in how the Margo DTL was managing memory.

By allocating a buffer and calling margo_bulk_create for every send/recv, the DTL was essentially incurring the cost of memory allocation + memory pinning on the NIC repeatedly throughout a run. To add onto that, when I started using DYAD on Tuo, the Margo bulk objects were never being freed, resulting in massive memory leaks. Those memory leaks were recently fixed by other PRs, but the underlying alloc + pin cost still exists.

To fix this issue, this PR mimics the memory reuse scheme used by the UCX DTL. It now creates a single allocation and calls margo_bulk_create once during initialization, and it frees and unpins the memory during finalization. On the sender side, the pinned memory buffer is returned to the broker module via dyad_dtl_margo_get_buffer to ensure the module reads directly from local storage into the pinned memory. Similarly, on the receiver side, the pinned memory is passed directly into the margo_bulk_transfer call to avoid excessive allocations. The dyad_dtl_margo_recv function then copies from that pinned memory into a new allocation to return to the DYAD client. This PR also adds C11's stdatomic.h to help the dyad_dtl_margo_recv and data_ready_rpc synchronize on data availability.

This PR also makes the max transfer size for the UCX and Margo DTL backends configurable via the new CMake cache variable DYAD_DTL_MAX_TRANSFER_SIZE. The value of this variable defaults to 4294967296 (i.e., 4 GiB) which matches the existing behavior of the UCX DTL.

I tested these changes on Tuo. They result in much better performance. In fact, with these changes, DYAD can achieve results comparable to a data transfer approach using MPI_Isend and MPI_Irecv in certain workflow situations.

…tions on Tuo based on excessive memory pinning

… segfaults

… time

JaeseungYeom

I have not finished reviewing margo_dtl.c yet but I will come back later.

JaeseungYeom · 2026-06-05T19:53:39Z

    margo_handle->recv_buffer = NULL;
    margo_handle->recv_len = 0;
-    margo_handle->recv_ready = 0;
+    atomic_store_explicit (&margo_handle->recv_ready, false, memory_order_relaxed);


I am curious about this. It is good to know that this variable needs to be protected. However, was there incident?

Reusing buffer would improve performance. However, that also means that there will be no concurrent transfer.

stdatomic seems to be optional with c11. We either bump up the minimum C standard required if the protection is necessary. Otherwise, do test as in cmake/tests at cmake configure time and propagate the choice via dyad_config.h

JaeseungYeom · 2026-06-05T19:58:07Z


 # UCX implementation for DTL
-set(UCX_DTL_SRC ${CMAKE_CURRENT_SOURCE_DIR}/ucx_dtl.c ${CMAKE_CURRENT_SOURCE_DIR}/ucx_ep_cache.cpp)
+set(UCX_DTL_SRC ${CMAKE_CURRENT_BINARY_DIR}/ucx_dtl.c ${CMAKE_CURRENT_SOURCE_DIR}/ucx_ep_cache.cpp)


Instead of doing this, MAX_TRANSFER_SIZE as an env variable.
And parse size characters at the end.
We wouldn't need to change CMakeLists.txt and make it needlessly static.

JaeseungYeom · 2026-06-05T19:59:56Z

 #include <mercury_macros.h>

+// clang-format off
+#define MARGO_MAX_TRANSFER_SIZE (@DYAD_DTL_MAX_TRANSFER_SIZE@)


This sould be converted into the env variable DYAD_DTL_MAX_TX_SIZE_ENV "DYAD_DTL_MAX_TX_SIZE"

JaeseungYeom · 2026-06-05T20:32:41Z

    // clang-format on

    margo_instance_id mid = margo_hg_handle_get_instance (h);
-    margo_set_log_level (mid, MARGO_LOG_INFO);


#if define(DYAD_LOGGER_LEVEL_INFO)
margo_set_log_level (mid, MARGO_LOG_INFO);
#endif

JaeseungYeom · 2026-06-05T20:40:20Z

 {
    DYAD_C_FUNCTION_START ();
    dyad_rc_t rc = DYAD_RC_OK;
-    if (data_buf == NULL || *data_buf == NULL) {


I see that static buffer is already checked but it is still safer to check especially if we are going to reallocate.
In UCX case, you had a different problem, an unsolved mystery. That is something to revisit as well.

JaeseungYeom · 2026-06-05T20:44:03Z


    // dyad_rc_t rc = DYAD_RC_OK;
    dyad_dtl_margo_t *margo_handle = NULL;
+    dyad_rc_t rc = 0;


Initialize it with a return code macro.

JaeseungYeom · 2026-06-05T20:45:31Z


 #include <dyad/dtl/dyad_dtl_api.h>
 #include <margo.h>
+#include <stdatomic.h>


This is guaranteed to be available in C11.

JaeseungYeom · 2026-06-05T20:48:54Z

 {
    dyad_mod_ctx_t *mod_ctx = (dyad_mod_ctx_t *)arg;
-    flux_msg_handler_delvec (mod_ctx->handlers);
+    if (mod_ctx->handlers != NULL) {


JaeseungYeom · 2026-06-05T20:49:33Z

 endif()

-
+set(DYAD_DTL_MAX_TRANSFER_SIZE "4294967296" CACHE STRING "Maximum transfer size supported by DYAD in bytes")


We want this be rather runtime configurable.
CMake can set the deafult value.
DYAD_DTL_MAX_TX_SIZE_DEFAULT, which should be written to dyad_config.h
This should be commonly used for both UCX and MARGO, client and server.

This one is from shuffle test.
bool parse_size (const char* str, size_t& out)
{
char* end;
long long val = strtoll (str, &end, 10);
if (val <= 0 || end == str)
return false;

switch (*end) { case 'K': case 'k': out = static_cast<size_t> (val) * 1024ULL; break; case 'M': case 'm': out = static_cast<size_t> (val) * 1024ULL * 1024ULL; break; case 'G': case 'g': out = static_cast<size_t> (val) * 1024ULL * 1024ULL * 1024ULL; break; case '\0': out = static_cast<size_t> (val); break; default: return false; // unrecognized suffix } return true;

}

JaeseungYeom · 2026-06-05T21:05:12Z

Minimum to merge is to make the transfer size parameter environment variable controllable.

JaeseungYeom · 2026-06-05T21:56:57Z

+    }
+    addr_str = malloc (addr_str_size);
+    if (addr_str == NULL) {
+        rc = DYAD_RC_SYSFAIL;


TODO: Need to define a RC macro for this as well as other cases where assert is called.

JaeseungYeom · 2026-06-05T22:04:53Z

            margo_addr_free (margo_handle->mid, margo_handle->remote_addr);
        margo_finalize (margo_handle->mid);
    }
+


Move this up before if (margo_handle->mid != MARGO_INSTANCE_NULL) {

JaeseungYeom · 2026-06-05T22:05:48Z

    if (margo_handle->mid != MARGO_INSTANCE_NULL) {
        margo_addr_free (margo_handle->mid, margo_handle->local_addr);
        if (margo_handle->remote_addr != NULL)
            margo_addr_free (margo_handle->mid, margo_handle->remote_addr);


set margo_handle->remote_addr = HG_ADDR_NULL after margo_addr_free()

JaeseungYeom · 2026-06-05T22:06:41Z

+    }
+
    if (margo_handle->mid != MARGO_INSTANCE_NULL) {
        margo_addr_free (margo_handle->mid, margo_handle->local_addr);


margo_handle->local_addr = HG_ADDR_NULL

JaeseungYeom · 2026-06-05T22:09:11Z

+
    if (margo_handle->mid != MARGO_INSTANCE_NULL) {
        margo_addr_free (margo_handle->mid, margo_handle->local_addr);
        if (margo_handle->remote_addr != NULL)


margo_handle->remote_addr != HG_ADDR_NULL

JaeseungYeom · 2026-06-05T22:11:20Z


    // both margo client and server
    margo_addr_self (margo_handle->mid, &margo_handle->local_addr);
    margo_handle->remote_addr = NULL;


margo_handle->remote_addr = HG_ADDR_NULL

ilumsden added 4 commits June 5, 2026 13:29

Update Margo DTL to avoid memory leaks and massive performance limita…

d0cbbc0

…tions on Tuo based on excessive memory pinning

Add a pointer check around flux_msg_handler_delvec to avoid potential…

8b50a60

… segfaults

Eliminate a redundant goto/jump

e737b89

Make the max transfer size for the DTL backends configurable at build…

2c70594

… time

ilumsden requested a review from JaeseungYeom June 5, 2026 18:07

ilumsden added bug Something isn't working enhancement New feature or request labels Jun 5, 2026

ilumsden added 2 commits June 5, 2026 14:22

Fix a couple of minor syntax mistakes

c27ed68

Fix small bug where there was a missing variable declaration

b75acd5

JaeseungYeom assigned JaeseungYeom and ilumsden and unassigned JaeseungYeom and ilumsden Jun 5, 2026

JaeseungYeom requested changes Jun 5, 2026

View reviewed changes

JaeseungYeom reviewed Jun 5, 2026

View reviewed changes

		endif()


		set(DYAD_DTL_MAX_TRANSFER_SIZE "4294967296" CACHE STRING "Maximum transfer size supported by DYAD in bytes")

Conversation

ilumsden commented Jun 5, 2026

Uh oh!

JaeseungYeom left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

JaeseungYeom commented Jun 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

JaeseungYeom Jun 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

JaeseungYeom commented Jun 5, 2026 •

edited

Loading

JaeseungYeom Jun 5, 2026 •

edited

Loading