Message spilling too

Author: TomAugspurger

cpp/src/shuffler/shuffler.cpp

@@ -138,10 +138,14 @@ std::size_t postbox_spilling(
             continue;
         }
         // We extract the chunk, spilled it, and insert it back into the PostBox.
-        auto chunk = postbox.extract(pid, cid);
-        chunk.set_data_buffer(br->move(chunk.release_data_buffer(), host_reservation));
-        postbox.insert(std::move(chunk));
-        RAPIDSMPF_NVTX_MARKER("postbox_spilling::chunk_spilled_bytes", size);
+        {
+            RAPIDSMPF_NVTX_SCOPED_RANGE("postbox_spilling::chunk", size);
+            auto chunk = postbox.extract(pid, cid);
+            chunk.set_data_buffer(
+                br->move(chunk.release_data_buffer(), host_reservation)
+            );
+            postbox.insert(std::move(chunk));
+        }
         if ((total_spilled += size) >= amount) {
             break;
         }

cpp/src/streaming/core/context.cpp

@@ -7,6 +7,7 @@
 #include <utility>
 
 #include <rapidsmpf/buffer/resource.hpp>
+#include <rapidsmpf/nvtx.hpp>
 #include <rapidsmpf/streaming/core/context.hpp>
 #include <rapidsmpf/utils.hpp>
 
@@ -33,6 +34,7 @@ std::size_t spill_messages(
     std::shared_ptr<BufferResource> br,
     std::size_t amount
 ) {
+    RAPIDSMPF_NVTX_SCOPED_RANGE("spill_messages::amount", amount);
     // Recall that std::map is sorted by key by default, so iteration follows the
     // order in which messages were inserted into `spillable_messages`.
     std::map<SpillableMessages::MessageId, ContentDescription> cds =
@@ -45,9 +47,11 @@ std::size_t spill_messages(
             break;
         }
         if (cd.spillable()) {
+            RAPIDSMPF_NVTX_SCOPED_RANGE("spill_messages::spill", id);
             total_spilled += spillable_messages.spill(id, br.get());
         }
     }
+    RAPIDSMPF_NVTX_MARKER("spill_messages::total_spilled_bytes", total_spilled);
     return total_spilled;
 }
 }  // namespace