rccl: add multithread support to NetIbMPI validation suite

projects/rccl/.claude/skills/feature-unit-testing/SKILL.md

@@ -14,6 +14,12 @@ feature, no stack archaeology. The suite is the acceptance criterion for every P
 **Central discipline:** measure coverage first, plan second. Verify assumptions against
 real data before writing a single test.
 
+## Modules
+
+- [multithread-validation.md](multithread-validation.md) — covering concurrent use of a
+  net plugin: which production threading to reproduce, which concurrency is not worth
+  testing, worker-option parsing, and how to make overlap an assertion.
+
 ---
 
 ## Feature Testing Lifecycle

projects/rccl/.claude/skills/feature-unit-testing/multithread-validation.md

@@ -0,0 +1,92 @@
+# Multithread Validation of a Net Plugin
+
+Module of the *feature-unit-testing* skill. Read it when a transport test suite
+needs to cover concurrent use of a net plugin, or when deciding which
+concurrency is worth testing at all.
+
+## Pick the concurrency the product actually performs
+
+A plugin-level test replaces a production thread with a host worker thread, so
+first establish which production thread makes each call:
+
+- the proxy progress thread of a communicator drives the data path
+  (`isend` / `irecv` / `test`);
+- the proxy service thread of the same communicator performs connection setup
+  and memory registration (`listen` / `connect` / `accept` / `regMr`);
+- the plugin context is created per communicator, and communicators derived
+  from a parent share both the context and the proxy state.
+
+So one process with several independent communicators runs several plugin
+contexts, each driven by its own proxy threads. That — N independent contexts
+and connections, used concurrently inside one process — is the shape worth
+testing, and it is what frameworks produce when a single process owns several
+devices or several communication groups.
+
+Two shapes are *not* worth testing:
+
+- **Concurrent calls on one communicator.** One progress thread owns a
+  communicator, so the product never does this. The send and receive comms hold
+  shared FIFO, request, and work-request state with no per-comm serialization; a
+  test would be a data race whose green result carries no information.
+- **Concurrent plugin `init` / `finalize`.** The core serializes these under its
+  own plugin mutex, so the race is unreachable from the product. Verify this in
+  the core before writing such a test — if it holds, the test only exercises the
+  test harness.
+
+Confirm both claims in the current source rather than assuming; they are
+properties of the core, and they decide whether a test is meaningful.
+
+## Worker option
+
+`--net_ib_nthreads=N` selects how many host workers each MPI rank uses in the
+tests that support it. Parse and strip a custom flag before GTest sees argv,
+since `InitGoogleTest` does not know about it.
+
+Requirements worth enforcing in the parser, because a silently ignored value
+turns a multithread configuration into a single-threaded run:
+
+- reject anything that is not a positive integer, and reject conflicting
+  repeated values;
+- clamp to a maximum so a typo cannot exhaust the machine;
+- reduce the value across all ranks and fail before the suite starts unless
+  every rank resolved the same one.
+
+## Harness structure
+
+- The main thread creates every context and connection, exchanges handles over
+  MPI, and tears everything down after joining. Workers never call MPI: a
+  worker-side collective deadlocks the peer rank when a sibling fails.
+- Workers return a result struct instead of using GTest assertions, and the main
+  thread reports those results after joining. GTest fatal assertions are
+  thread-unsafe and a `return` inside a worker cannot fail the test.
+- Per-thread MPI tag namespaces keep the main-thread handshakes separate. Assert
+  at compile time that the worst-case tag fits the range every MPI
+  implementation must provide.
+- Synchronize a rank-wide failure flag between phases, so both ranks abandon a
+  connection handshake together instead of one waiting forever.
+
+## Prove that the workers overlapped
+
+A threaded test that a scheduler serialized still passes while covering nothing.
+Make overlap a test assertion:
+
+- release the workers from a common gate, then hold them at a second gate so no
+  worker can finish its body before the others have entered theirs — the
+  observed maximum of concurrently active workers becomes deterministic instead
+  of timing-dependent;
+- fail the test when the distinct thread count is below the requested worker
+  count, or when the observed overlap is less than two;
+- cancel the gates when thread creation fails, so already-started workers do not
+  wait forever.
+
+Keep the single-worker path byte-for-byte as it was, and gate the threaded path
+behind the option. The default run then stays a regression check for the
+original behavior.
+
+## What such tests are worth
+
+They exercise the state the plugin shares between contexts on one device —
+protection domain references, the memory-registration cache, device tables,
+asynchronous port recovery — under real concurrency, plus the per-connection
+data path in parallel. They say nothing about collective semantics or
+performance; those belong to suite-level and framework-level tests.

