Fix ObjectFetchTimedOutError

Signed-off-by: Jiajun Yao <jeromeyjj@gmail.com>

src/ray/core_worker/reference_count.cc

@@ -1527,6 +1527,11 @@ void ReferenceCounter::UpdateObjectReady(const ObjectID &object_id) {
   UpdateObjectPendingCreationInternal(object_id, /*pending_creation*/ false);
 }
 
+void ReferenceCounter::UpdateObjectPendingCreation(const ObjectID &object_id) {
+  absl::MutexLock lock(&mutex_);
+  UpdateObjectPendingCreationInternal(object_id, /*pending_creation*/ true);
+}
+
 bool ReferenceCounter::IsObjectPendingCreation(const ObjectID &object_id) const {
   absl::MutexLock lock(&mutex_);
   auto it = object_id_refs_.find(object_id);

src/ray/core_worker/reference_count.h

@@ -562,6 +562,9 @@ class ReferenceCounter : public ReferenceCounterInterface,
   /// Update that the object is ready to be fetched.
   void UpdateObjectReady(const ObjectID &object_id);
 
+  /// Update that the object's creator task is scheduled/executing.
+  void UpdateObjectPendingCreation(const ObjectID &object_id);
+
   /// Whether the object is pending creation (the task that creates it is
   /// scheduled/executing).
   bool IsObjectPendingCreation(const ObjectID &object_id) const;

src/ray/core_worker/task_manager.cc

@@ -30,6 +30,15 @@ const int64_t kTaskFailureThrottlingThreshold = 50;
 // Throttle task failure logs to once this interval.
 const int64_t kTaskFailureLoggingFrequencyMillis = 5000;
 
+absl::flat_hash_set<ObjectID> ObjectRefStream::GetItemsConsumed() const {
+  absl::flat_hash_set<ObjectID> result;
+  for (int64_t index = 0; index < next_index_; index++) {
+    const auto &object_id = GetObjectRefAtIndex(index);
+    result.emplace(object_id);
+  }
+  return result;
+}
+
 absl::flat_hash_set<ObjectID> ObjectRefStream::GetItemsUnconsumed() const {
   absl::flat_hash_set<ObjectID> result;
   for (int64_t index = 0; index <= max_index_seen_; index++) {
@@ -1016,6 +1025,16 @@ bool TaskManager::RetryTaskIfPossible(const TaskID &task_id,
                                  spec.AttemptNumber(),
                                  RayConfig::instance().task_oom_retry_delay_base_ms())
                            : RayConfig::instance().task_retry_delay_ms();
+    if (spec.IsStreamingGenerator()) {
+      absl::MutexLock lock(&object_ref_stream_ops_mu_);
+      const auto generator_id = spec.ReturnId(0);
+      auto stream_it = object_ref_streams_.find(generator_id);
+      if (stream_it != object_ref_streams_.end()) {
+        for (ObjectID obj : stream_it->second.GetItemsConsumed()) {
+          reference_counter_->UpdateObjectPendingCreation(obj);
+        }
+      }
+    }
     retry_task_callback_(spec, /*object_recovery*/ false, update_seqno, delay_ms);
     return true;
   } else {

src/ray/core_worker/task_manager.h

@@ -161,6 +161,8 @@ class ObjectRefStream {
   /// \return A list of object IDs that are not read yet.
   absl::flat_hash_set<ObjectID> GetItemsUnconsumed() const;
 
+  absl::flat_hash_set<ObjectID> GetItemsConsumed() const;
+
   /// Pop all ObjectIDs that are not read yet via
   /// TryReadNextItem.
   ///