Fixed race condition on shared staging buffers in tree broadcast resharding.

tpu_raiden/core/raiden_manager_base.cc

@@ -284,18 +284,34 @@ absl::StatusOr<std::vector<int>> RaidenManagerBase::H2hReadDirect(
   return server_->SyncPull(peers, src_block_ids, {}, {}, parallelism_);
 }
 
-absl::Status RaidenManagerBase::PushWeightsChunk(absl::string_view peer,
-                                                 size_t dst_shard_idx,
-                                                 size_t dst_offset_bytes,
-                                                 const uint8_t* data_ptr,
-                                                 size_t size_bytes) {
+absl::Status RaidenManagerBase::PushWeightsChunk(
+    absl::string_view peer, size_t dst_shard_idx, size_t dst_offset_bytes,
+    const uint8_t* data_ptr, size_t size_bytes, uint64_t uuid) {
   InitTransportServer();
   absl::MutexLock lock(server_init_mu_);
   if (!server_) {
     return absl::FailedPreconditionError("Transport server is not running");
   }
   return server_->PushBuffer(peer, /*buffer_id=*/0, dst_shard_idx,
-                             dst_offset_bytes, data_ptr, size_bytes);
+                             dst_offset_bytes, data_ptr, size_bytes, uuid);
+}
+
+absl::Status RaidenManagerBase::RegisterExpectedChunks(
+    uint64_t uuid, uint32_t expected_chunks) {
+  InitTransportServer();
+  absl::MutexLock lock(server_init_mu_);
+  if (!server_) {
+    return absl::FailedPreconditionError("Transport server is not running");
+  }
+  return server_->RegisterExpectedChunks(uuid, expected_chunks);
+}
+
+void RaidenManagerBase::ForgetPushProgress(uint64_t uuid) {
+  InitTransportServer();
+  absl::MutexLock lock(server_init_mu_);
+  if (server_) {
+    server_->ForgetPushProgress(uuid);
+  }
 }
 
 size_t RaidenManagerBase::bytes_per_block() const { return slice_byte_size_; }

tpu_raiden/core/raiden_manager_base.h

@@ -79,7 +79,12 @@ class RaidenManagerBase : public tpu_raiden::transport::BlockTransportDelegate {
 
   absl::Status PushWeightsChunk(absl::string_view peer, size_t dst_shard_idx,
                                 size_t dst_offset_bytes,
-                                const uint8_t* data_ptr, size_t size_bytes);
+                                const uint8_t* data_ptr, size_t size_bytes,
+                                uint64_t uuid = 0);
+
+  absl::Status RegisterExpectedChunks(uint64_t uuid, uint32_t expected_chunks);
+
+  void ForgetPushProgress(uint64_t uuid);
 
   virtual std::optional<int> local_port() const;
   virtual std::string local_ip() const;

tpu_raiden/rpc/raiden_controller.py