projects/rccl/test/common/MPIEnvironment.hpp

@@ -80,6 +80,25 @@ class MPIEnvironment : public ::testing::Environment
      */
     inline static int world_size{0};
 
+    /**
+     * @brief Upper bound accepted for --net_ib_nthreads
+     *
+     * Mirrors the MAX_THREADS cap NetSocketTests.cpp applies to
+     * NCCL_SOCKET_NTHREADS: guards against a typo'd value spawning enough
+     * threads to exhaust the machine.
+     */
+    static constexpr int kMaxThreads{16};
+
+    /**
+     * @brief Number of worker threads each MPI rank should fan out into
+     *
+     * Set from the --net_ib_nthreads=N CLI flag (parsed in main_mpi.cpp before
+     * GTest consumes argv), validated consistently across ranks, and clamped
+     * to [1, kMaxThreads]. Defaults to 1, which preserves single-threaded
+     * behavior identical to before multithread support was added.
+     */
+    inline static int nThreads{1};
+
     /**
      * @brief Aggregated return code for test results
      *

projects/rccl/test/common/main_mpi.cpp

@@ -13,8 +13,12 @@
  * its own main() function and use MPIHelpers for common functionality.
  */
 
+#include <charconv>
 #include <cstdio>
 #include <cstdlib>
+#include <cstring>
+#include <string>
+#include <system_error>
 #include <gtest/gtest.h>
 
 #ifdef MPI_TESTS_ENABLED
@@ -23,11 +27,138 @@
     #include "MPITestBase.hpp"
     #include "MPIEnvironment.hpp"
 
