replace raw copy threads with a persistent copy_executor_

fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py

@@ -137,6 +137,9 @@ class UVMCacheStatsIndex(enum.IntEnum):
 class RESParams:
     res_server_port: int = 0  # the port of the res server
     res_store_shards: int = 1  # the number of shards to store the raw embeddings
+    res_chunk_size: int = 500000  # max rows copied into one enqueued chunk
+    res_num_consumers: int = 8  # threads draining the stream queue
+    res_num_copy_threads: int = 4  # parallel chunk-copy threads per stream() call
     table_names: list[str] = field(default_factory=list)  # table names the TBE holds
     table_offsets: list[int] = field(
         default_factory=list
@@ -1560,6 +1563,9 @@ def __init__(  # noqa C901
                 self.enable_raw_embedding_streaming,
                 self.res_params.res_store_shards,
                 self.res_params.res_server_port,
+                self.res_params.res_chunk_size,
+                self.res_params.res_num_consumers,
+                self.res_params.res_num_copy_threads,
                 self.res_params.table_names,
                 self.res_params.table_offsets,
                 self.res_params.table_sizes,
@@ -4529,7 +4535,7 @@ def raw_embedding_stream(self) -> None:
             f"## uvm_lookup_prefetched_rows {self.timestep} {self.uuid} ##"
         ):
             if not self._res_sync_copy and self._res_require_copy:
-                self._raw_embedding_streamer.join_stream_tensor_copy_thread()
+                self._raw_embedding_streamer.join_dispatch()
             prefetched_info = self.prefetched_info_list.pop(0)
             updated_locations = torch.ops.fbgemm.lxu_cache_lookup(
                 prefetched_info.linear_unique_cache_indices,
@@ -4580,7 +4586,7 @@ def raw_embedding_stream(self) -> None:
                         # Lazy resize: runtime_meta shape/dtype is not known until
                         # the first data arrives from the MC module. Must use UVM
                         # (new_unified_tensor) because the C++ RawEmbeddingStreamer
-                        # reads this buffer via raw CPU pointers in tensor_copy().
+                        # reads this buffer via raw CPU pointers in tensor_copy_chunk().
                         self.register_buffer(
                             "res_runtime_meta",
                             torch.ops.fbgemm.new_unified_tensor(

fbgemm_gpu/include/fbgemm_gpu/split_embeddings_cache/raw_embedding_streamer.h

@@ -10,11 +10,17 @@
 #include <ATen/ATen.h>
 #ifdef FBGEMM_FBCODE
 #include <folly/coro/Task.h>
+#include <folly/futures/Future.h>
 #endif
 
 #include <utility>
+#include <vector>
 
 #ifdef FBGEMM_FBCODE
+namespace folly {
+class CPUThreadPoolExecutor;
+} // namespace folly
+
 namespace facebook::aiplatform::gmpp::experimental::training_ps {
 class TrainingPsOdsLogger;
 } // namespace facebook::aiplatform::gmpp::experimental::training_ps
@@ -48,6 +54,9 @@ class RawEmbeddingStreamer : public torch::jit::CustomClassHolder {
       bool enable_raw_embedding_streaming,
       int64_t res_store_shards,
       int64_t res_server_port,
+      int64_t res_chunk_size,
+      int64_t res_num_consumers,
+      int64_t res_num_copy_threads,
       std::vector<std::string> table_names,
       std::vector<int64_t> table_offsets,
       const std::vector<int64_t>& table_sizes);
@@ -56,10 +65,12 @@ class RawEmbeddingStreamer : public torch::jit::CustomClassHolder {
 
   /// Stream out non-negative elements in <indices> and its paired embeddings
   /// from <weights> for the first <count> elements in the tensor.
-  /// It spins up a thread that will copy all 3 tensors to CPU and inject them
-  /// into the background queue which will be picked up by another set of thread
-  /// pools for streaming out to the thrift server (co-located on same host
-  /// now).
+  /// It spins up a dispatcher thread that copies the 4 tensors (indices,
+  /// weights, and the optional identities / runtime_meta) to CPU and injects
+  /// them into the background queue, which is drained by a pool of consumer
+  /// threads that stream out to the thrift server (co-located on same host
+  /// now). The copy is split into <= res_chunk_size-row chunks that run
+  /// concurrently across the res_num_copy_threads-worker copy pool.
   ///
   /// This is used in cuda stream callback, which doesn't require to be
   /// serialized with other callbacks, thus a separate thread is used to
@@ -86,64 +97,125 @@ class RawEmbeddingStreamer : public torch::jit::CustomClassHolder {
       std::optional<at::Tensor> copy_done_flag = std::nullopt);
 
   /*
-   * Join the stream tensor copy thread, make sure the thread is properly
+   * Join the pending dispatch future (which resolves only after all copies for
+   * the previous iteration have been enqueued), making sure it is properly
    * finished before creating new.
    */
-  void join_stream_tensor_copy_thread();
+  void join_dispatch();
 
 #ifdef FBGEMM_FBCODE
   folly::coro::Task<void> tensor_stream(
       const at::Tensor& indices,
       const at::Tensor& weights,
       std::optional<at::Tensor> identities,
       std::optional<at::Tensor> runtime_meta);
-  /*
-   * Copy the indices, weights and count tensors and enqueue them for
-   * asynchronous stream.
-   */
-  void copy_and_enqueue_stream_tensors(
-      const at::Tensor& indices,
-      const at::Tensor& weights,
-      std::optional<at::Tensor> identities,
-      std::optional<at::Tensor> runtime_meta,
-      const at::Tensor& count);
 
   /*
-   * FOR TESTING: Join the weight stream thread, make sure the thread is
-   * properly finished for destruction and testing.
+   * FOR TESTING ONLY: drops the ship executor to 0 worker threads so a test can
+   * read a stable queue size (via get_weights_to_stream_queue_size()). Not
+   * reversible -- the executor stops shipping after this call.
    */
   void join_weights_stream_thread();
   // FOR TESTING: get queue size.
   uint64_t get_weights_to_stream_queue_size();
 #endif
  private:
-  std::atomic<bool> stop_{false};
   std::string unique_id_;
   bool enable_raw_embedding_streaming_;
 #ifdef FBGEMM_FBCODE
   int64_t res_store_shards_;
   int64_t res_server_port_;
+  size_t res_chunk_size_;
+  size_t res_num_consumers_;
+  size_t res_num_copy_threads_;
 #endif
   std::vector<std::string> table_names_;
   std::vector<int64_t> table_offsets_;
   at::Tensor table_sizes_;
 #ifdef FBGEMM_FBCODE
-  std::unique_ptr<std::thread> weights_stream_thread_;
-  folly::UMPSCQueue<StreamQueueItem, true> weights_to_stream_queue_;
-  std::unique_ptr<std::thread> stream_tensor_copy_thread_;
+  // Named executor that ships enqueued StreamQueueItems to the PS. Push model:
+  // producers submit one ship task per item and workers wake on submit (no
+  // polling). Sized to res_num_consumers_.
+  std::unique_ptr<folly::CPUThreadPoolExecutor> consumer_executor_;
+  // Persistent pool that runs the per-chunk tensor copies (CPU-bound
+  // std::copy), sized res_num_copy_threads_ so chunks run concurrently. Shared
+  // by the blocking and non-blocking stream() paths; this assumes a given table
+  // streams in a single mode at a time (blocking OR non-blocking), never
+  // concurrently.
+  std::unique_ptr<folly::CPUThreadPoolExecutor> copy_executor_;
+  // Persistent size-1 executor that runs the per-iteration dispatch (poll_flag
+  // + chunked_copy_and_enqueue). Named so its thread is identifiable in traces.
+  std::unique_ptr<folly::CPUThreadPoolExecutor> dispatch_executor_;
+  folly::SemiFuture<folly::Unit> dispatch_future_{folly::makeSemiFuture()};
   // OBC logger for RES silent-failure counters (res.fail.*). Emits to the
   // host-level OBC agent, so it reaches ODS from the trainer process without
   // per-process fb303 scrape config. Only constructed when streaming is on.
   std::unique_ptr<facebook::aiplatform::gmpp::experimental::training_ps::
                       TrainingPsOdsLogger>
       ods_logger_;
+
+  // Submit one ship task (blockingWait(tensor_stream(...))) for `item` onto
+  // consumer_executor_.
+  void submit_stream_item(StreamQueueItem item);
+  // Copies `count` rows of the source tensors into CPU chunks and enqueues each
+  // via submit_stream_item, fanning the per-chunk copies out across
+  // copy_executor_ and awaiting them all (the torn-row barrier). A coroutine so
+  // the non-blocking dispatch can co_await it; args are by value so nothing
+  // dangles once a chunk is scheduled on copy_executor_.
+  folly::coro::Task<void> chunked_copy_and_enqueue(
+      at::Tensor indices,
+      at::Tensor weights,
+      std::optional<at::Tensor> identities,
+      std::optional<at::Tensor> runtime_meta,
+      at::Tensor count);
+  // Copies one chunk [start, end) and enqueues it via submit_stream_item. Runs
+  // to completion on a single copy_executor_ worker; the per-chunk try/catch
+  // keeps an exception from escaping into collectAllRange (which awaits every
+  // sibling to completion and then rethrows one exception onto the blocking
+  // caller / dispatch future). By value for the same reason as
+  // chunked_copy_and_enqueue.
+  folly::coro::Task<void> copy_chunk_task(
+      at::Tensor indices,
+      at::Tensor weights,
+      std::optional<at::Tensor> identities,
+      std::optional<at::Tensor> runtime_meta,
+      int64_t start,
+      int64_t end);
+
+  // Waits (spinning) for copy_done_flag to signal the source tensors are safe
+  // to read, resetting it. Returns false on timeout. Shared by the blocking
+  // stream() path and dispatch_copy_task.
+  bool poll_copy_done_flag(const std::optional<at::Tensor>& copy_done_flag);
+
+  // Coroutine form of the non-blocking dispatch (poll_copy_done_flag +
+  // chunked_copy_and_enqueue). Args are taken by value so nothing dangles once
+  // it is scheduled on dispatch_executor_ -- a capturing lambda coroutine would
+  // risk a use-after-free (clang-tidy cppcoreguidelines-avoid-capturing-lambda-
+  // coroutines).
+  folly::coro::Task<void> dispatch_copy_task(
+      at::Tensor indices,
+      at::Tensor weights,
+      std::optional<at::Tensor> identities,
+      std::optional<at::Tensor> runtime_meta,
+      at::Tensor count,
+      std::optional<at::Tensor> copy_done_flag);
 #endif
 };
 
-fbgemm_gpu::StreamQueueItem tensor_copy(
+fbgemm_gpu::StreamQueueItem tensor_copy_chunk(
     const at::Tensor& indices,
     const at::Tensor& weights,
     std::optional<at::Tensor> identities,
     std::optional<at::Tensor> runtime_meta,
-    const at::Tensor& count);
+    int64_t start_row,
+    int64_t end_row);
+
+// Tiles [0, num_rows) into flat [start, end) chunks of size <= chunk_size,
+// contiguous and non-overlapping (their union is the whole range). One task per
+// chunk is submitted to copy_executor_, which load-balances them across its
+// workers. Pure/build-agnostic so it is unit-testable without a GPU or
+// FBGEMM_FBCODE.
+std::vector<std::pair<int64_t, int64_t>> computeChunks(
+    int64_t num_rows,
+    size_t chunk_size);
 } // namespace fbgemm_gpu