fix(umbp): pool the gather streams so the standalone server stops leaking them

src/umbp/distributed/transfer/hbm_copy_engine.cpp

@@ -264,26 +264,55 @@ void LogPlanGeometry(const char* path, size_t plan_index, const TransferPlan& pl
   }
 }
 
-// One non-blocking stream per device, per thread.  Submit runs on PoolClient's
-// executor threads, so a shared stream would serialize copies that the
-// hipMemcpy path ran independently.
-std::map<int, hipStream_t>& GatherStreams() {
-  static thread_local std::map<int, hipStream_t>* streams = new std::map<int, hipStream_t>();
-  return *streams;
-}
+// Non-blocking gather streams, pooled per device and lent to one gather at a
+// time, so concurrent copies still run on separate streams.  They are not tied
+// to the calling thread: the standalone server calls Submit from gRPC's sync
+// pool, which retires and spawns threads under load, and every HIP stream holds
+// ~6.4 MiB of device memory for its kernarg pool.  Kept per thread, they leaked
+// until the GPU was full and stream creation failed (which HIP 7.2 turns into a
+// segfault).  A device's pool grows only to its peak number of concurrent
+// gathers.  Never destroyed: worker threads can still be copying during static
+// teardown.
+class GatherStreamPool {
+ public:
+  static GatherStreamPool& Instance() {
+    static GatherStreamPool* pool = new GatherStreamPool();
+    return *pool;
+  }
 
-hipStream_t GatherStream(int device_id) {
-  auto& streams = GatherStreams();
-  auto it = streams.find(device_id);
-  if (it != streams.end()) return it->second;
-  hipStream_t stream = nullptr;
-  if (hipStreamCreateWithFlags(&stream, hipStreamNonBlocking) != hipSuccess) {
-    (void)hipGetLastError();
-    return nullptr;
+  // The current device must already be `device_id`.  nullptr if no stream is
+  // free and a new one cannot be created.
+  hipStream_t Acquire(int device_id) {
+    {
+      std::lock_guard<std::mutex> lock(mutex_);
+      auto& idle = idle_[device_id];
+      if (!idle.empty()) {
+        hipStream_t stream = idle.back();
+        idle.pop_back();
+        return stream;
+      }
+    }
+    hipStream_t stream = nullptr;
+    if (hipStreamCreateWithFlags(&stream, hipStreamNonBlocking) != hipSuccess) {
+      (void)hipGetLastError();
+      return nullptr;
+    }
+    created_.fetch_add(1, std::memory_order_relaxed);
+    return stream;
   }
-  streams.emplace(device_id, stream);
-  return stream;
-}
+
+  void Release(int device_id, hipStream_t stream) {
+    std::lock_guard<std::mutex> lock(mutex_);
+    idle_[device_id].push_back(stream);
+  }
+
+  uint64_t Created() const { return created_.load(std::memory_order_relaxed); }
+
+ private:
+  std::mutex mutex_;
+  std::map<int, std::vector<hipStream_t>> idle_;
+  std::atomic<uint64_t> created_{0};
+};
 
 // Latches on the first gather failure and never tries again.  A kernel that
 // cannot launch here will not start working later, and the fallback is always
@@ -398,6 +427,8 @@ void* HbmCopyEngine::HostRegionDeviceAddress(const void* ptr, size_t size, int d
   return nullptr;
 }
 
+uint64_t HbmCopyEngine::GatherStreamsCreated() { return GatherStreamPool::Instance().Created(); }
+
 // Run the gather kernel over every eligible segment in the batch, and report
 // which plans it completed.
 //
@@ -558,7 +589,7 @@ std::vector<char> HbmCopyEngine::GatherEligiblePlans(const std::vector<TransferP
       note_bucket(GatherSkip::kSetDeviceFailed);
       continue;
     }
-    hipStream_t stream = GatherStream(device_id);
+    hipStream_t stream = GatherStreamPool::Instance().Acquire(device_id);
     if (stream == nullptr) {
       note_bucket(GatherSkip::kNoStream);
       continue;
@@ -569,21 +600,23 @@ std::vector<char> HbmCopyEngine::GatherEligiblePlans(const std::vector<TransferP
     // perturb what is being measured.
     const auto started = std::chrono::steady_clock::now();
     if (!LaunchDeviceGather(bucket.fragments.data(), bucket.fragments.size(), device_id, stream)) {
+      GatherStreamPool::Instance().Release(device_id, stream);
       note_bucket(GatherSkip::kLaunchFailed);
       continue;
     }
     const hipError_t sync = hipStreamSynchronize(stream);
     if (sync != hipSuccess) {
       // Latch off: a kernel that cannot complete here will not start working
-      // later, and every caller still has the fallback.
+      // later, and every caller still has the fallback.  The stream is dropped
+      // rather than returned to the pool.
       (void)hipGetLastError();
       GatherDisabled().store(true, std::memory_order_relaxed);
-      GatherStreams().erase(device_id);
       MORI_UMBP_ERROR("[HbmCopyEngine] gather sync dev={} failed: {}", device_id,
                       hipGetErrorString(sync));
       note_bucket(GatherSkip::kSyncFailed);
       continue;
     }
+    GatherStreamPool::Instance().Release(device_id, stream);
     for (size_t p : bucket.plan_indices) done[p] = 1;
     if (!debug) continue;
 

src/umbp/include/umbp/common/device_gather.h

@@ -53,6 +53,9 @@ struct DeviceGatherFragment {
 // cannot dereference plain mmap memory at all, it faults. Returns false before
 // enqueueing any work when the launch could not be issued, so the caller can
 // fall back to hipMemcpy.
+//
+// The descriptors are staged in a buffer kept per `stream`, so the caller must
+// not use `stream` from another thread until it has synchronized this launch.
 bool LaunchDeviceGather(const DeviceGatherFragment* fragments, size_t count, int device_id,
                         hipStream_t stream);
 

src/umbp/include/umbp/distributed/transfer/hbm_copy_engine.h

@@ -121,6 +121,12 @@ class HbmCopyEngine final : public TransferEngine {
   // where a multi-MiB copy dwarfs the per-call launch overhead.
   std::unique_ptr<TransferHandle> Submit(std::vector<TransferPlan> plans) override;
 
+  // Gather streams created since process start, across all engines.  They are
+  // pooled per device and outlive the thread that first used them, so this
+  // tracks the peak number of gathers in flight at once, not the number of
+  // threads that have called Submit.  Lets a test assert the reuse.
+  static uint64_t GatherStreamsCreated();
+
  private:
   // Why a plan did not take the gather kernel.  The hot path never computes
   // this: it is filled in only when debug mode is on (UMBP_HBM_COPY_DEBUG), so

src/umbp/local/tiers/device_gather.hip

@@ -26,8 +26,9 @@
 #include <cstdlib>
 #include <cstring>
 #include <limits>
-#include <map>
+#include <mutex>
 #include <string>
+#include <unordered_map>
 
 #include "mori/utils/mori_log.hpp"
 #include "umbp/common/device_gather.h"
@@ -71,25 +72,31 @@ __global__ void GatherFragmentsKernel(const DeviceGatherFragment* __restrict__ f
   for (size_t i = index; i < fragment.bytes; i += stride) dst[i] = src[i];
 }
 
-// Per-thread, per-device descriptor staging.
+// Per-stream descriptor staging.
 //
-// Per-thread because two threads serving the same device share one stream (see
-// DeviceStream), so a shared staging buffer would be a data race even though
-// the stream itself is ordered. Never freed, for the same reason DeviceStream
-// never destroys its streams: releasing HIP resources from a thread or static
-// destructor races with runtime teardown. The count is bounded by the server's
-// worker threads times the devices they touch.
+// The caller holds `stream` exclusively from the upload to the synchronize
+// (see LaunchDeviceGather), so a stream's buffer is never in use twice at once,
+// and the count is bounded by the caller's streams rather than by however many
+// threads have ever called in. Never freed, for the same reason the streams are
+// never destroyed: releasing HIP resources from a static destructor races with
+// runtime teardown.
 struct DescriptorStaging {
   DeviceGatherFragment* host = nullptr;
   DeviceGatherFragment* device = nullptr;
   size_t capacity = 0;
 };
 
-DescriptorStaging* StagingFor(int device_id, size_t count) {
-  static thread_local std::map<int, DescriptorStaging>* staging_by_device =
-      new std::map<int, DescriptorStaging>();
+DescriptorStaging* StagingFor(hipStream_t stream, int device_id, size_t count) {
+  static std::mutex* staging_mutex = new std::mutex();
+  static auto* staging_by_stream = new std::unordered_map<hipStream_t, DescriptorStaging>();
 
-  DescriptorStaging& staging = (*staging_by_device)[device_id];
+  DescriptorStaging* staging_entry = nullptr;
+  {
+    // References into an unordered_map stay valid across rehashing.
+    std::lock_guard<std::mutex> lock(*staging_mutex);
+    staging_entry = &(*staging_by_stream)[stream];
+  }
+  DescriptorStaging& staging = *staging_entry;
   if (staging.capacity >= count) return &staging;
 
   // Grow with headroom so a slowly increasing batch size does not reallocate
@@ -154,7 +161,7 @@ bool LaunchDeviceGather(const DeviceGatherFragment* fragments, size_t count, int
       (chunks + kThreadsPerBlock - 1) / kThreadsPerBlock, 1, kMaxBlocksPerFragment));
   if (count > std::numeric_limits<unsigned>::max() / blocks_per_fragment) return false;
 
-  DescriptorStaging* staging = StagingFor(device_id, count);
+  DescriptorStaging* staging = StagingFor(stream, device_id, count);
   if (staging == nullptr) return false;
 
   std::memcpy(staging->host, fragments, count * sizeof(DeviceGatherFragment));

tests/cpp/umbp/distributed/test_transfer_engine.cpp

@@ -26,6 +26,7 @@
 
 #include <fcntl.h>
 #include <gtest/gtest.h>
+#include <hip/hip_runtime_api.h>
 #include <unistd.h>
 
 #include <cstdint>
@@ -34,9 +35,11 @@
 #include <numeric>
 #include <set>
 #include <string>
+#include <thread>
 #include <vector>
 
 #include "umbp/common/aligned_buffer.h"
+#include "umbp/common/device_gather.h"
 #include "umbp/common/host_registration.h"
 #include "umbp/distributed/transfer/composite_transfer_engine.h"
 #include "umbp/distributed/transfer/hbm_copy_engine.h"
@@ -292,6 +295,59 @@ TEST(HbmCopyEngine, BuilderSplitsAtHostRegistrationPieces) {
   ExpectPlanContract(items, planned);
 }
 
+// Gathers submitted from a succession of short-lived threads, the way the
+// standalone server's gRPC pool retires and spawns threads under load, share a
+// pooled stream.  A stream per thread grew until HIP could not create another.
+TEST(HbmCopyEngine, GatherStreamsOutliveTheirThreads) {
+  int devices = 0;
+  if (hipGetDeviceCount(&devices) != hipSuccess || devices == 0) GTEST_SKIP() << "no GPU";
+  if (!DeviceGatherEnabled()) GTEST_SKIP() << "gather kernel disabled";
+
+  constexpr size_t kBytes = 1 << 20;
+  constexpr size_t kFragment = 4096;
+  constexpr int kThreads = 16;
+  AlignedBuffer host(kBytes);
+  for (size_t i = 0; i < kBytes; ++i) host.data()[i] = static_cast<char>(i * 131 + 7);
+  HbmCopyEngine engine;
+  engine.AddHostGatherRegion(host.data(), host.size());
+  void* device = nullptr;
+  ASSERT_EQ(hipMalloc(&device, kBytes), hipSuccess);
+  ASSERT_EQ(hipMemset(device, 0, kBytes), hipSuccess);
+  const auto src = TransferRef::HostBytes(host.data(), host.size());
+  const auto dst = TransferRef::HostBytes(device, kBytes, mori::io::MemoryLocationType::GPU, 0);
+
+  const uint64_t launches_before = DeviceGatherLaunchCount();
+  const uint64_t streams_before = HbmCopyEngine::GatherStreamsCreated();
+  for (int t = 0; t < kThreads; ++t) {
+    std::thread([&, t] {
+      // Two fragments far apart, so they stay two segments: one is hipMemcpy's case.
+      const uint64_t first = t * 2 * kFragment;
+      const uint64_t second = kBytes / 2 + first;
+      auto planned = engine.Plan({MakeItem(src, first, dst, first, kFragment, 1),
+                                  MakeItem(src, second, dst, second, kFragment, 2)});
+      std::vector<TransferFailure> failures;
+      if (auto handle = engine.Submit(std::move(planned.plans))) handle->Wait(&failures);
+      EXPECT_TRUE(failures.empty());
+    }).join();
+  }
+  if (DeviceGatherLaunchCount() == launches_before) {
+    (void)hipFree(device);
+    GTEST_SKIP() << "host region not gatherable on this device";
+  }
+  EXPECT_EQ(DeviceGatherLaunchCount(), launches_before + kThreads);
+  EXPECT_LE(HbmCopyEngine::GatherStreamsCreated() - streams_before, 1u);
+
+  std::vector<char> copied(kBytes);
+  ASSERT_EQ(hipMemcpy(copied.data(), device, kBytes, hipMemcpyDeviceToHost), hipSuccess);
+  for (int t = 0; t < kThreads; ++t) {
+    for (const size_t offset : {t * 2 * kFragment, kBytes / 2 + t * 2 * kFragment}) {
+      EXPECT_EQ(std::memcmp(copied.data() + offset, host.data() + offset, kFragment), 0)
+          << "thread " << t << " offset " << offset;
+    }
+  }
+  (void)hipFree(device);
+}
+
 // ---------------------------------------------------------------------------
 //  MoriIoEngine (planning only: the engine is never Init()ed)
 // ---------------------------------------------------------------------------