+namespace
+{
+struct NetIbThreadConfig
+{
+    int         nThreads{1};
+    bool        valid{true};
+    bool        clamped{false};
+    std::string error;
+};
+
+/**
+ * @brief Parse and strip --net_ib_nthreads=N from argv before GTest sees it
+ *
+ * Mirrors rccl-tests' own -t/--nthreads handling (parsed in its hand-written
+ * main() before any collective/threading logic runs): a custom flag must be
+ * consumed here since GTest's InitGoogleTest() does not know about it.
+ *
+ * Parsing is deliberately strict: accepting a malformed value as one worker
+ * would make a multithread configuration silently run single-threaded.
+ */
+NetIbThreadConfig parseAndStripNThreads(int* argc, char** argv)
+{
+    constexpr const char* kFlagPrefix = "--net_ib_nthreads=";
+    const size_t          kPrefixLen  = std::strlen(kFlagPrefix);
+    NetIbThreadConfig     config;
+    bool                  sawFlag = false;
+    // Conflict detection compares the value as written on the command line,
+    // before clamping. Comparing post-clamp values would let two distinct
+    // out-of-range requests (e.g. 99 and 100) both collapse to kMaxThreads and
+    // pass as non-conflicting, so the reject-conflicting-values contract would
+    // not hold above the cap.
+    int                   rawSeen = 0;
+
+    int writeIdx = 1;
+    for(int readIdx = 1; readIdx < *argc; ++readIdx)
+    {
+        if(std::strncmp(argv[readIdx], kFlagPrefix, kPrefixLen) != 0)
+        {
+            argv[writeIdx++] = argv[readIdx];
+            continue;
+        }
+
+        const char* value = argv[readIdx] + kPrefixLen;
+        const char* end   = value + std::strlen(value);
+        int         parsed = 0;
+        auto [ptr, ec] = std::from_chars(value, end, parsed);
+        if(value == end || ec != std::errc{} || ptr != end || parsed < 1)
+        {
+            config.valid = false;
+            if(config.error.empty())
+                config.error = "expected a positive integer for --net_ib_nthreads";
+            continue;
+        }
+
+        const int raw = parsed;
+
+        if(sawFlag && raw != rawSeen)
+        {
+            config.valid = false;
+            if(config.error.empty())
+                config.error = "conflicting --net_ib_nthreads values";
+            continue;
+        }
+
+        if(parsed > MPIEnvironment::kMaxThreads)
+        {
+            parsed         = MPIEnvironment::kMaxThreads;
+            config.clamped = true;
+        }
+
+        sawFlag         = true;
+        rawSeen         = raw;
+        config.nThreads = parsed;
+    }
+
+    *argc = writeIdx;
+    // C/POSIX guarantee argv[argc] == NULL; MPI_Init_thread() receives this
+    // argv, so restore the terminator after compacting.
+    argv[writeIdx] = nullptr;
+    return config;
+}
+
+bool validateNThreadsAcrossRanks(const NetIbThreadConfig& config, const MPIHelpers::MPIContext& mpiCtx)
+{
+    int localValid  = config.valid ? 1 : 0;
+    int allValid    = 0;
+    int minNThreads = 0;
+    int maxNThreads = 0;
+    MPI_Allreduce(&localValid, &allValid, 1, MPI_INT, MPI_MIN, MPI_COMM_WORLD);
+    MPI_Allreduce(&config.nThreads, &minNThreads, 1, MPI_INT, MPI_MIN, MPI_COMM_WORLD);
+    MPI_Allreduce(&config.nThreads, &maxNThreads, 1, MPI_INT, MPI_MAX, MPI_COMM_WORLD);
+
+    if(allValid && minNThreads == maxNThreads) return true;
+
+    if(mpiCtx.world_rank == 0)
+    {
+        if(!allValid)
+        {
+            std::fprintf(stderr,
+                         "ERROR: invalid --net_ib_nthreads configuration on at least one MPI rank");
+            if(!config.error.empty()) std::fprintf(stderr, " (%s)", config.error.c_str());
+            std::fprintf(stderr, "\n");
+        }
+        if(minNThreads != maxNThreads)
+        {
+            std::fprintf(stderr,
+                         "ERROR: --net_ib_nthreads must resolve to one value on all ranks "
+                         "(observed %d through %d)\n",
+                         minNThreads,
+                         maxNThreads);
+        }
+    }
+    return false;
+}
+} // namespace
+
 int main(int argc, char* argv[])
 {
+    // Parse our own custom flag before anything else touches argv
+    const NetIbThreadConfig threadConfig = parseAndStripNThreads(&argc, argv);
+
     // Initialize MPI using shared helper
     auto mpi_ctx = MPIHelpers::initializeMPI(&argc, &argv);
 
+    if(!validateNThreadsAcrossRanks(threadConfig, mpi_ctx))
+    {
+        MPI_Finalize();
+        MPIEnvironment::mpi_initialized = false;
+        return EXIT_FAILURE;
+    }
+    MPIEnvironment::nThreads = threadConfig.nThreads;
+
     const auto world_rank = mpi_ctx.world_rank;
     const auto world_size = mpi_ctx.world_size;
 
@@ -41,6 +172,10 @@ int main(int argc, char* argv[])
         TEST_INFO("MPI initialized - World size: %d, Thread support: %d",
                   world_size,
                   mpi_ctx.thread_support);
+        if(threadConfig.clamped)
+            TEST_INFO("--net_ib_nthreads exceeds maximum %d; using %d workers per rank",
+                      MPIEnvironment::kMaxThreads,
+                      MPIEnvironment::nThreads);
     }
 
     // Initialize Google Test