@@ -2043,7 +2043,7 @@ async def _execute_slice_broadcast(
 
     available_sources = [src_unit]
     node_slice_offsets = {
-        src_unit: (src_block_id, src_block_offset, src_stride)
+        src_unit: (shard_idx, src_block_id, src_block_offset, src_stride)
     }
 
     pending_targets = list(targets)
@@ -2068,7 +2068,9 @@ async def _execute_slice_broadcast(
           active_pushes[s] += 1
           scheduled_any = True
 
-          s_block_id, s_block_offset, s_stride = node_slice_offsets[s]
+          s_shard_idx, s_block_id, s_block_offset, s_stride = (
+              node_slice_offsets[s]
+          )
 
           entry = (
               dst_peer,
@@ -2083,9 +2085,9 @@ async def _execute_slice_broadcast(
               count,
           )
 
-          sub_schedule = {s: {shard_idx if s == src_unit else 0: [entry]}}
           hop_uuid = random.randint(1, 2**63 - 1)
           hop_req_id = f"{req_id}_{hop_uuid}"
+          sub_schedule = {s: {s_shard_idx: [entry]}}
           sub_plan = TransferPlan(
               src_units=[s],
               dst_units=[dst_unit],
@@ -2097,7 +2099,7 @@ async def _execute_slice_broadcast(
               dst_mem_type=dst_mem_type,
               use_block_chunks=True,
               is_sender=True,
-              expected_block_count=1,
+              expected_block_count=count,
               req_id=hop_req_id,
           )
 
@@ -2115,7 +2117,7 @@ async def _run_single_transfer(s_node, d_node, plan):
                   [d_node],
                   plan.req_id,
                   True,
-                  False,
+                  s_node not in self._registered_shards,
                   plan.expected_block_count,
                   plan.uuid,
                   dst_controller_address,
@@ -2170,6 +2172,7 @@ async def _run_single_transfer(s_node, d_node, plan):
           active_pushes[s] -= 1
           available_sources.append(dst_unit)
           node_slice_offsets[dst_unit] = (
+              t[2],  # dst_shard_idx
               dst_block_id,
               dst_block_offset,
               dst_stride,
@@ -2905,16 +2908,22 @@ async def _execute_transfer() -> None:
           )
 
           # 4. Trigger COMMAND_START_TRANSFER (is_sender=False) on local workers
-          logging.info(
-              "Triggering preparation RPCs on local destination workers: %s,"
-              " expected blocks: %d",
-              local_dst_units,
-              expected_block_count,
-          )
-          await asyncio.gather(*[
-              self.worker_rpc_client.start_transfer(unit, receiver_plan)
-              for unit in local_dst_units
-          ])
+          if expected_block_count > 0:
+            logging.info(
+                "Triggering preparation RPCs on local destination workers: %s,"
+                " expected blocks: %d",
+                local_dst_units,
+                expected_block_count,
+            )
+            await asyncio.gather(*[
+                self.worker_rpc_client.start_transfer(unit, receiver_plan)
+                for unit in local_dst_units
+            ])
+          else:
+            logging.info(
+                "Skipping preparation RPCs on local destination workers because "
+                "expected_block_count is 0"
+            )
           logging.info(
               "Symmetric preparation complete on all local destination workers."
           )
@@ -3382,7 +3391,8 @@ async def _execute_transfer() -> None:
                       " controller",
                       req_id,
                   )
-              await asyncio.gather(*tree_broadcast_tasks)
+              for task in tree_broadcast_tasks:
+                await task
 
       else:
         # === OLD PLAN-BASED WORKFLOW (Backward Compatibility) ===
@@ -3838,7 +3848,7 @@ def decode_entries(schedule_proto):
               )
             else:
               shard_push_schedules = {}
-              if len(srcs) == 1 and len(start_req.shard_push_schedules) > 1:
+              if len(srcs) == 1:
                 unit_schedules = {}
                 for (
                     key_idx,
@@ -4263,7 +4273,7 @@ def register_transfer_schedule(
           else:
             key_idx = (
                 int(src_unit.job_replica_id)
-                if num_src_shards == 1
+                if (len(src_units) > 1 and num_src_shards == 1)
                 else shard_idx
             )
           schedule_proto = self._raiden_proto_module.ShardPushScheduleProto()

tpu_raiden/transport/block_transport.cc

@@ -913,6 +913,7 @@ void BlockTransport::H2hWriteWorker(int stream_idx, absl::string_view peer,
 }
 
 void BlockTransport::ForgetPushProgress(uint64_t uuid) {
+  RawBufferTransport::ForgetPushProgress(uuid);
   absl::MutexLock lock(progress_mu_);
   for (auto it = layer_progress_.begin(); it != layer_progress_.end();) {
     if (it->first.first == uuid) {

tpu_raiden/transport/block_transport.h

@@ -247,7 +247,7 @@ class BlockTransport : public lib::RawBufferTransport {
 
   // Drops receive-progress counters belonging to a finished, failed, or
   // timed-out plan so the UUID can be safely reused.
-  void ForgetPushProgress(uint64_t uuid);
+  void ForgetPushProgress(uint64_t uuid) override;
 
  private:
   absl::Status HandleCustomRequest(int client_fd,

tpu_raiden/transport/lib/raw_buffer_transport.cc

@@ -128,7 +128,7 @@ RawBufferTransport::~RawBufferTransport() {
     close(server_fd_);
   }
   {
-    absl::MutexLock _( mu_ );
+    absl::MutexLock _(mu_);
     for (int fd : active_client_fds_) {
       shutdown(fd, SHUT_RDWR);
       close(fd);
@@ -186,7 +186,7 @@ void RawBufferTransport::ReturnConnection(bool ok, int fd,
     return;
   }
 
-  absl::MutexLock lock( pool_mu_ );
+  absl::MutexLock lock(pool_mu_);
   if (stopping_) {
     shutdown(fd, SHUT_RDWR);
     close(fd);
@@ -198,7 +198,7 @@ void RawBufferTransport::ReturnConnection(bool ok, int fd,
 }
 
 void RawBufferTransport::ClosePooledConnections() {
-  absl::MutexLock lock( pool_mu_ );
+  absl::MutexLock lock(pool_mu_);
   for (auto& entry : conn_pool_) {
     for (int fd : entry.second) {
       shutdown(fd, SHUT_RDWR);
@@ -227,6 +227,30 @@ absl::Status RawBufferTransport::ProcessSingleRequest(int client_fd) {
     uint8_t* dest_ptr = base_host_ptr + dst_offset;
     RETURN_IF_ERROR(ReadExact(client_fd, dest_ptr, size_bytes));
 
+    bool trigger_h2d = false;
+    if (header.uuid > 0) {
+      absl::MutexLock lock(raw_progress_mu_);
+      auto& prog = raw_progress_[header.uuid];
+      prog.completed_chunks++;
+      VLOG(1) << "Received chunk for uuid=" << header.uuid
+              << " shard=" << dst_shard_idx << " offset=" << dst_offset
+              << " size=" << size_bytes << " progress=" << prog.completed_chunks
+              << "/"
+              << (prog.expected_chunks.has_value()
+                      ? std::to_string(*prog.expected_chunks)
+                      : "unknown");
+      if (prog.expected_chunks.has_value() &&
+          prog.completed_chunks == *prog.expected_chunks) {
+        raw_progress_.erase(header.uuid);
+        trigger_h2d = true;
+        VLOG(1) << "Triggering H2D for uuid=" << header.uuid;
+      }
+    }
+
+    if (trigger_h2d) {
+      RETURN_IF_ERROR(raw_delegate_->OnDataReceived());
+    }
+
     uint8_t ack = 1;
     RETURN_IF_ERROR(WriteExact(client_fd, &ack, 1));
   } else if (header.op == 3) {
@@ -277,7 +301,7 @@ void RawBufferTransport::ConnectionWorker(int client_fd) {
   }
   close(client_fd);
   {
-    absl::MutexLock _( mu_ );
+    absl::MutexLock _(mu_);
     active_client_fds_.erase(client_fd);
   }
 }
@@ -310,7 +334,7 @@ void RawBufferTransport::ListenerLoop() {
     setsockopt(client_fd, SOL_SOCKET, SO_RCVBUF, &buf_opt, sizeof(buf_opt));
 
     {
-      absl::MutexLock _( mu_ );
+      absl::MutexLock _(mu_);
       active_client_fds_.insert(client_fd);
     }
 
@@ -356,9 +380,40 @@ absl::Status RawBufferTransport::PullBuffer(
   return absl::OkStatus();
 }
 
-absl::Status RawBufferTransport::PushBuffer(
-    absl::string_view peer, size_t buffer_id, size_t dst_shard_idx,
-    size_t dst_offset_bytes, const uint8_t* data_ptr, size_t size_bytes) {
+absl::Status RawBufferTransport::RegisterExpectedChunks(
+    uint64_t uuid, uint32_t expected_chunks) {
+  if (expected_chunks == 0) {
+    return absl::InvalidArgumentError("expected_chunks must be positive");
+  }
+
+  bool trigger_h2d = false;
+  {
+    absl::MutexLock lock(raw_progress_mu_);
+    auto& prog = raw_progress_[uuid];
+    prog.expected_chunks = expected_chunks;
+    VLOG(1) << "RegisterExpectedChunks: uuid=" << uuid
+            << " expected_chunks=" << expected_chunks
+            << " completed_chunks=" << prog.completed_chunks;
+    if (prog.completed_chunks == expected_chunks) {
+      raw_progress_.erase(uuid);
+      trigger_h2d = true;
+      VLOG(1) << "RegisterExpectedChunks triggering H2D for uuid=" << uuid;
+    }
+  }
+
+  if (trigger_h2d) {
+    RETURN_IF_ERROR(raw_delegate_->OnDataReceived());
+  }
+
+  return absl::OkStatus();
+}
+
+absl::Status RawBufferTransport::PushBuffer(absl::string_view peer,
+                                            size_t buffer_id,
+                                            size_t dst_shard_idx,
+                                            size_t dst_offset_bytes,
+                                            const uint8_t* data_ptr,
+                                            size_t size_bytes, uint64_t uuid) {
   if (peer.empty()) {
     return absl::InvalidArgumentError(
         "Destination peer address cannot be empty");
@@ -375,7 +430,12 @@ absl::Status RawBufferTransport::PushBuffer(
   header.remote_id = static_cast<uint32_t>(dst_offset_bytes);
   header.local_id = static_cast<uint32_t>(dst_shard_idx);
   header.count_or_size = static_cast<uint32_t>(size_bytes);
+  header.uuid = uuid;
+  header.reserved = 0;
 
+  VLOG(1) << "Pushing chunk to peer=" << peer << " uuid=" << uuid
+          << " dst_shard=" << dst_shard_idx
+          << " dst_offset=" << dst_offset_bytes << " size=" << size_bytes;
   RETURN_IF_ERROR(WriteExact(fd, &header, sizeof(header)));
   RETURN_IF_ERROR(WriteExact(fd, data_ptr, size_bytes));
 
@@ -389,4 +449,9 @@ absl::Status RawBufferTransport::PushBuffer(
   return absl::OkStatus();
 }
 
+void RawBufferTransport::ForgetPushProgress(uint64_t uuid) {
+  absl::MutexLock lock(raw_progress_mu_);
+  raw_progress_.erase(uuid);
+}
+
 }  // namespace tpu_raiden::transport::lib

tpu_raiden/transport/lib/raw_buffer_transport.h

@@ -20,6 +20,7 @@
 #include <atomic>
 #include <cstddef>
 #include <cstdint>
+#include <optional>
 #include <string>
 #include <thread>  // NOLINT
 #include <vector>
@@ -34,6 +35,11 @@
 
 namespace tpu_raiden::transport::lib {
 
+struct RawProgress {
+  uint32_t completed_chunks = 0;
+  std::optional<uint32_t> expected_chunks;
+};
+
 // Foundational delegate interface for RawBufferTransport to query base host
 // memory.
 class RawBufferTransportDelegate {
@@ -76,7 +82,8 @@ class RawBufferTransport {
   // of a remote peer's buffer.
   absl::Status PushBuffer(absl::string_view peer, size_t buffer_id,
                           size_t dst_shard_idx, size_t dst_offset_bytes,
-                          const uint8_t* data_ptr, size_t size_bytes);
+                          const uint8_t* data_ptr, size_t size_bytes,
+                          uint64_t uuid = 0);
 
   // Synchronously requests an arbitrary continuous byte slice from a remote
   // peer's staging memory.
@@ -85,6 +92,10 @@ class RawBufferTransport {
                           size_t dst_shard_idx, size_t dst_offset_bytes,
                           size_t size_bytes);
 
+  absl::Status RegisterExpectedChunks(uint64_t uuid, uint32_t expected_chunks);
+
+  virtual void ForgetPushProgress(uint64_t uuid);
+
   int local_port() const { return local_port_; }
   const std::string& bound_ip() const { return bound_ip_; }
 
@@ -116,6 +127,10 @@ class RawBufferTransport {
   absl::flat_hash_map<std::string, std::vector<int>> conn_pool_
       ABSL_GUARDED_BY(pool_mu_);
 
+  absl::Mutex raw_progress_mu_;
+  absl::flat_hash_map<uint64_t, RawProgress> raw_progress_
+      ABSL_GUARDED_BY(raw_progress_mu_);
+
   std::thread listener_thread_;
   std::vector<std::thread> worker_threads_;
 };