kernel-proxy: address PR #514 review findings

Author: MC952-arch

flagcx/core/flagcx_hetero.cc

@@ -789,7 +789,8 @@ flagcxResult_t flagcxHeteroFlushRma(flagcxHeteroComm_t comm, int peer,
       return flagcxRemoteError;
     usleep(100);
   }
-  // Check rmaError after wait: doneSeqs may have been advanced despite failure
+  // Check rmaError after wait: kernel proxy direct-post path advances doneSeqs
+  // unconditionally (even on failure) to prevent flush hangs.
   if (__atomic_load_n(&proxy->rmaError, __ATOMIC_ACQUIRE))
     return flagcxRemoteError;
   return flagcxSuccess;
@@ -830,6 +831,10 @@ flagcxResult_t flagcxHeteroFlushAllRma(flagcxHeteroComm_t comm) {
       usleep(100);
     }
   }
+  // Check rmaError after wait: kernel proxy direct-post path advances doneSeqs
+  // unconditionally (even on failure) to prevent flush hangs.
+  if (__atomic_load_n(&proxy->rmaError, __ATOMIC_ACQUIRE))
+    return flagcxRemoteError;
   return flagcxSuccess;
 }
 

flagcx/core/proxy.cc

@@ -24,6 +24,7 @@
 #include <string>
 #include <sys/syscall.h>
 #include <sys/time.h>
+#include <time.h>
 #include <unistd.h>
 using namespace std;
 
