[direct task] Retry tasks on failure and turn on RAY_FORCE_DIRECT for test_multinode_failures.py (#6306)

* multinode failures direct

* Add number of retries allowed for tasks

* Retry tasks

* Add failing test for object reconstruction

* Handle return status and debug

* update

* Retry task unit test

* update

* update

* todo

* Fix max_retries decorator, fix test

* Fix test that flaked

* lint

* comments

src/ray/core_worker/core_worker.cc

@@ -165,7 +165,10 @@ CoreWorker::CoreWorker(const WorkerType worker_type, const Language language,
       },
       ref_counting_enabled ? reference_counter_ : nullptr, raylet_client_));
 
-  task_manager_.reset(new TaskManager(memory_store_));
+  task_manager_.reset(
+      new TaskManager(memory_store_, [this](const TaskSpecification &spec) {
+        RAY_CHECK_OK(direct_task_submitter_->SubmitTask(spec));
+      }));
   resolver_.reset(new LocalDependencyResolver(memory_store_));
 
   // Create an entry for the driver task in the task table. This task is
@@ -589,7 +592,7 @@ void CoreWorker::PinObjectReferences(const TaskSpecification &task_spec,
 Status CoreWorker::SubmitTask(const RayFunction &function,
                               const std::vector<TaskArg> &args,
                               const TaskOptions &task_options,
-                              std::vector<ObjectID> *return_ids) {
+                              std::vector<ObjectID> *return_ids, int max_retries) {
   TaskSpecBuilder builder;
   const int next_task_index = worker_context_.GetNextTaskIndex();
   const auto task_id =
@@ -605,7 +608,7 @@ Status CoreWorker::SubmitTask(const RayFunction &function,
       return_ids);
   TaskSpecification task_spec = builder.Build();
   if (task_options.is_direct_call) {
-    task_manager_->AddPendingTask(task_spec);
+    task_manager_->AddPendingTask(task_spec, max_retries);
     PinObjectReferences(task_spec, TaskTransportType::DIRECT);
     return direct_task_submitter_->SubmitTask(task_spec);
   } else {

src/ray/core_worker/core_worker.h

@@ -280,7 +280,8 @@ class CoreWorker {
   /// \param[out] return_ids Ids of the return objects.
   /// \return Status error if task submission fails, likely due to raylet failure.
   Status SubmitTask(const RayFunction &function, const std::vector<TaskArg> &args,
-                    const TaskOptions &task_options, std::vector<ObjectID> *return_ids);
+                    const TaskOptions &task_options, std::vector<ObjectID> *return_ids,
+                    int max_retries);
 
   /// Create an actor.
   ///

src/ray/core_worker/task_manager.cc

@@ -2,10 +2,11 @@
 
 namespace ray {
 
-void TaskManager::AddPendingTask(const TaskSpecification &spec) {
+void TaskManager::AddPendingTask(const TaskSpecification &spec, int max_retries) {
   RAY_LOG(DEBUG) << "Adding pending task " << spec.TaskId();
   absl::MutexLock lock(&mu_);
-  RAY_CHECK(pending_tasks_.emplace(spec.TaskId(), spec.NumReturns()).second);
+  std::pair<TaskSpecification, int> entry = {spec, max_retries};
+  RAY_CHECK(pending_tasks_.emplace(spec.TaskId(), std::move(entry)).second);
 }
 
 void TaskManager::CompletePendingTask(const TaskID &task_id,
@@ -48,18 +49,48 @@ void TaskManager::CompletePendingTask(const TaskID &task_id,
   }
 }
 
-void TaskManager::FailPendingTask(const TaskID &task_id, rpc::ErrorType error_type) {
+void TaskManager::PendingTaskFailed(const TaskID &task_id, rpc::ErrorType error_type) {
+  if (error_type == rpc::ErrorType::ACTOR_DIED) {
+    // Note that this might be the __ray_terminate__ task, so we don't log
+    // loudly with ERROR here.
+    RAY_LOG(INFO) << "Task " << task_id << " failed with error "
+                  << rpc::ErrorType_Name(error_type);
+  } else {
+    RAY_LOG(ERROR) << "Task " << task_id << " failed with error "
+                   << rpc::ErrorType_Name(error_type);
+  }
+
   RAY_LOG(DEBUG) << "Failing task " << task_id;
-  int64_t num_returns;
+  int num_retries_left = 0;
+  TaskSpecification spec;
   {
     absl::MutexLock lock(&mu_);
     auto it = pending_tasks_.find(task_id);
     RAY_CHECK(it != pending_tasks_.end())
         << "Tried to complete task that was not pending " << task_id;
-    num_returns = it->second;
-    pending_tasks_.erase(it);
+    spec = it->second.first;
+    num_retries_left = it->second.second;
+    if (num_retries_left == 0) {
+      pending_tasks_.erase(it);
+    } else {
+      RAY_CHECK(num_retries_left > 0);
+      it->second.second--;
+    }
   }
 
+  // We should not hold the lock during these calls because they may trigger
+  // callbacks in this or other classes.
+  if (num_retries_left > 0) {
+    RAY_LOG(ERROR) << num_retries_left << " retries left for task " << spec.TaskId()
+                   << ", attempting to resubmit.";
+    retry_task_callback_(spec);
+  } else {
+    MarkPendingTaskFailed(task_id, spec.NumReturns(), error_type);
+  }
+}
+
+void TaskManager::MarkPendingTaskFailed(const TaskID &task_id, int64_t num_returns,
+                                        rpc::ErrorType error_type) {
   RAY_LOG(DEBUG) << "Treat task as failed. task_id: " << task_id
                  << ", error_type: " << ErrorType_Name(error_type);
   for (int i = 0; i < num_returns; i++) {

src/ray/core_worker/task_manager.h

@@ -18,21 +18,26 @@ class TaskFinisherInterface {
   virtual void CompletePendingTask(const TaskID &task_id,
                                    const rpc::PushTaskReply &reply) = 0;
 
-  virtual void FailPendingTask(const TaskID &task_id, rpc::ErrorType error_type) = 0;
+  virtual void PendingTaskFailed(const TaskID &task_id, rpc::ErrorType error_type) = 0;
 
   virtual ~TaskFinisherInterface() {}
 };
 
+using RetryTaskCallback = std::function<void(const TaskSpecification &spec)>;
+
 class TaskManager : public TaskFinisherInterface {
  public:
-  TaskManager(std::shared_ptr<CoreWorkerMemoryStore> in_memory_store)
-      : in_memory_store_(in_memory_store) {}
+  TaskManager(std::shared_ptr<CoreWorkerMemoryStore> in_memory_store,
+              RetryTaskCallback retry_task_callback)
+      : in_memory_store_(in_memory_store), retry_task_callback_(retry_task_callback) {}
 
   /// Add a task that is pending execution.
   ///
   /// \param[in] spec The spec of the pending task.
+  /// \param[in] max_retries Number of times this task may be retried
+  /// on failure.
   /// \return Void.
-  void AddPendingTask(const TaskSpecification &spec);
+  void AddPendingTask(const TaskSpecification &spec, int max_retries = 0);
 
   /// Return whether the task is pending.
   ///
@@ -50,23 +55,38 @@ class TaskManager : public TaskFinisherInterface {
   void CompletePendingTask(const TaskID &task_id,
                            const rpc::PushTaskReply &reply) override;
 
-  /// Treat a pending task as failed.
+  /// A pending task failed. This will either retry the task or mark the task
+  /// as failed if there are no retries left.
   ///
   /// \param[in] task_id ID of the pending task.
   /// \param[in] error_type The type of the specific error.
-  /// \return Void.
-  void FailPendingTask(const TaskID &task_id, rpc::ErrorType error_type) override;
+  void PendingTaskFailed(const TaskID &task_id, rpc::ErrorType error_type) override;
 
  private:
+  /// Treat a pending task as failed. The lock should not be held when calling
+  /// this method because it may trigger callbacks in this or other classes.
+  void MarkPendingTaskFailed(const TaskID &task_id, int64_t num_returns,
+                             rpc::ErrorType error_type) LOCKS_EXCLUDED(mu_);
+
   /// Used to store task results.
   std::shared_ptr<CoreWorkerMemoryStore> in_memory_store_;
 
+  /// Called when a task should be retried.
+  const RetryTaskCallback retry_task_callback_;
+
   /// Protects below fields.
   absl::Mutex mu_;
 
-  /// Map from task ID to the task's number of return values. This map contains
-  /// one entry per pending task that we submitted.
-  absl::flat_hash_map<TaskID, int64_t> pending_tasks_ GUARDED_BY(mu_);
+  /// Map from task ID to a pair of:
+  ///   {task spec, number of allowed retries left}
+  /// This map contains one entry per pending task that we submitted.
+  /// TODO(swang): The TaskSpec protobuf must be copied into the
+  /// PushTaskRequest protobuf when sent to a worker so that we can retry it if
+  /// the worker fails. We could avoid this by either not caching the full
+  /// TaskSpec for tasks that cannot be retried (e.g., actor tasks), or by
+  /// storing a shared_ptr to a PushTaskRequest protobuf for all tasks.
+  absl::flat_hash_map<TaskID, std::pair<TaskSpecification, int>> pending_tasks_
+      GUARDED_BY(mu_);
 };
 
 }  // namespace ray

src/ray/core_worker/test/core_worker_test.cc

@@ -259,7 +259,8 @@ void CoreWorkerTest::TestNormalTask(std::unordered_map<std::string, double> &res
       options.is_direct_call = true;
 
       std::vector<ObjectID> return_ids;
-      RAY_CHECK_OK(driver.SubmitTask(func, args, options, &return_ids));
+      RAY_CHECK_OK(
+          driver.SubmitTask(func, args, options, &return_ids, /*max_retries=*/0));
 
       ASSERT_EQ(return_ids.size(), 1);
 

src/ray/core_worker/test/direct_actor_transport_test.cc

@@ -42,7 +42,7 @@ class MockTaskFinisher : public TaskFinisherInterface {
   MockTaskFinisher() {}
 
   MOCK_METHOD2(CompletePendingTask, void(const TaskID &, const rpc::PushTaskReply &));
-  MOCK_METHOD2(FailPendingTask, void(const TaskID &task_id, rpc::ErrorType error_type));
+  MOCK_METHOD2(PendingTaskFailed, void(const TaskID &task_id, rpc::ErrorType error_type));
 };
 
 TaskSpecification CreateActorTaskHelper(ActorID actor_id, int64_t counter) {
@@ -86,7 +86,7 @@ TEST_F(DirectActorTransportTest, TestSubmitTask) {
 
   EXPECT_CALL(*task_finisher_, CompletePendingTask(TaskID::Nil(), _))
       .Times(worker_client_->callbacks.size());
-  EXPECT_CALL(*task_finisher_, FailPendingTask(_, _)).Times(0);
+  EXPECT_CALL(*task_finisher_, PendingTaskFailed(_, _)).Times(0);
   while (!worker_client_->callbacks.empty()) {
     ASSERT_TRUE(worker_client_->ReplyPushTask());
   }
@@ -163,7 +163,7 @@ TEST_F(DirectActorTransportTest, TestActorFailure) {
   ASSERT_EQ(worker_client_->callbacks.size(), 2);
 
   // Simulate the actor dying. All submitted tasks should get failed.
-  EXPECT_CALL(*task_finisher_, FailPendingTask(_, _)).Times(2);
+  EXPECT_CALL(*task_finisher_, PendingTaskFailed(_, _)).Times(2);
   EXPECT_CALL(*task_finisher_, CompletePendingTask(_, _)).Times(0);
   while (!worker_client_->callbacks.empty()) {
     ASSERT_TRUE(worker_client_->ReplyPushTask(Status::IOError("")));

src/ray/core_worker/test/direct_task_transport_test.cc

@@ -44,7 +44,7 @@ class MockTaskFinisher : public TaskFinisherInterface {
   void CompletePendingTask(const TaskID &, const rpc::PushTaskReply &) override {
     num_tasks_complete++;
   }
-  void FailPendingTask(const TaskID &task_id, rpc::ErrorType error_type) override {
+  void PendingTaskFailed(const TaskID &task_id, rpc::ErrorType error_type) override {
     num_tasks_failed++;
   }
 

src/ray/core_worker/test/task_manager_test.cc

@@ -18,10 +18,14 @@ class TaskManagerTest : public ::testing::Test {
  public:
   TaskManagerTest()
       : store_(std::shared_ptr<CoreWorkerMemoryStore>(new CoreWorkerMemoryStore())),
-        manager_(store_) {}
+        manager_(store_, [this](const TaskSpecification &spec) {
+          num_retries_++;
+          return Status::OK();
+        }) {}
 
   std::shared_ptr<CoreWorkerMemoryStore> store_;
   TaskManager manager_;
+  int num_retries_ = 0;
 };
 
 TEST_F(TaskManagerTest, TestTaskSuccess) {
@@ -47,6 +51,7 @@ TEST_F(TaskManagerTest, TestTaskSuccess) {
   ASSERT_EQ(std::memcmp(results[0]->GetData()->Data(), return_object->data().data(),
                         return_object->data().size()),
             0);
+  ASSERT_EQ(num_retries_, 0);
 }
 
 TEST_F(TaskManagerTest, TestTaskFailure) {
@@ -58,7 +63,7 @@ TEST_F(TaskManagerTest, TestTaskFailure) {
   WorkerContext ctx(WorkerType::WORKER, JobID::FromInt(0));
 
   auto error = rpc::ErrorType::WORKER_DIED;
-  manager_.FailPendingTask(spec.TaskId(), error);
+  manager_.PendingTaskFailed(spec.TaskId(), error);
   ASSERT_FALSE(manager_.IsTaskPending(spec.TaskId()));
 
   std::vector<std::shared_ptr<RayObject>> results;
@@ -67,6 +72,36 @@ TEST_F(TaskManagerTest, TestTaskFailure) {
   rpc::ErrorType stored_error;
   ASSERT_TRUE(results[0]->IsException(&stored_error));
   ASSERT_EQ(stored_error, error);
+  ASSERT_EQ(num_retries_, 0);
+}
+
+TEST_F(TaskManagerTest, TestTaskRetry) {
+  auto spec = CreateTaskHelper(1);
+  ASSERT_FALSE(manager_.IsTaskPending(spec.TaskId()));
+  int num_retries = 3;
+  manager_.AddPendingTask(spec, num_retries);
+  ASSERT_TRUE(manager_.IsTaskPending(spec.TaskId()));
+  auto return_id = spec.ReturnId(0, TaskTransportType::DIRECT);
+  WorkerContext ctx(WorkerType::WORKER, JobID::FromInt(0));
+
+  auto error = rpc::ErrorType::WORKER_DIED;
+  for (int i = 0; i < num_retries; i++) {
+    manager_.PendingTaskFailed(spec.TaskId(), error);
+    ASSERT_TRUE(manager_.IsTaskPending(spec.TaskId()));
+    std::vector<std::shared_ptr<RayObject>> results;
+    ASSERT_FALSE(store_->Get({return_id}, 1, 0, ctx, false, &results).ok());
+    ASSERT_EQ(num_retries_, i + 1);
+  }
+
+  manager_.PendingTaskFailed(spec.TaskId(), error);
+  ASSERT_FALSE(manager_.IsTaskPending(spec.TaskId()));
+
+  std::vector<std::shared_ptr<RayObject>> results;
+  RAY_CHECK_OK(store_->Get({return_id}, 1, -0, ctx, false, &results));
+  ASSERT_EQ(results.size(), 1);
+  rpc::ErrorType stored_error;
+  ASSERT_TRUE(results[0]->IsException(&stored_error));
+  ASSERT_EQ(stored_error, error);
 }
 
 }  // namespace ray

src/ray/core_worker/transport/direct_actor_transport.cc

@@ -20,7 +20,10 @@ Status CoreWorkerDirectActorTaskSubmitter::SubmitTask(TaskSpecification task_spe
     const auto task_id = task_spec.TaskId();
 
     auto request = std::unique_ptr<rpc::PushTaskRequest>(new rpc::PushTaskRequest);
-    request->mutable_task_spec()->Swap(&task_spec.GetMutableMessage());
+    // NOTE(swang): CopyFrom is needed because if we use Swap here and the task
+    // fails, then the task data will be gone when the TaskManager attempts to
+    // access the task.
+    request->mutable_task_spec()->CopyFrom(task_spec.GetMessage());
 
     std::unique_lock<std::mutex> guard(mutex_);
 
@@ -45,7 +48,7 @@ Status CoreWorkerDirectActorTaskSubmitter::SubmitTask(TaskSpecification task_spe
     } else {
       // Actor is dead, treat the task as failure.
       RAY_CHECK(iter->second.state_ == ActorTableData::DEAD);
-      task_finisher_->FailPendingTask(task_id, rpc::ErrorType::ACTOR_DIED);
+      task_finisher_->PendingTaskFailed(task_id, rpc::ErrorType::ACTOR_DIED);
     }
   });
 
@@ -85,7 +88,7 @@ void CoreWorkerDirectActorTaskSubmitter::HandleActorUpdate(
         auto request = std::move(head->second);
         head = pending_it->second.erase(head);
         auto task_id = TaskID::FromBinary(request->task_spec().task_id());
-        task_finisher_->FailPendingTask(task_id, rpc::ErrorType::ACTOR_DIED);
+        task_finisher_->PendingTaskFailed(task_id, rpc::ErrorType::ACTOR_DIED);
       }
       pending_requests_.erase(pending_it);
     }
@@ -123,21 +126,15 @@ void CoreWorkerDirectActorTaskSubmitter::PushActorTask(
       << "Counter was " << task_number << " expected " << next_sequence_number_[actor_id];
   next_sequence_number_[actor_id]++;
 
-  auto status = client.PushActorTask(
+  RAY_CHECK_OK(client.PushActorTask(
       std::move(request),
       [this, task_id](Status status, const rpc::PushTaskReply &reply) {
         if (!status.ok()) {
-          // Note that this might be the __ray_terminate__ task, so we don't log
-          // loudly with ERROR here.
-          RAY_LOG(INFO) << "Task failed with error: " << status;
-          task_finisher_->FailPendingTask(task_id, rpc::ErrorType::ACTOR_DIED);
+          task_finisher_->PendingTaskFailed(task_id, rpc::ErrorType::ACTOR_DIED);
         } else {
           task_finisher_->CompletePendingTask(task_id, reply);
         }
-      });
-  if (!status.ok()) {
-    task_finisher_->FailPendingTask(task_id, rpc::ErrorType::ACTOR_DIED);
-  }
+      }));
 }
 
 bool CoreWorkerDirectActorTaskSubmitter::IsActorAlive(const ActorID &actor_id) const {

src/ray/core_worker/transport/direct_task_transport.cc

@@ -5,7 +5,9 @@
 namespace ray {
 
 Status CoreWorkerDirectTaskSubmitter::SubmitTask(TaskSpecification task_spec) {
+  RAY_LOG(DEBUG) << "Submit task " << task_spec.TaskId();
   resolver_.ResolveDependencies(task_spec, [this, task_spec]() {
+    RAY_LOG(DEBUG) << "Task dependencies resolved " << task_spec.TaskId();
     absl::MutexLock lock(&mu_);
     // Note that the dependencies in the task spec are mutated to only contain
     // plasma dependencies after ResolveDependencies finishes.
@@ -138,11 +140,15 @@ void CoreWorkerDirectTaskSubmitter::RequestNewWorkerIfNeeded(
 void CoreWorkerDirectTaskSubmitter::PushNormalTask(const rpc::WorkerAddress &addr,
                                                    rpc::CoreWorkerClientInterface &client,
                                                    const SchedulingKey &scheduling_key,
-                                                   TaskSpecification &task_spec) {
+                                                   const TaskSpecification &task_spec) {
   auto task_id = task_spec.TaskId();
   auto request = std::unique_ptr<rpc::PushTaskRequest>(new rpc::PushTaskRequest);
-  request->mutable_task_spec()->Swap(&task_spec.GetMutableMessage());
-  auto status = client.PushNormalTask(
+  RAY_LOG(DEBUG) << "Pushing normal task " << task_spec.TaskId();
+  // NOTE(swang): CopyFrom is needed because if we use Swap here and the task
+  // fails, then the task data will be gone when the TaskManager attempts to
+  // access the task.
+  request->mutable_task_spec()->CopyFrom(task_spec.GetMessage());
+  RAY_CHECK_OK(client.PushNormalTask(
       std::move(request), [this, task_id, scheduling_key, addr](
                               Status status, const rpc::PushTaskReply &reply) {
         {
@@ -150,14 +156,14 @@ void CoreWorkerDirectTaskSubmitter::PushNormalTask(const rpc::WorkerAddress &add
           OnWorkerIdle(addr, scheduling_key, /*error=*/!status.ok());
         }
         if (!status.ok()) {
-          task_finisher_->FailPendingTask(task_id, rpc::ErrorType::WORKER_DIED);
+          // TODO: It'd be nice to differentiate here between process vs node
+          // failure (e.g., by contacting the raylet). If it was a process
+          // failure, it may have been an application-level error and it may
+          // not make sense to retry the task.
+          task_finisher_->PendingTaskFailed(task_id, rpc::ErrorType::WORKER_DIED);
         } else {
           task_finisher_->CompletePendingTask(task_id, reply);
         }
-      });
-  if (!status.ok()) {
-    // TODO(swang): add unit test for this.
-    task_finisher_->FailPendingTask(task_id, rpc::ErrorType::WORKER_DIED);
-  }
+      }));
 }
 };  // namespace ray

src/ray/core_worker/transport/direct_task_transport.h

@@ -76,7 +76,8 @@ class CoreWorkerDirectTaskSubmitter {
   /// Push a task to a specific worker.
   void PushNormalTask(const rpc::WorkerAddress &addr,
                       rpc::CoreWorkerClientInterface &client,
-                      const SchedulingKey &task_queue_key, TaskSpecification &task_spec);
+                      const SchedulingKey &task_queue_key,
+                      const TaskSpecification &task_spec);
 
   // Client that can be used to lease and return workers from the local raylet.
   std::shared_ptr<WorkerLeaseInterface> local_lease_client_;

src/ray/protobuf/common.proto

@@ -77,6 +77,8 @@ message TaskSpec {
   ActorTaskSpec actor_task_spec = 15;
   // Whether this task is a direct call task.
   bool is_direct_call = 16;
+  // Number of times this task may be retried on worker failure.
+  int32 max_retries = 17;
 }
 
 // Argument in the task.