fix(mpi): 🐛 do not fail a completed partition barrier round on a later poison

cpp/monoprop/detail/mpi/PartitionBarrier.h

@@ -43,31 +43,36 @@ class PartitionBarrier {
     PartitionBarrier(const PartitionBarrier &) = delete;
     auto operator=(const PartitionBarrier &) -> PartitionBarrier & = delete;
 
+    // Throws ShmCommPoisoned only for a round that cannot complete. Once every participant has arrived the
+    // round has completed, and a poison set afterwards reports a later failure: its poisoner never arrives at
+    // the next round, which surfaces it. Failing the completed round instead would fail a collective that
+    // succeeded, and under the MPI hybrid partition 0 turns that into MPI_Abort.
     auto sync() -> void {
         const unsigned g = gen_.load(std::memory_order_acquire);
         if (arrived_.fetch_add(1, std::memory_order_acq_rel) + 1 == participants_) {
             arrived_.store(0, std::memory_order_relaxed);
             gen_.store(g + 1, std::memory_order_release);
+            return;
         }
-        else {
-            // Bounded on-core spin first (pinned partitions ⇒ the release store lands in the pause window,
-            // no syscall); only long waits (imbalance, oversubscription) fall to yield.
-            int spins = 0;
-            while (gen_.load(std::memory_order_acquire) == g) {
-                if (poisoned_.load(std::memory_order_acquire)) {
-                    throw ShmCommPoisoned();
-                }
-                if (spins < detail::kSpinPauseIters) {
-                    ++spins;
-                    detail::cpu_relax();
-                }
-                else {
-                    std::this_thread::yield();
+        // Bounded on-core spin first (pinned partitions ⇒ the release store lands in the pause window,
+        // no syscall); only long waits (imbalance, oversubscription) fall to yield.
+        int spins = 0;
+        while (gen_.load(std::memory_order_acquire) == g) {
+            if (poisoned_.load(std::memory_order_acquire)) {
+                // A poisoner that left this round saw gen_ advance first, so seeing its poison means seeing
+                // the advance too: re-read it before failing a round that may have just completed.
+                if (gen_.load(std::memory_order_acquire) != g) {
+                    return;
                 }
+                throw ShmCommPoisoned();
+            }
+            if (spins < detail::kSpinPauseIters) {
+                ++spins;
+                detail::cpu_relax();
+            }
+            else {
+                std::this_thread::yield();
             }
-        }
-        if (poisoned_.load(std::memory_order_acquire)) {
-            throw ShmCommPoisoned();
         }
     }
 

cpp/monoprop/detail/partition/PartitionGroup.h

@@ -113,8 +113,7 @@ class PartitionGroup {
     // that caused the failure: the lowest-ranked one that is not ShmCommPoisoned. A throw on one master never
     // hangs the rest, since peers are released via poison.
     // ShmCommPoisoned only ever follows another master's throw, so it is rethrown only when nothing else was
-    // raised. It can even reach a master whose collective already completed, if a peer throws between the
-    // barrier's release and that master's post-barrier check.
+    // raised. A master whose collective already completed never sees it (see PartitionBarrier::sync()).
     auto run_on_all(const std::function<void(int)> &body) -> void {
         {
             std::lock_guard lk(m_);

cpp/tests/shm_comm_tests.cpp

@@ -26,9 +26,11 @@
 #include "monoprop/detail/mpi/Comm.h"
 #include "monoprop/detail/mpi/Exchange.h"
 #include "monoprop/detail/mpi/MPICompat.h"
+#include "monoprop/detail/mpi/PartitionBarrier.h"
 #include "monoprop/detail/mpi/ShmComm.h"
 
 using monoprop::mpi::Comm;
+using monoprop::mpi::PartitionBarrier;
 using monoprop::mpi::ShmComm;
 using monoprop::mpi::ShmCommPoisoned;
 
@@ -349,3 +351,34 @@ BOOST_AUTO_TEST_CASE(shm_comm_poison_releases_waiters) {
         }
     }
 }
+
+// A round every participant reached has completed, so it returns normally even when poison is already set:
+// the poison reports a later failure, which the next round surfaces. The last arriver completes the round.
+BOOST_AUTO_TEST_CASE(partition_barrier_completed_round_ignores_later_poison) {
+    PartitionBarrier barrier(1);
+    barrier.poison();
+    BOOST_CHECK_NO_THROW(barrier.sync());
+}
+
+// Every partition finishes a collective, then partition 1 fails and poisons, as a PartitionGroup worker
+// does when its job throws. Its siblings completed the collective and must not see ShmCommPoisoned: under
+// the MPI hybrid, partition 0 turns that throw into MPI_Abort. The race needs a sibling still between the
+// final barrier's release and its return when the poison lands, hence the repetitions.
+BOOST_AUTO_TEST_CASE(shm_comm_poison_after_completed_collective_spares_siblings) {
+    constexpr int kRepetitions = 2000;
+    for (const int S : {2, 4}) {
+        for (int rep = 0; rep < kRepetitions; ++rep) {
+            auto errs = run_shm(S, [&](ShmComm &sh, int r) {
+                (void)sh.allreduce_sum<int>(r, 1);
+                if (r == 1) {
+                    sh.poison();
+                }
+            });
+            for (int r = 0; r < S; ++r) {
+                BOOST_REQUIRE_MESSAGE(
+                    errs[static_cast<size_t>(r)] == nullptr,
+                    "S=" << S << " repetition " << rep << ": partition " << r << " failed a collective it completed");
+            }
+        }
+    }
+}

cpp/tests/term_expectation_values_tests.cpp

@@ -185,7 +185,8 @@ BOOST_AUTO_TEST_CASE(schrodinger_values_are_the_evolved_state_and_match_heisenbe
 
 // Heisenberg throws for a term absent from the evolved operator. The term is past the cutoff and not fully
 // paired (fully-paired terms are exempt from truncation), so it cannot be present. With partitions, every
-// partition throws right after a collective; this case exposed the poison race fixed in run_on_all().
+// partition throws right after a collective, which must not fail a sibling's completed collective
+// (shm_comm_poison_after_completed_collective_spares_siblings covers that race directly).
 BOOST_AUTO_TEST_CASE(heisenberg_absent_term_throws) {
     const auto data = load_case_data<kNumModes>(kCase);
     const VecZ absent{0, 2, 4, 6, 8, 10, 12, 14};