[core][task-manager/02] consolidate TaskManager interface (#54317)

This PR is smaller than it looks.

The `TaskManager` class currently exposes two interfaces: `TaskFinisher`
and `TaskResubmission`. While these interfaces are well-intentioned,
they are only implemented by `TaskManager` itself, and the methods they
define are not fully independent. As a result, it’s unlikely that these
interfaces could be meaningfully separated or implemented in isolation.

This change consolidates them into a single `TaskManager` interface,
which can be reused where needed. The goal is to reduce the number of
concepts and components required to reason about the Ray core, and to
simplify the overall design.

Test:
- CI

Signed-off-by: Cuong Nguyen <[email protected]>

Author: can-anyscale

src/mock/ray/core_worker/task_manager.h renamed to src/mock/ray/core_worker/task_manager_interface.h

@@ -14,11 +14,20 @@
 
 #pragma once
 #include "gmock/gmock.h"
+#include "ray/core_worker/task_manager_interface.h"
+
 namespace ray {
 namespace core {
 
-class MockTaskFinisherInterface : public TaskFinisherInterface {
+class MockTaskManagerInterface : public TaskManagerInterface {
  public:
+  MOCK_METHOD(std::vector<rpc::ObjectReference>,
+              AddPendingTask,
+              (const rpc::Address &caller_address,
+               const TaskSpecification &spec,
+               const std::string &call_site,
+               int max_retries),
+              (override));
   MOCK_METHOD(void,
               CompletePendingTask,
               (const TaskID &task_id,
@@ -42,6 +51,10 @@ class MockTaskFinisherInterface : public TaskFinisherInterface {
                bool mark_task_object_failed,
                bool fail_immediately),
               (override));
+  MOCK_METHOD(std::optional<rpc::ErrorType>,
+              ResubmitTask,
+              (const TaskID &task_id, std::vector<ObjectID> *task_deps),
+              (override));
   MOCK_METHOD(void,
               OnTaskDependenciesInlined,
               (const std::vector<ObjectID> &inlined_dependency_ids,

src/ray/core_worker/BUILD.bazel

@@ -262,8 +262,8 @@ ray_cc_library(
 )
 
 ray_cc_library(
-    name = "task_finisher",
-    hdrs = ["task_finisher.h"],
+    name = "task_manager_interface",
+    hdrs = ["task_manager_interface.h"],
     deps = [
         "//src/ray/common:id",
         "//src/ray/common:status",
@@ -282,7 +282,7 @@ ray_cc_library(
         ":actor_manager",
         ":memory_store",
         ":task_event_buffer",
-        ":task_finisher",
+        ":task_manager_interface",
         "//src/ray/gcs:gcs_pb_util",
         "//src/ray/common:id",
         "//src/ray/common:ray_object",
@@ -306,7 +306,7 @@ ray_cc_library(
     deps = [
         ":actor_creator",
         ":memory_store",
-        ":task_finisher",
+        ":task_manager_interface",
         "//src/ray/common:id",
         "//src/ray/common:task_common",
         "@com_google_absl//absl/container:flat_hash_map",

src/ray/core_worker/object_recovery_manager.cc

@@ -183,7 +183,7 @@ void ObjectRecoveryManager::ReconstructObject(const ObjectID &object_id) {
   // after ResubmitTask, then it will remain true forever.
   // see https://github.com/ray-project/ray/issues/47606 for more details.
   reference_counter_.UpdateObjectPendingCreation(object_id, true);
-  auto error_type_optional = task_resubmitter_.ResubmitTask(task_id, &task_deps);
+  auto error_type_optional = task_manager_.ResubmitTask(task_id, &task_deps);
 
   if (!error_type_optional.has_value()) {
     // Try to recover the task's dependencies.

src/ray/core_worker/object_recovery_manager.h

@@ -48,11 +48,11 @@ class ObjectRecoveryManager {
       std::shared_ptr<PinObjectsInterface> local_object_pinning_client,
       std::function<Status(const ObjectID &object_id,
                            const ObjectLookupCallback &callback)> object_lookup,
-      TaskResubmissionInterface &task_resubmitter,
+      TaskManagerInterface &task_manager,
       ReferenceCounter &reference_counter,
       CoreWorkerMemoryStore &in_memory_store,
       ObjectRecoveryFailureCallback recovery_failure_callback)
-      : task_resubmitter_(task_resubmitter),
+      : task_manager_(task_manager),
         reference_counter_(reference_counter),
         rpc_address_(std::move(rpc_address)),
         client_factory_(std::move(client_factory)),
@@ -112,7 +112,7 @@ class ObjectRecoveryManager {
   void ReconstructObject(const ObjectID &object_id);
 
   /// Used to resubmit tasks.
-  TaskResubmissionInterface &task_resubmitter_;
+  TaskManagerInterface &task_manager_;
 
   /// Used to check whether we own an object.
   ReferenceCounter &reference_counter_;

src/ray/core_worker/task_finisher.h

@@ -26,7 +26,7 @@
 #include "ray/common/id.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_finisher.h"
+#include "ray/core_worker/task_manager_interface.h"
 #include "ray/stats/metric_defs.h"
 #include "ray/util/counter_map.h"
 #include "src/ray/protobuf/common.pb.h"
@@ -38,14 +38,6 @@ namespace core {
 
 class ActorManager;
 
-class TaskResubmissionInterface {
- public:
-  virtual std::optional<rpc::ErrorType> ResubmitTask(
-      const TaskID &task_id, std::vector<ObjectID> *task_deps) = 0;
-
-  virtual ~TaskResubmissionInterface() = default;
-};
-
 using TaskStatusCounter = CounterMap<std::tuple<std::string, rpc::TaskStatus, bool>>;
 using PutInLocalPlasmaCallback =
     std::function<void(const RayObject &object, const ObjectID &object_id)>;
@@ -173,7 +165,7 @@ class ObjectRefStream {
   int64_t total_num_object_consumed_{};
 };
 
-class TaskManager : public TaskFinisherInterface, public TaskResubmissionInterface {
+class TaskManager : public TaskManagerInterface {
  public:
   TaskManager(CoreWorkerMemoryStore &in_memory_store,
               ReferenceCounter &reference_counter,
@@ -208,36 +200,11 @@ class TaskManager : public TaskFinisherInterface, public TaskResubmissionInterfa
         });
   }
 
-  /// Add a task that is pending execution.
-  ///
-  /// The local ref count for all return refs (excluding actor creation tasks)
-  /// will be initialized to 1 so that the ref is considered in scope before
-  /// returning to the language frontend. The caller is responsible for
-  /// decrementing the ref count once the frontend ref has gone out of scope.
-  ///
-  /// \param[in] caller_address The rpc address of the calling task.
-  /// \param[in] spec The spec of the pending task.
-  /// \param[in] max_retries Number of times this task may be retried
-  /// on failure.
-  /// \return ObjectRefs returned by this task.
   std::vector<rpc::ObjectReference> AddPendingTask(const rpc::Address &caller_address,
                                                    const TaskSpecification &spec,
                                                    const std::string &call_site,
-                                                   int max_retries = 0);
-
-  /// Resubmit a task that has completed execution before. This is used to
-  /// reconstruct objects stored in Plasma that were lost.
-  ///
-  /// \param[in] task_id The ID of the task to resubmit.
-  /// \param[out] task_deps The object dependencies of the resubmitted task,
-  /// i.e. all arguments that were not inlined in the task spec. The caller is
-  /// responsible for making sure that these dependencies become available, so
-  /// that the resubmitted task can run. This is only populated if the task was
-  /// not already pending and was successfully resubmitted.
-  /// \return nullopt if the task was successfully resubmitted (task or actor being
-  /// scheduled, but no guarantee on completion), or was already pending. Return the
-  /// appopriate error type to propagate for the object if the task was not successfully
-  /// resubmitted.
+                                                   int max_retries = 0) override;
+
   std::optional<rpc::ErrorType> ResubmitTask(const TaskID &task_id,
                                              std::vector<ObjectID> *task_deps) override;
 
@@ -246,13 +213,6 @@ class TaskManager : public TaskFinisherInterface, public TaskResubmissionInterfa
   /// \param shutdown The shutdown callback to call.
   void DrainAndShutdown(std::function<void()> shutdown);
 
-  /// Write return objects for a pending task to the memory store.
-  ///
-  /// \param[in] task_id ID of the pending task.
-  /// \param[in] reply Proto response to a direct actor or task call.
-  /// \param[in] worker_addr Address of the worker that executed the task.
-  /// \param[in] is_application_error Whether this is an Exception return.
-  /// \return Void.
   void CompletePendingTask(const TaskID &task_id,
                            const rpc::PushTaskReply &reply,
                            const rpc::Address &worker_addr,
@@ -434,50 +394,18 @@ class TaskManager : public TaskFinisherInterface, public TaskResubmissionInterfa
   std::pair<ObjectID, bool> PeekObjectRefStream(const ObjectID &generator_id)
       ABSL_LOCKS_EXCLUDED(mu_);
 
-  /// Called by submitter when a generator task marked for resubmission for intermediate
-  /// object recovery comes back from the executing worker. We mark the attempt as failed
-  /// and resubmit it, so we can recover the intermediate return.
   void MarkGeneratorFailedAndResubmit(const TaskID &task_id) override;
 
-  /// Returns true if task can be retried.
-  ///
-  /// \param[in] task_id ID of the task to be retried.
-  /// \return true if task is scheduled to be retried.
   bool RetryTaskIfPossible(const TaskID &task_id,
                            const rpc::RayErrorInfo &error_info) override;
 
-  /// 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.
-  /// \param[in] status Optional status message.
-  /// \param[in] ray_error_info The error information of a given error type.
-  /// Nullptr means that there's no error information.
-  /// TODO(sang): Remove nullptr case. Every error message should have metadata.
-  /// \param[in] mark_task_object_failed whether or not it marks the task
-  /// return object as failed. If this is set to false, then the caller is
-  /// responsible for later failing or completing the task.
-  /// \param[in] fail_immediately whether to fail the task and ignore
-  /// the retries that are available.
-  /// \return Whether the task will be retried or not.
   bool FailOrRetryPendingTask(const TaskID &task_id,
                               rpc::ErrorType error_type,
                               const Status *status = nullptr,
                               const rpc::RayErrorInfo *ray_error_info = nullptr,
                               bool mark_task_object_failed = true,
                               bool fail_immediately = false) override;
 
-  /// A pending task failed. This will mark the task as failed.
-  /// This doesn't always mark the return object as failed
-  /// depending on mark_task_object_failed.
-  ///
-  /// \param[in] task_id ID of the pending task.
-  /// \param[in] error_type The type of the specific error.
-  /// \param[in] status Optional status message.
-  /// \param[in] ray_error_info The error information of a given error type.
-  /// \param[in] mark_task_object_failed whether or not it marks the task
-  /// return object as failed.
   void FailPendingTask(const TaskID &task_id,
                        rpc::ErrorType error_type,
                        const Status *status = nullptr,
@@ -495,24 +423,11 @@ class TaskManager : public TaskFinisherInterface, public TaskResubmissionInterfa
       const rpc::RayErrorInfo *ray_error_info,
       const absl::flat_hash_set<ObjectID> &store_in_plasma_ids) ABSL_LOCKS_EXCLUDED(mu_);
 
-  /// A task's dependencies were inlined in the task spec. This will decrement
-  /// the ref count for the dependency IDs. If the dependencies contained other
-  /// ObjectIDs, then the ref count for these object IDs will be incremented.
-  ///
-  /// \param[in] inlined_dependency_ids The args that were originally passed by
-  /// reference into the task, but have now been inlined.
-  /// \param[in] contained_ids Any ObjectIDs that were newly inlined in the
-  /// task spec, because a serialized copy of the ID was contained in one of
-  /// the inlined dependencies.
   void OnTaskDependenciesInlined(const std::vector<ObjectID> &inlined_dependency_ids,
                                  const std::vector<ObjectID> &contained_ids) override;
 
-  /// Set the task state to be canceled. Set the number of retries to zero.
-  ///
-  /// \param[in] task_id to cancel.
   void MarkTaskCanceled(const TaskID &task_id) override;
 
-  /// Return the spec for a pending task.
   std::optional<TaskSpecification> GetTaskSpec(const TaskID &task_id) const override;
 
   /// Return specs for pending children tasks of the given parent task.
@@ -524,10 +439,6 @@ class TaskManager : public TaskFinisherInterface, public TaskResubmissionInterfa
   /// \return Whether the task can be submitted for execution.
   bool IsTaskSubmissible(const TaskID &task_id) const;
 
-  /// Return whether the task is pending.
-  ///
-  /// \param[in] task_id ID of the task to query.
-  /// \return Whether the task is pending.
   bool IsTaskPending(const TaskID &task_id) const override;
 
   /// Return whether the task is scheduled adn waiting for execution.
@@ -549,17 +460,8 @@ class TaskManager : public TaskFinisherInterface, public TaskResubmissionInterfa
     return total_lineage_footprint_bytes_;
   }
 
-  /// Record that the given task's dependencies have been created and the task
-  /// can now be scheduled for execution.
-  ///
-  /// \param[in] task_id The task that is now scheduled.
   void MarkDependenciesResolved(const TaskID &task_id) override;
 
-  /// Record that the given task is scheduled and wait for execution.
-  ///
-  /// \param[in] task_id The task that is will be running.
-  /// \param[in] node_id The node id that this task wil be running.
-  /// \param[in] worker_id The worker id that this task wil be running.
   void MarkTaskWaitingForExecution(const TaskID &task_id,
                                    const NodeID &node_id,
                                    const WorkerID &worker_id) override;