[core] Fix RAY_CHECK failure during shutdown due to plasma store race condition (#55367)

codope · dstrodtman · commit af6a1dcb9674 · 2025-10-06T16:34:29.000-04:00
## Why are these changes needed? Workers crash with a fatal `RAY_CHECK` failure when the plasma store connection is broken during shutdown, causing the following error: ``` RAY_CHECK failed: PutInLocalPlasmaStore(object, object_id, true) Status not OK: IOError: Broken pipe ``` Stacktrace: ``` core_worker.cc:720 C Check failed: PutInLocalPlasmaStore(object, object_id, true) Status not OK: IOError: Broken pipe *** StackTrace Information *** /home/ray/anaconda3/lib/python3.11/site-packages/ray/_raylet.so(+0x141789a) [0x7924dd2c689a] ray::operator<<() /home/ray/anaconda3/lib/python3.11/site-packages/ray/_raylet.so(_ZN3ray6RayLogD1Ev+0x479) [0x7924dd2c9319] ray::RayLog::~RayLog() /home/ray/anaconda3/lib/python3.11/site-packages/ray/_raylet.so(+0x95cc8a) [0x7924dc80bc8a] ray::core::CoreWorker::CoreWorker()::{lambda()#13}::operator()() /home/ray/anaconda3/lib/python3.11/site-packages/ray/_raylet.so(_ZN3ray4core11TaskManager27MarkTaskReturnObjectsFailedERKNS_17TaskSpecificationENS_3rpc9ErrorTypeEPKNS5_12RayErrorInfoERKN4absl12lts_2023080213flat_hash_setINS_8ObjectIDENSB_13hash_internal4HashISD_EESt8equal_toISD_ESaISD_EEE+0x679) [0x7924dc868f29] ray::core::TaskManager::MarkTaskReturnObjectsFailed() /home/ray/anaconda3/lib/python3.11/site-packages/ray/_raylet.so(_ZN3ray4core11TaskManager15FailPendingTaskERKNS_6TaskIDENS_3rpc9ErrorTypeEPKNS_6StatusEPKNS5_12RayErrorInfoE+0x416) [0x7924dc86f186] ray::core::TaskManager::FailPendingTask() /home/ray/anaconda3/lib/python3.11/site-packages/ray/_raylet.so(+0x9a90e6) [0x7924dc8580e6] ray::core::NormalTaskSubmitter::RequestNewWorkerIfNeeded()::{lambda()#1}::operator()() /home/ray/anaconda3/lib/python3.11/site-packages/ray/_raylet.so(_ZN3ray3rpc14ClientCallImplINS0_23RequestWorkerLeaseReplyEE15OnReplyReceivedEv+0x68) [0x7924dc94aa48] ray::rpc::ClientCallImpl<>::OnReplyReceived() /home/ray/anaconda3/lib/python3.11/site-packages/ray/_raylet.so(_ZNSt17_Function_handlerIFvvEZN3ray3rpc17ClientCallManager29PollEventsFromCompletionQueueEiEUlvE_E9_M_invokeERKSt9_Any_data+0x15) [0x7924dc79e285] std::_Function_handler<>::_M_invoke() /home/ray/anaconda3/lib/python3.11/site-packages/ray/_raylet.so(+0xd9b4c8) [0x7924dcc4a4c8] EventTracker::RecordExecution() /home/ray/anaconda3/lib/python3.11/site-packages/ray/_raylet.so(+0xd4648e) [0x7924dcbf548e] std::_Function_handler<>::_M_invoke() /home/ray/anaconda3/lib/python3.11/site-packages/ray/_raylet.so(+0xd46906) [0x7924dcbf5906] boost::asio::detail::completion_handler<>::do_complete() /home/ray/anaconda3/lib/python3.11/site-packages/ray/_raylet.so(+0x13f417b) [0x7924dd2a317b] boost::asio::detail::scheduler::do_run_one() /home/ray/anaconda3/lib/python3.11/site-packages/ray/_raylet.so(+0x13f5af9) [0x7924dd2a4af9] boost::asio::detail::scheduler::run() /home/ray/anaconda3/lib/python3.11/site-packages/ray/_raylet.so(+0x13f6202) [0x7924dd2a5202] boost::asio::io_context::run() /home/ray/anaconda3/lib/python3.11/site-packages/ray/_raylet.so(_ZN3ray4core10CoreWorker12RunIOServiceEv+0x91) [0x7924dc793a61] ray::core::CoreWorker::RunIOService() /home/ray/anaconda3/lib/python3.11/site-packages/ray/_raylet.so(+0xcba0b0) [0x7924dcb690b0] thread_proxy /lib/x86_64-linux-gnu/libc.so.6(+0x94ac3) [0x7924dde71ac3] /lib/x86_64-linux-gnu/libc.so.6(+0x126850) [0x7924ddf03850] ``` Stack trace flow: 1. Task lease request fails -> `NormalTaskSubmitter::RequestNewWorkerIfNeeded()` callback. 2. Triggers `TaskManager::FailPendingTask()` -> `MarkTaskReturnObjectsFailed()`. 3. System attempts to store error objects in plasma via `put_in_local_plasma_callback_`. 4. Plasma connection is broken (raylet/plasma store already shut down). 5. `RAY_CHECK_OK()` in the callback causes fatal crash instead of graceful handling. Root Cause: This is a shutdown ordering race condition: 1. Raylet shuts down first: The raylet stops its IO context ([main_service_.stop()](https://github.com/ray-project/ray/blob/77c5475195e56a26891d88460973198391d20edf/src/ray/object_manager/plasma/store_runner.cc#L146)) which closes plasma store connections. 2. Worker still processes callbacks: Core worker continues processing pending callbacks on separate threads. 3. Broken connection: When the callback tries to store error objects in plasma, the connection is already closed. 4. Fatal crash: The `RAY_CHECK_OK()` treats this as an unexpected error and crashes the process. Fix: 1. Shutdown-aware plasma operations - Add `CoreWorker::IsShuttingDown()` method to check shutdown state. - Skip plasma operations entirely when shutdown is in progress. - Prevents attempting operations on already-closed connections. 2. Targeted error handling for connection failures - Replace blanket `RAY_CHECK_OK()` with specific error type checking. - Handle connection errors (Broken pipe, Connection reset, Bad file descriptor) as warnings during shutdown scenarios. - Maintain `RAY_CHECK_OK()` for other error types to catch real issues. --------- Signed-off-by: Sagar Sumit <sagarsumit09@gmail.com> Signed-off-by: Douglas Strodtman <douglas@anyscale.com>
diff --git a/src/ray/core_worker/core_worker.h b/src/ray/core_worker/core_worker.h
@@ -1494,6 +1494,19 @@ class CoreWorker {
       std::string *application_error);
 
   /// Put an object in the local plasma store.
+  ///
+  /// Return status semantics:
+  /// - Status::OK(): The object was created (or already existed) and bookkeeping was
+  ///   updated. Note: an internal ObjectExists from the plasma provider is treated
+  ///   as OK and does not surface here.
+  /// - Status::ObjectStoreFull(): The local plasma store is out of memory (or out of
+  ///   disk when spilling). The error message contains context and a short memory
+  ///   report.
+  /// - Status::IOError(): IPC/connection failures while talking to the plasma store
+  ///   (e.g., broken pipe/connection reset during shutdown, store not reachable).
+  ///
+  /// Call sites that run during shutdown may choose to tolerate IOError specifically,
+  /// but should treat all other statuses as real failures.
   Status PutInLocalPlasmaStore(const RayObject &object,
                                const ObjectID &object_id,
                                bool pin_object);
diff --git a/src/ray/core_worker/core_worker_process.cc b/src/ray/core_worker/core_worker_process.cc
@@ -423,8 +423,14 @@ std::shared_ptr<CoreWorker> CoreWorkerProcessImpl::CreateCoreWorker(
       /*put_in_local_plasma_callback=*/
       [this](const RayObject &object, const ObjectID &object_id) {
         auto core_worker = GetCoreWorker();
-        RAY_CHECK_OK(
-            core_worker->PutInLocalPlasmaStore(object, object_id, /*pin_object=*/true));
+        auto put_status =
+            core_worker->PutInLocalPlasmaStore(object, object_id, /*pin_object=*/true);
+        if (!put_status.ok()) {
+          RAY_LOG(WARNING).WithField(object_id)
+              << "Failed to put object in plasma store: " << put_status;
+          return put_status;
+        }
+        return Status::OK();
       },
       /* retry_task_callback= */
       [this](TaskSpecification &spec, bool object_recovery, uint32_t delay_ms) {
diff --git a/src/ray/core_worker/task_manager.cc b/src/ray/core_worker/task_manager.cc
@@ -533,10 +533,10 @@ size_t TaskManager::NumPendingTasks() const {
   return num_pending_tasks_;
 }
 
-bool TaskManager::HandleTaskReturn(const ObjectID &object_id,
-                                   const rpc::ReturnObject &return_object,
-                                   const NodeID &worker_node_id,
-                                   bool store_in_plasma) {
+StatusOr<bool> TaskManager::HandleTaskReturn(const ObjectID &object_id,
+                                             const rpc::ReturnObject &return_object,
+                                             const NodeID &worker_node_id,
+                                             bool store_in_plasma) {
   bool direct_return = false;
   reference_counter_.UpdateObjectSize(object_id, return_object.size());
   RAY_LOG(DEBUG) << "Task return object " << object_id << " has size "
@@ -579,7 +579,15 @@ bool TaskManager::HandleTaskReturn(const ObjectID &object_id,
                      /*copy_data=*/false,
                      tensor_transport.value_or(rpc::TensorTransport::OBJECT_STORE));
     if (store_in_plasma) {
-      put_in_local_plasma_callback_(object, object_id);
+      Status s = put_in_local_plasma_callback_(object, object_id);
+      int retry_count = 0;
+      while (!s.ok() && s.IsTransientObjectStoreFull() && retry_count < 3) {
+        retry_count++;
+        s = put_in_local_plasma_callback_(object, object_id);
+      }
+      if (!s.ok()) {
+        return s;
+      }
     } else {
       in_memory_store_.Put(object, object_id);
       direct_return = true;
@@ -813,10 +821,15 @@ bool TaskManager::HandleReportGeneratorItemReturns(
     }
     // When an object is reported, the object is ready to be fetched.
     reference_counter_.UpdateObjectPendingCreation(object_id, false);
-    HandleTaskReturn(object_id,
-                     return_object,
-                     NodeID::FromBinary(request.worker_addr().node_id()),
-                     /*store_in_plasma=*/store_in_plasma_ids.contains(object_id));
+    StatusOr<bool> put_res =
+        HandleTaskReturn(object_id,
+                         return_object,
+                         NodeID::FromBinary(request.worker_addr().node_id()),
+                         /*store_in_plasma=*/store_in_plasma_ids.contains(object_id));
+    if (!put_res.ok()) {
+      RAY_LOG(WARNING).WithField(object_id)
+          << "Failed to handle streaming dynamic return: " << put_res.status();
+    }
   }
 
   // Handle backpressure if needed.
@@ -900,23 +913,54 @@ void TaskManager::CompletePendingTask(const TaskID &task_id,
         reference_counter_.AddDynamicReturn(object_id, generator_id);
         dynamic_return_ids.push_back(object_id);
       }
-      if (!HandleTaskReturn(object_id,
-                            return_object,
-                            NodeID::FromBinary(worker_addr.node_id()),
-                            store_in_plasma_ids.contains(object_id))) {
-        if (first_execution) {
-          dynamic_returns_in_plasma.push_back(object_id);
+      StatusOr<bool> direct_or =
+          HandleTaskReturn(object_id,
+                           return_object,
+                           NodeID::FromBinary(worker_addr.node_id()),
+                           store_in_plasma_ids.contains(object_id));
+      if (!direct_or.ok()) {
+        RAY_LOG(WARNING).WithField(object_id)
+            << "Failed to handle dynamic task return: " << direct_or.status();
+        Status st = direct_or.status();
+        rpc::ErrorType err_type = rpc::ErrorType::WORKER_DIED;
+        if (st.IsObjectStoreFull() || st.IsTransientObjectStoreFull()) {
+          err_type = rpc::ErrorType::OUT_OF_MEMORY;
         }
+        rpc::RayErrorInfo err_info;
+        err_info.set_error_message(st.ToString());
+        FailOrRetryPendingTask(task_id,
+                               err_type,
+                               &st,
+                               /*ray_error_info=*/&err_info,
+                               /*mark_task_object_failed=*/true,
+                               /*fail_immediately=*/true);
+        return;
+      } else if (!direct_or.value() && first_execution) {
+        dynamic_returns_in_plasma.push_back(object_id);
       }
     }
   }
 
   for (const auto &return_object : reply.return_objects()) {
     const auto object_id = ObjectID::FromBinary(return_object.object_id());
-    if (HandleTaskReturn(object_id,
-                         return_object,
-                         NodeID::FromBinary(worker_addr.node_id()),
-                         store_in_plasma_ids.contains(object_id))) {
+    StatusOr<bool> direct_or = HandleTaskReturn(object_id,
+                                                return_object,
+                                                NodeID::FromBinary(worker_addr.node_id()),
+                                                store_in_plasma_ids.contains(object_id));
+    if (!direct_or.ok()) {
+      RAY_LOG(WARNING).WithField(object_id)
+          << "Failed to handle task return: " << direct_or.status();
+      // If storing return in plasma failed, treat as system failure for this attempt.
+      // Do not proceed with normal completion. Mark task failed immediately.
+      Status st = direct_or.status();
+      FailOrRetryPendingTask(task_id,
+                             rpc::ErrorType::WORKER_DIED,
+                             &st,
+                             /*ray_error_info=*/nullptr,
+                             /*mark_task_object_failed=*/true,
+                             /*fail_immediately=*/true);
+      return;
+    } else if (direct_or.value()) {
       direct_return_ids.push_back(object_id);
     }
   }
@@ -1040,10 +1084,16 @@ void TaskManager::CompletePendingTask(const TaskID &task_id,
           const auto generator_return_id = spec.StreamingGeneratorReturnId(i);
           RAY_CHECK_EQ(reply.return_objects_size(), 1);
           const auto &return_object = reply.return_objects(0);
-          HandleTaskReturn(generator_return_id,
-                           return_object,
-                           NodeID::FromBinary(worker_addr.node_id()),
-                           store_in_plasma_ids.contains(generator_return_id));
+          StatusOr<bool> res =
+              HandleTaskReturn(generator_return_id,
+                               return_object,
+                               NodeID::FromBinary(worker_addr.node_id()),
+                               store_in_plasma_ids.contains(generator_return_id));
+          if (!res.ok()) {
+            RAY_LOG(WARNING).WithField(generator_return_id)
+                << "Failed to handle generator return during app error propagation: "
+                << res.status();
+          }
         }
       }
     }
@@ -1454,18 +1504,26 @@ void TaskManager::MarkTaskReturnObjectsFailed(
   int64_t num_returns = spec.NumReturns();
   for (int i = 0; i < num_returns; i++) {
     const auto object_id = ObjectID::FromIndex(task_id, /*index=*/i + 1);
+    // Always place an error marker in local memory to unblock waiters quickly.
+    in_memory_store_.Put(error, object_id);
+    // Best-effort plasma put if the object was meant to be in plasma.
     if (store_in_plasma_ids.contains(object_id)) {
-      put_in_local_plasma_callback_(error, object_id);
-    } else {
-      in_memory_store_.Put(error, object_id);
+      Status s = put_in_local_plasma_callback_(error, object_id);
+      if (!s.ok()) {
+        RAY_LOG(WARNING).WithField(object_id)
+            << "Failed to put error object in plasma: " << s;
+      }
     }
   }
   if (spec.ReturnsDynamic()) {
     for (const auto &dynamic_return_id : spec.DynamicReturnIds()) {
+      in_memory_store_.Put(error, dynamic_return_id);
       if (store_in_plasma_ids.contains(dynamic_return_id)) {
-        put_in_local_plasma_callback_(error, dynamic_return_id);
-      } else {
-        in_memory_store_.Put(error, dynamic_return_id);
+        Status s = put_in_local_plasma_callback_(error, dynamic_return_id);
+        if (!s.ok()) {
+          RAY_LOG(WARNING).WithField(dynamic_return_id)
+              << "Failed to put error object in plasma: " << s;
+        }
       }
     }
   }
@@ -1488,7 +1546,11 @@ void TaskManager::MarkTaskReturnObjectsFailed(
     for (size_t i = 0; i < num_streaming_generator_returns; i++) {
       const auto generator_return_id = spec.StreamingGeneratorReturnId(i);
       if (store_in_plasma_ids.contains(generator_return_id)) {
-        put_in_local_plasma_callback_(error, generator_return_id);
+        Status s = put_in_local_plasma_callback_(error, generator_return_id);
+        if (!s.ok()) {
+          RAY_LOG(WARNING).WithField(generator_return_id)
+              << "Failed to put error object in plasma: " << s;
+        }
       } else {
         in_memory_store_.Put(error, generator_return_id);
       }
diff --git a/src/ray/core_worker/task_manager.h b/src/ray/core_worker/task_manager.h
@@ -14,6 +14,7 @@
 
 #pragma once
 
+#include <functional>
 #include <memory>
 #include <string>
 #include <tuple>
@@ -25,6 +26,7 @@
 #include "absl/container/flat_hash_map.h"
 #include "absl/synchronization/mutex.h"
 #include "ray/common/id.h"
+#include "ray/common/status.h"
 #include "ray/core_worker/store_provider/memory_store/memory_store.h"
 #include "ray/core_worker/task_event_buffer.h"
 #include "ray/core_worker/task_manager_interface.h"
@@ -42,7 +44,7 @@ class ActorManager;
 
 using TaskStatusCounter = CounterMap<std::tuple<std::string, rpc::TaskStatus, bool>>;
 using PutInLocalPlasmaCallback =
-    std::function<void(const RayObject &object, const ObjectID &object_id)>;
+    std::function<Status(const RayObject &object, const ObjectID &object_id)>;
 using RetryTaskCallback =
     std::function<void(TaskSpecification &spec, bool object_recovery, uint32_t delay_ms)>;
 using ReconstructObjectCallback = std::function<void(const ObjectID &object_id)>;
@@ -608,12 +610,12 @@ class TaskManager : public TaskManagerInterface {
       ABSL_LOCKS_EXCLUDED(mu_);
 
   /// Update nested ref count info and store the in-memory value for a task's
-  /// return object. Returns true if the task's return object was returned
-  /// directly by value.
-  bool HandleTaskReturn(const ObjectID &object_id,
-                        const rpc::ReturnObject &return_object,
-                        const NodeID &worker_node_id,
-                        bool store_in_plasma) ABSL_LOCKS_EXCLUDED(mu_);
+  /// return object. On success, sets direct_return_out to true if the object's value
+  /// was returned directly by value (not stored in plasma).
+  StatusOr<bool> HandleTaskReturn(const ObjectID &object_id,
+                                  const rpc::ReturnObject &return_object,
+                                  const NodeID &worker_node_id,
+                                  bool store_in_plasma) ABSL_LOCKS_EXCLUDED(mu_);
 
   /// Remove a lineage reference to this object ID. This should be called
   /// whenever a task that depended on this object ID can no longer be retried.
diff --git a/src/ray/core_worker/tests/task_manager_test.cc b/src/ray/core_worker/tests/task_manager_test.cc