@@ -1402,6 +1403,8 @@ void *flagcxProxyService(void *args) {
 // Bypasses RMA Proxy: posts IB ops directly from kernel proxy thread.
 // ============================================================================
 
+// Note: each kernel proxy thread allocates nRanks × (256 × 16B + 8B) ≈
+// 4KB/peer. With contextCount threads, total is contextCount × nRanks × 4KB.
 #define FLAGCX_KPROXY_MAX_INFLIGHT 256
 #define FLAGCX_KPROXY_RING_MASK (FLAGCX_KPROXY_MAX_INFLIGHT - 1)
 
@@ -1426,7 +1429,9 @@ static void flagcxKernelProxyPoll(struct flagcxKernelProxyState *state,
                                   struct flagcxHeteroComm *comm) {
   struct flagcxRmaProxyState *proxy = comm->rmaProxy;
   struct flagcxNetAdaptor *net = comm->netAdaptor;
-  bool anyCompleted = false;
+  if (proxy == NULL || net == NULL)
+    return;
+  bool anyRetired = false;
   for (int p = 0; p < state->nRanks; p++) {
     struct flagcxKernelProxyPeerState *ps = &state->peers[p];
     while (ps->head != ps->tail) {
@@ -1451,6 +1456,7 @@ static void flagcxKernelProxyPoll(struct flagcxKernelProxyState *state,
       if (!done)
         break;
       ps->head++;
+      anyRetired = true;
       // Always advance doneSeqs (even on failure) to prevent flush hangs.
       __atomic_store_n((uint64_t *)&proxy->doneSeqs[p], inf->opSeq,
                        __ATOMIC_RELEASE);
@@ -1459,13 +1465,11 @@ static void flagcxKernelProxyPoll(struct flagcxKernelProxyState *state,
                          __ATOMIC_RELEASE);
       if (!failed) {
         __atomic_fetch_add(&proxy->completionCount, 1ULL, __ATOMIC_RELEASE);
-        anyCompleted = true;
       }
     }
   }
-  // Batch the broadcast: signal waiters once per poll sweep, not per
-  // completion.
-  if (anyCompleted && !proxy->useStreamOps) {
+  // Batch the broadcast: signal waiters once per poll sweep on any retirement.
+  if (anyRetired && !proxy->useStreamOps) {
     pthread_mutex_lock(&proxy->doneMutex);
     pthread_cond_broadcast(&proxy->doneCond);
     pthread_mutex_unlock(&proxy->doneMutex);
@@ -1504,12 +1508,18 @@ flagcxKernelProxyPost(struct flagcxKernelProxyState *state,
   }
 
   // Back-pressure: if ring is full, poll until a slot frees.
-  int spins = 0;
+  struct timespec deadline;
+  clock_gettime(CLOCK_MONOTONIC, &deadline);
+  deadline.tv_sec += 30; // 30s timeout
   while ((ps->tail - ps->head) >= FLAGCX_KPROXY_MAX_INFLIGHT) {
     flagcxKernelProxyPoll(state, comm);
     if ((ps->tail - ps->head) >= FLAGCX_KPROXY_MAX_INFLIGHT) {
-      if (++spins > 10000000) {
-        WARN("flagcxKernelProxyPost: back-pressure timeout peer=%d", peer);
+      struct timespec now;
+      clock_gettime(CLOCK_MONOTONIC, &now);
+      if (now.tv_sec > deadline.tv_sec ||
+          (now.tv_sec == deadline.tv_sec && now.tv_nsec >= deadline.tv_nsec)) {
+        WARN("flagcxKernelProxyPost: back-pressure timeout (30s) peer=%d",
+             peer);
         __atomic_store_n((int *)&proxy->rmaError, 1, __ATOMIC_RELEASE);
         return flagcxInternalError;
       }
@@ -1518,13 +1528,22 @@ flagcxKernelProxyPost(struct flagcxKernelProxyState *state,
   }
 
   void *sendComm = comm->oneSideHandles[0]->fullSendComms[peer];
+  if (sendComm == NULL) {
+    WARN("flagcxKernelProxyPost: sendComm is NULL for peer=%d", peer);
+    return flagcxInternalError;
+  }
   void **srcHandles = NULL, **dstHandles = NULL;
-  if (size > 0 && srcMrIdx >= 0) {
+  if (size > 0 && srcMrIdx >= 0 && dstMrIdx >= 0) {
     srcHandles = (void **)comm->oneSideHandles[srcMrIdx];
     dstHandles = (void **)comm->oneSideHandles[dstMrIdx];
   }
 
-  // Assign opSeq (shared atomic counter with RMA Proxy for ordering)
+  // Assign opSeq (shared atomic counter with RMA Proxy for ordering).
+  // Single-writer invariant: for a given peer, only the kernel proxy thread
+  // (this path) or the RMA Proxy thread (stream path) posts ops at any time.
+  // This ensures opSeqs are assigned and complete in order, so doneSeqs
+  // advances monotonically. Violating this invariant would require per-peer
+  // max-tracking instead of direct store.
   uint64_t opSeq =
       __atomic_add_fetch((uint64_t *)&proxy->opSeqs[peer], 1, __ATOMIC_RELAXED);
 
@@ -1550,6 +1569,28 @@ flagcxKernelProxyPost(struct flagcxKernelProxyState *state,
       break;
     }
     case FLAGCX_RMA_PUT_VALUE: {
+      // Serialize PUT_VALUE per-peer: drain all prior ops to ensure the staging
+      // slot is idle. The NIC reads staging asynchronously, so concurrent
+      // PUT_VALUE to the same peer would corrupt the buffer.
+      struct timespec pvDeadline;
+      clock_gettime(CLOCK_MONOTONIC, &pvDeadline);
+      pvDeadline.tv_sec += 30;
+      while (ps->head != ps->tail) {
+        flagcxKernelProxyPoll(state, comm);
+        if (ps->head != ps->tail) {
+          struct timespec now;
+          clock_gettime(CLOCK_MONOTONIC, &now);
+          if (now.tv_sec > pvDeadline.tv_sec ||
+              (now.tv_sec == pvDeadline.tv_sec &&
+               now.tv_nsec >= pvDeadline.tv_nsec)) {
+            WARN("flagcxKernelProxyPost: PUT_VALUE drain timeout (30s) peer=%d",
+                 peer);
+            __atomic_store_n((int *)&proxy->rmaError, 1, __ATOMIC_RELEASE);
+            return flagcxInternalError;
+          }
+          sched_yield();
+        }
+      }
       struct flagcxOneSideHandleInfo *stagingH = comm->stagingHandle;
       if (stagingH == NULL || stagingH->baseVas == NULL) {
         WARN("flagcxKernelProxyPost: staging handles not initialized");
@@ -1593,6 +1634,8 @@ static void flagcxKernelProxyDrain(struct flagcxKernelProxyState *state,
     return;
   struct flagcxNetAdaptor *net = comm->netAdaptor;
   struct flagcxRmaProxyState *proxy = comm->rmaProxy;
+  if (net == NULL || proxy == NULL)
+    return;
   for (int p = 0; p < state->nRanks; p++) {
     struct flagcxKernelProxyPeerState *ps = &state->peers[p];
     while (ps->head != ps->tail) {
@@ -1609,6 +1652,8 @@ static void flagcxKernelProxyDrain(struct flagcxKernelProxyState *state,
             done = 1;
             break;
           }
+          if (!done)
+            sched_yield();
         }
         succeeded = !testFailed;
       }

test/unittest/rma/coll_rma_get.cpp

@@ -74,8 +74,9 @@ TEST_F(RmaTest, GetSmall) {
     for (size_t i = 0; i < testSize; ++i) {
       if (received[i] != 0xCD) {
         mismatches++;
-        if (mismatches == 1)
+        if (mismatches == 1) {
           EXPECT_EQ(received[i], 0xCD) << "Mismatch at byte " << i;
+        }
       }
     }
     EXPECT_EQ(mismatches, 0);
@@ -130,8 +131,9 @@ TEST_F(RmaTest, GetLarge) {
       uint8_t expected = static_cast<uint8_t>((i * 7) & 0xFF);
       if (received[i] != expected) {
         mismatches++;
-        if (mismatches == 1)
+        if (mismatches == 1) {
           EXPECT_EQ(received[i], expected) << "Mismatch at byte " << i;
+        }
       }
     }
     EXPECT_EQ(mismatches, 0);
@@ -184,9 +186,10 @@ TEST_F(RmaTest, GetBidirectional) {
     uint8_t expected = static_cast<uint8_t>((peer + 1 + i) & 0xFF);
     if (received[i] != expected) {
       mismatches++;
-      if (mismatches == 1)
+      if (mismatches == 1) {
         EXPECT_EQ(received[i], expected)
             << "Mismatch at byte " << i << " reading from rank " << peer;
+      }
     }
   }
   EXPECT_EQ(mismatches, 0);

test/unittest/rma/coll_rma_put.cpp

@@ -77,8 +77,9 @@ TEST_F(RmaTest, PutSignalSmall) {
     for (size_t i = 0; i < testSize; ++i) {
       if (received[i] != 0xAB) {
         mismatches++;
-        if (mismatches == 1)
+        if (mismatches == 1) {
           EXPECT_EQ(received[i], 0xAB) << "Mismatch at byte " << i;
+        }
       }
     }
     EXPECT_EQ(mismatches, 0);
@@ -132,8 +133,9 @@ TEST_F(RmaTest, PutSignalLarge) {
       uint8_t expected = static_cast<uint8_t>(i & 0xFF);
       if (received[i] != expected) {
         mismatches++;
-        if (mismatches == 1)
+        if (mismatches == 1) {
           EXPECT_EQ(received[i], expected) << "Mismatch at byte " << i;
+        }
       }
     }
     EXPECT_EQ(mismatches, 0);