Enabling concurrent spilling

[nirandaperera]

[C++] Allow concurrent SpillManager::spill() invocations

SpillManager serialized all spill() calls behind a single std::mutex, even though spill work is long-running (D2H copies, I/O) and only the rarely-touched function registry actually needed mutual exclusion. This switches the lock to a std::shared_mutex so spill() takes a shared lock and runs in parallel, while add_spill_function / remove_spill_function still take an exclusive lock — which preserves the drain-on-remove guarantee that callers like Shuffler rely on for safe teardown.

• spill() now takes a shared lock; concurrent callers no longer serialize on the registry mutex.
• add_spill_function / remove_spill_function take an exclusive lock, so in-flight spillers are drained before either returns.
• Added SpillManager.ConcurrentSpill test that uses a rendezvous to prove spill() calls actually overlap (it would fail on the old exclusive mutex).
• Audited the three registered C++ spill functions (Shuffler::spill, AllGather::spill, spill_messages); all already lock their own state for concurrent self-invocation, so no caller-side changes are required.

[madsbk]

Overall looks good @nirandaperera

[wence-]

With the changes Mads proposes to the docstring

[wence-]

I note that any given spill function will probably run in a mutex-protected section (since it needs to modify a container). Hence although SpillManager::spill can now run concurrently, multiple threads will potentially still block in spill because they serialise when calling the spill functions. Suppose we have three spill functions registered and want to call spill from two threads. Thread-1 comes in, starts spilling from fid-0. Then thread-2 comes in and blocks trying to spill from fid-0 (the same container that thread-1 is spilling from), even if it could concurrently spill from fid-1.

One could do something like:

diff --git a/cpp/src/memory/spill_manager.cpp b/cpp/src/memory/spill_manager.cpp
index 6b20150d..15e9c384 100644
--- a/cpp/src/memory/spill_manager.cpp
+++ b/cpp/src/memory/spill_manager.cpp
@@ -3,6 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+#include <atomic>
 #include <utility>
 
 #include <rapidsmpf/memory/buffer_resource.hpp>
@@ -39,7 +40,7 @@ SpillManager::SpillFunctionID SpillManager::add_spill_function(
         spill_functions_.insert({id, std::move(spill_function)}).second,
         "corrupted id counter"
     );
-    spill_function_priorities_.insert({priority, id});
+    spill_function_priorities_.insert({priority, std::pair{id, false}});
 
     // Make sure the spill thread is running.
     if (periodic_spill_thread_.has_value()) {
@@ -69,11 +70,18 @@ std::size_t SpillManager::spill(std::size_t amount) {
     RAPIDSMPF_NVTX_FUNC_RANGE();
     std::size_t spilled{0};
     std::unique_lock<std::mutex> lock(mutex_);
-    for (auto const [_, fid] : spill_function_priorities_) {
+    for (auto& [_, state] : spill_function_priorities_) {
         if (spilled >= amount) {
             break;
         }
-        spilled += spill_functions_.at(fid)(amount - spilled);
+        auto& [fid, in_use] = state;
+        bool expect = false;
+        if (in_use.compare_exchange_strong(expect, true, std::memory_order_acq_rel)) {
+            spilled += spill_functions_.at(fid)(amount - spilled);
+            in_use.store(false, std::memory_order_release);
+        } else {
+            continue;
+        }
     }
     return spilled;
 }

But that has the disadvantage that if a call to spill-function-0 requests a small amount (from a container with lots to spill), and then a call to spill-function-1 requests a large amount (from a container with little to spill). We can spuriously fail to spill enough.

[nirandaperera]

@wence- sorry for not responding earlier. I had this PR on the back-burner until now.

I like the idea.

But that has the disadvantage that if a call to spill-function-0 requests a small amount (from a container with lots to spill), and then a call to spill-function-1 requests a large amount (from a container with little to spill). We can spuriously fail to spill enough.

To tackle this, how about we guarantee that all spill functions have been called by revisiting the ones that were skipped?

[nirandaperera]

I want to hold off this PR, because it didnt show immediate spilling perf improvement for PDX 1GPU 3k TPCH runs. Its interesting why this would be? 🤔

Some immediate observations.

1. Irrespective of this PR, Q7 suddenly spills on the second iteration (not the first). Should we sync streams between iterations?
2. Irrespective of this PR, Q21 OOMs (I was able to run this with 80% spill device before)
3. With this PR, Q9 iter 2 goes OOM (I suspect, 825e4df because now, spilling might not be sufficient)