Add StatusOr to transfer server BulkTransportInterface on the bond id to

forward errors from bond connection failures to the control plane connection.

PiperOrigin-RevId: 820783819

Author: pschuh

xla/python/transfer/socket-server.cc

@@ -299,10 +299,14 @@ class SocketServer::SocketNetworkState : public PollEventLoop::Handler {
         msg.data = const_cast<void*>(data);
         msg.size = size;
         msg.on_send = [val = tsl::FormRef(this), offset, req_id, is_largest](
-                          int bond_id, size_t size) {
+                          absl::StatusOr<int> bond_id, size_t size) {
+          if (!bond_id.ok()) {
+            val->SendError(req_id, offset, size, is_largest, bond_id.status());
+            return;
+          }
           SocketTransferRequest response;
           auto* packet = response.mutable_packet();
-          packet->set_bulk_transport_id(bond_id);
+          packet->set_bulk_transport_id(*bond_id);
           packet->set_offset(offset);
           packet->set_size(size);
           packet->set_req_id(req_id);

xla/python/transfer/socket_bulk_transport.cc

@@ -180,6 +180,7 @@ class SendConnectionHandler : public PollEventLoop::Handler {
         artificial_send_limit_(artificial_send_limit) {}
 
   ~SendConnectionHandler() override {
+    msg_queue_->Poison(absl::InternalError("A send connection has failed."));
 #ifdef MSG_ZEROCOPY
     table_.ClearAll();
 #endif
@@ -356,6 +357,19 @@ std::shared_ptr<SharedSendWorkQueue> SharedSendWorkQueue::Start() {
   return result;
 }
 
+void SharedSendMsgQueue::Poison(absl::Status s) {
+  mu_.lock();
+  poison_status_ = s;
+  auto work_items = std::move(work_items_);
+  mu_.unlock();
+  while (!work_items.empty()) {
+    auto work = std::move(work_items.front());
+    work_items.pop_front();
+    std::move(work.on_send)(s, work.size);
+    std::move(work.on_done)();
+  }
+}
+
 void SharedSendMsgQueue::ReportReadyToSend(SendConnectionHandler* handler) {
   mu_.lock();
   if (!work_items_.empty()) {
@@ -375,6 +389,13 @@ void SharedSendMsgQueue::ReportReadyToSend(SendConnectionHandler* handler) {
 void SharedSendMsgQueue::ScheduleSendWork(
     aux::BulkTransportInterface::SendMessage msg) {
   mu_.lock();
+  if (!poison_status_.ok()) {
+    auto s = poison_status_;
+    mu_.unlock();
+    std::move(msg.on_send)(std::move(s), msg.size);
+    std::move(msg.on_done)();
+    return;
+  }
   DCHECK(!shutdown_);
   if (work_items_.empty() && !handlers_.empty()) {
     auto* handler = handlers_.front();

xla/python/transfer/socket_bulk_transport.h

@@ -129,13 +129,16 @@ class SharedSendMsgQueue {
       std::shared_ptr<SharedSendWorkQueue> work_queue,
       size_t artificial_send_limiti = std::numeric_limits<size_t>::max());
 
+  void Poison(absl::Status s);
+
  private:
   friend class SendConnectionHandler;
 
   void ReportReadyToSend(SendConnectionHandler* handler);
 
   absl::Mutex mu_;
   bool shutdown_ = false;
+  absl::Status poison_status_;
   std::deque<SendConnectionHandler*> handlers_;
   std::deque<aux::BulkTransportInterface::SendMessage> work_items_;
 };

xla/python/transfer/socket_bulk_transport_test.cc

@@ -85,7 +85,7 @@ TEST(SendQueue, TestZeroCopyQueueCleanRemoteShutdown) {
   BulkTransportInterface::SendMessage msg;
   msg.data = txt_msg.data();
   msg.size = txt_msg.size();
-  msg.on_send = [](int id, size_t size) {};
+  msg.on_send = [](absl::StatusOr<int> id, size_t size) {};
   msg.on_done = [&notify]() { notify.Notify(); };
   msg_queue->ScheduleSendWork(std::move(msg));
   notify.WaitForNotification();
@@ -124,7 +124,7 @@ TEST(SendQueue, SendAndRecvQueuesArtificialLimit) {
     BulkTransportInterface::SendMessage msg;
     msg.data = txt_msg.data();
     msg.size = txt_msg.size();
-    msg.on_send = [](int id, size_t size) {};
+    msg.on_send = [](absl::StatusOr<int> id, size_t size) {};
     msg.on_done = [&mu, &send_count]() {
       absl::MutexLock l(mu);
       --send_count;
@@ -230,9 +230,9 @@ TEST(SocketBulkTransportFactoryTest, SendAndRecvWithFactory) {
     BulkTransportInterface::SendMessage msg;
     msg.data = txt_msgs[i].data();
     msg.size = txt_msgs[i].size();
-    msg.on_send = [&, i](int id, size_t size) {
+    msg.on_send = [&, i](absl::StatusOr<int> id, size_t size) {
       absl::MutexLock l(mu);
-      send_queue.push_back({i, id});
+      send_queue.push_back({i, id.value()});
     };
     msg.on_done = [&mu, &send_count]() {
       absl::MutexLock l(mu);

xla/python/transfer/streaming.cc

@@ -151,7 +151,10 @@ BulkTransportInterface::SendMessage BulkTransportInterface::MakeMessage(
   SendMessage result;
   result.data = tmp->data();
   result.size = tmp->size();
-  result.on_send = std::move(on_send);
+  result.on_send = [on_send = std::move(on_send)](absl::StatusOr<int> bond_id,
+                                                  size_t size) mutable {
+    std::move(on_send)(bond_id.value(), size);
+  };
   result.on_done = [tmp = std::move(tmp)]() {};
   return result;
 }

xla/python/transfer/streaming.h

@@ -107,7 +107,8 @@ class BulkTransportInterface {
     // There may be some delay between Send() and when the message
     // is actually sent. on_send gets called when the message actually
     // gets sent.
-    absl::AnyInvocable<void(int bond_id, size_t size) &&> on_send;
+    absl::AnyInvocable<void(absl::StatusOr<int> bond_id, size_t size) &&>
+        on_send;
   };
 
   // Schedules a send over a BulkTransportInterface connection.

xla/python/transfer/streaming_ifrt.cc

@@ -180,7 +180,9 @@ void PremappedCopierState::StartWorkUnlocked(const WorkList& work_list) {
             --num_parallel_copies_;
             work_item->is_ready = true;
             work_item->result_status = s;
-            FlushReadyWorkItemsInOrder();
+            if (!currently_flushing_) {
+              FlushReadyWorkItemsInOrder();
+            }
             work_list2 = FindWorkLocked();
           }
           StartWorkUnlocked(work_list2);
@@ -194,14 +196,20 @@ void PremappedCopierState::FlushReadyWorkItemsInOrder() {
     if (!work_item->is_ready) {
       return;
     }
+    if (!work_item->result_status.ok()) {
+      available_copy_offsets_.push_back(work_item->dest_buffer);
+    }
+    currently_flushing_ = true;
+    mu_.unlock();
     if (work_item->result_status.ok()) {
       std::move(work_item->on_done)(this, work_item->dest_buffer,
                                     work_item->work);
     } else {
       std::move(work_item->on_done)(this, work_item->result_status,
                                     work_item->work);
-      available_copy_offsets_.push_back(work_item->dest_buffer);
     }
+    mu_.lock();
+    currently_flushing_ = false;
     work_queue_.pop_front();
     ++base_seq_id_;
   }

xla/python/transfer/streaming_ifrt.h

@@ -117,6 +117,7 @@ class PremappedCopierState
   size_t num_parallel_copies_ = 0;
   std::deque<WorkQueueItem> work_queue_ ABSL_GUARDED_BY(mu_);
   std::shared_ptr<absl::Span<uint8_t>> scratch_;
+  bool currently_flushing_ ABSL_GUARDED_BY(mu_) = false;
   size_t max_num_parallel_copies_;
   size_t xfer_size_;
   size_t max_copies_;