propagate status, update tests and minor refactor

codope · codope · commit 136c43c5d6ac · 2025-08-18T15:59:09.000Z
Signed-off-by: Sagar Sumit &lt;sagarsumit09@gmail.com&gt;
diff --git a/src/ray/core_worker/core_worker_process.cc b/src/ray/core_worker/core_worker_process.cc
@@ -424,21 +424,12 @@ std::shared_ptr<CoreWorker> CoreWorkerProcessImpl::CreateCoreWorker(
       /*put_in_local_plasma_callback=*/
       [this](const RayObject &object, const ObjectID &object_id) {
         auto core_worker = GetCoreWorker();
-
         auto status =
             core_worker->PutInLocalPlasmaStore(object, object_id, /*pin_object=*/true);
         if (!status.ok()) {
-          // During shutdown, plasma store connections can be closed causing IOError.
-          // Tolerate only IOError during shutdown to avoid masking real errors.
-          if (status.IsIOError() && core_worker->IsShuttingDown()) {
-            // Double-check shutdown state - this handles the race where shutdown
-            // began after our first check but before the plasma operation.
-            RAY_LOG(WARNING) << "Failed to put error object " << object_id
-                             << " in plasma store during shutdown: " << status.ToString();
-            return Status::OK();
-          } else {
-            return status;
-          }
+          RAY_LOG(WARNING).WithField(object_id)
+              << "Failed to put object in plasma store: " << status;
+          return status;
         }
         return Status::OK();
       },
diff --git a/src/ray/core_worker/task_manager.cc b/src/ray/core_worker/task_manager.cc
@@ -533,11 +533,10 @@ size_t TaskManager::NumPendingTasks() const {
   return num_pending_tasks_;
 }
 
-Status TaskManager::HandleTaskReturn(const ObjectID &object_id,
-                                     const rpc::ReturnObject &return_object,
-                                     const NodeID &worker_node_id,
-                                     bool store_in_plasma,
-                                     bool *direct_return_out) {
+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 "
@@ -581,6 +580,11 @@ Status TaskManager::HandleTaskReturn(const ObjectID &object_id,
                      tensor_transport.value_or(rpc::TensorTransport::OBJECT_STORE));
     if (store_in_plasma) {
       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;
       }
@@ -599,10 +603,7 @@ Status TaskManager::HandleTaskReturn(const ObjectID &object_id,
     }
     reference_counter_.AddNestedObjectIds(object_id, nested_ids, owner_address);
   }
-  if (direct_return_out != nullptr) {
-    *direct_return_out = direct_return;
-  }
-  return Status::OK();
+  return direct_return;
 }
 
 bool TaskManager::TryDelObjectRefStream(const ObjectID &generator_id) {
@@ -820,13 +821,15 @@ bool TaskManager::HandleReportGeneratorItemReturns(
     }
     // When an object is reported, the object is ready to be fetched.
     reference_counter_.UpdateObjectPendingCreation(object_id, false);
-    bool _direct_unused = false;
-    RAY_UNUSED(
+    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),
-                         &_direct_unused));
+                         /*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.
@@ -910,41 +913,49 @@ void TaskManager::CompletePendingTask(const TaskID &task_id,
         reference_counter_.AddDynamicReturn(object_id, generator_id);
         dynamic_return_ids.push_back(object_id);
       }
-      bool direct = false;
-      Status s = HandleTaskReturn(object_id,
-                                  return_object,
-                                  NodeID::FromBinary(worker_addr.node_id()),
-                                  store_in_plasma_ids.contains(object_id),
-                                  &direct);
-      if (!s.ok()) {
+      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: " << s;
-      } else if (!direct && first_execution) {
+            << "Failed to handle dynamic task return: " << direct_or.status();
+        // Treat as system failure for this attempt and fail immediately to avoid hangs.
+        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() && 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());
-    bool direct = false;
-    Status s = HandleTaskReturn(object_id,
-                                return_object,
-                                NodeID::FromBinary(worker_addr.node_id()),
-                                store_in_plasma_ids.contains(object_id),
-                                &direct);
-    if (!s.ok()) {
-      RAY_LOG(WARNING).WithField(object_id) << "Failed to handle task return: " << s;
+    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,
-                             &s,
+                             &st,
                              /*ray_error_info=*/nullptr,
                              /*mark_task_object_failed=*/true,
                              /*fail_immediately=*/true);
       return;
-    } else if (direct) {
+    } else if (direct_or.value()) {
       direct_return_ids.push_back(object_id);
     }
   }
@@ -1068,12 +1079,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);
-          bool _direct_unused = false;
-          RAY_UNUSED(HandleTaskReturn(generator_return_id,
-                                      return_object,
-                                      NodeID::FromBinary(worker_addr.node_id()),
-                                      store_in_plasma_ids.contains(generator_return_id),
-                                      &_direct_unused));
+          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();
+          }
         }
       }
     }
@@ -1484,26 +1499,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)) {
       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;
       }
-    } else {
-      in_memory_store_.Put(error, object_id);
     }
   }
   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)) {
         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;
         }
-      } else {
-        in_memory_store_.Put(error, dynamic_return_id);
       }
     }
   }
diff --git a/src/ray/core_worker/task_manager.h b/src/ray/core_worker/task_manager.h
@@ -612,11 +612,10 @@ class TaskManager : public TaskManagerInterface {
   /// Update nested ref count info and store the in-memory value for a task's
   /// return object. On success, sets direct_return_out to true if the object's value
   /// was returned directly by value (not stored in plasma).
-  Status HandleTaskReturn(const ObjectID &object_id,
-                          const rpc::ReturnObject &return_object,
-                          const NodeID &worker_node_id,
-                          bool store_in_plasma,
-                          bool *direct_return_out) ABSL_LOCKS_EXCLUDED(mu_);
+  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
