[core][refactor] Move `to_resubmit_` from CoreWorker to TaskManager to avoid an abstraction leak #52779

kevin85421 · 2025-05-05T02:27:57Z

Why are these changes needed?

Currently, CoreWorker has a member to_resubmit_ which is an abstraction leak for me. This PR moves to_resubmit_ from CoreWorker to TaskManager.

Related issue number

Checks

I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
I've run scripts/format.sh to lint the changes in this PR.
I've included any doc changes needed for https://docs.ray.io/en/master/.
- I've added any new APIs to the API Reference. For example, if I added a
  method in Tune, I've added it in doc/source/tune/api/ under the
  corresponding .rst file.
I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
Testing Strategy
- Unit tests
- Release tests
- This PR is not tested :(

Signed-off-by: kaihsun <[email protected]>

kevin85421 · 2025-05-05T08:57:26Z

src/ray/core_worker/test/task_manager_test.cc

    ASSERT_TRUE(manager_.IsTaskPending(spec.TaskId()));
    ASSERT_TRUE(reference_counter_->IsObjectPendingCreation(return_id));
    ASSERT_EQ(reference_counter_->NumObjectIDsInScope(), 3);
    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);


FailOrRetryPendingTask will not call the retry_task_execution callback because object_recovery is false. Hence, we can't use num_retries_ for testing.

I don't follow this. Isn't the number of attempts incremented in RetryTask and not retry_task_execution?

…tmux3-ray2 Signed-off-by: kaihsun <[email protected]>

Signed-off-by: kaihsun <[email protected]>

kevin85421 · 2025-05-15T19:21:05Z

src/ray/stats/metric.cc

@@ -136,7 +136,8 @@ void Metric::Record(double value, TagsType tags) {
  opencensus::stats::Record({{*measure_, value}}, std::move(combined_tags));
 }

-void Metric::Record(double value, std::unordered_map<std::string_view, std::string> tags) {


This is updated by the C++ linter.

kevin85421 · 2025-05-15T19:23:24Z

src/ray/core_worker/task_manager.h

@@ -172,6 +173,29 @@ class ObjectRefStream {
  int64_t total_num_object_consumed_{};
 };

+/// Represents a task that needs to be retried at a specific time.


copy from core_worker.h

kevin85421 · 2025-05-15T19:24:12Z

src/ray/core_worker/task_manager.h

@@ -863,6 +901,10 @@ class TaskManager : public TaskFinisherInterface, public TaskResubmissionInterfa
  /// error).
  worker::TaskEventBuffer &task_event_buffer_;

+  /// Queue of tasks to retry, ordered by their execution time.


copy from core_worker.h

kevin85421 · 2025-05-15T19:27:52Z

src/ray/core_worker/task_manager.cc

+void TaskManager::RetryTask(TaskSpecification &spec,
+                            bool object_recovery,
+                            uint32_t delay_ms) {
+  spec.GetMutableMessage().set_attempt_number(spec.AttemptNumber() + 1);


copy from core_worker.cc (retry_task_callback)

kevin85421 · 2025-05-15T19:29:05Z

src/ray/core_worker/task_manager.h

@@ -697,6 +726,15 @@ class TaskManager : public TaskFinisherInterface, public TaskResubmissionInterfa
                                 std::vector<ObjectID> *ids_to_release)
      ABSL_LOCKS_EXCLUDED(mu_);

+  /// A wrapper of `retry_task_callback_`. This function should be the only


new function

kevin85421 · 2025-05-15T19:29:11Z

src/ray/core_worker/task_manager.h

@@ -578,6 +602,11 @@ class TaskManager : public TaskFinisherInterface, public TaskResubmissionInterfa
  /// Record OCL metrics.
  void RecordMetrics();

+  /// Returns a vector of task specifications that are ready to be retried.


new function

Signed-off-by: kaihsun <[email protected]>

kevin85421 · 2025-05-16T00:09:41Z

cc @jjyao This is something we discussed a while ago.

edoakes · 2025-05-16T21:19:45Z

@israbbani can you review this pls? (not urgent)

israbbani

I saw this in the code recently and it's definitely an abstraction leak. Really happy to see that you're fixing it!

israbbani · 2025-05-16T00:47:48Z

src/ray/core_worker/task_manager.h

@@ -863,6 +901,10 @@ class TaskManager : public TaskFinisherInterface, public TaskResubmissionInterfa
  /// error).
  worker::TaskEventBuffer &task_event_buffer_;

+  /// Queue of tasks to retry, ordered by their execution time.
+  std::priority_queue<TaskToRetry, std::deque<TaskToRetry>, TaskToRetryDescComparator>
+      to_resubmit_ ABSL_GUARDED_BY(mu_);


You're using resubmit and retry interchangeably. E.g. PopTasksToRetry from the queue to_rebsubmit_. I have no preference between them but lets use either Retry or Resubmit.

israbbani · 2025-05-16T00:50:35Z

src/ray/core_worker/task_manager.h

+/// Sorts TaskToRetry in descending order of the execution time.
+/// Priority queue naturally sorts elements in descending order,
+/// in order to have the tasks ordered by execution time in
+/// ascending order we use a comparator that sorts elements in
+/// descending order. Per docs "Priority queues are a type of container
+/// adaptors, specifically designed such that its first element is always
+/// the greatest of the elements it contains".


Since you're refactoring, do you think we can make this comment more concise?

israbbani · 2025-05-16T00:51:07Z

src/ray/core_worker/task_manager.cc

+    // Retry after a delay to emulate the existing Raylet reconstruction
+    // behaviour. TODO(ekl) backoff exponentially.


Is this TODO still relevant? Can we figure it out and either create a ticket/issue or remove it?

israbbani · 2025-05-16T00:52:00Z

src/ray/core_worker/task_manager.h

@@ -863,6 +901,10 @@ class TaskManager : public TaskFinisherInterface, public TaskResubmissionInterfa
  /// error).
  worker::TaskEventBuffer &task_event_buffer_;

+  /// Queue of tasks to retry, ordered by their execution time.
+  std::priority_queue<TaskToRetry, std::deque<TaskToRetry>, TaskToRetryDescComparator>
+      to_resubmit_ ABSL_GUARDED_BY(mu_);


Can we use a separate mutex to guard this priority_queue since it's updated independently of any other state in task_manager?

israbbani · 2025-05-16T19:21:51Z

src/ray/core_worker/core_worker.cc

-          absl::MutexLock lock(&mutex_);
-          TaskToRetry task_to_retry{current_time_ms() + delay_ms, spec};
-          to_resubmit_.push(std::move(task_to_retry));
+      [this](TaskSpecification &spec) {


Not from your PR, but from before. Binding this to the scope of a lambda is also an abstraction leak too. Can we try to clean this up?

israbbani · 2025-05-19T20:23:06Z

src/ray/core_worker/task_manager.h

+/// Represents a task that needs to be retried at a specific time.
+struct TaskToRetry {
+  /// Time when the task should be retried.
+  int64_t execution_time_ms{};


I would remove the comment and change the name of the variable to something that describes its purpose more accurately.

israbbani · 2025-05-19T20:39:02Z

src/ray/core_worker/test/task_manager_test.cc

+  void AddTaskToResubmitQueue(int64_t execution_time_ms, const TaskSpecification &spec) {
+    absl::MutexLock lock(&manager_.mu_);
+    TaskToRetry task_to_retry = {execution_time_ms, spec};
+    manager_.to_resubmit_.push(std::move(task_to_retry));
+  }
+
+  size_t GetResubmitQueueSize() const {
+    absl::MutexLock lock(&manager_.mu_);
+    return manager_.to_resubmit_.size();
+  }
+


I know it's already a pattern that we use in the codebase, but we should not write tests against private variables and functions of a class.

Instead, use the Public API of the class for testing. See discussion here for more context. Can you use RetryTask to write the test?

israbbani · 2025-05-19T20:45:10Z

src/ray/core_worker/test/task_manager_test.cc

+  auto spec1 = CreateTaskHelper(1, {});
+  AddTaskToResubmitQueue(current_time_ms() - 100, spec1);
+
+  auto spec2 = CreateTaskHelper(1, {});
+  AddTaskToResubmitQueue(current_time_ms() + 100, spec2);
+
+  auto tasks = manager_.PopTasksToRetry();
+  ASSERT_EQ(tasks.size(), 1);
+  ASSERT_EQ(tasks[0].TaskId(), spec1.TaskId());
+
+  ASSERT_EQ(GetResubmitQueueSize(), 1);
+}


This test relies on timing and is going to be flaky.

If your testing process gets descheduled by the OS scheduler for more than 100ms, then PopTasksToRetry() can return both tasks. This is going to be non-trivial to mock, I'd suggest using a much larger time when you submit spec2.

israbbani · 2025-05-19T21:01:42Z

src/ray/core_worker/core_worker.cc

-          absl::MutexLock lock(&mutex_);
-          TaskToRetry task_to_retry{current_time_ms() + delay_ms, spec};
-          to_resubmit_.push(std::move(task_to_retry));
+      [this](TaskSpecification &spec) {


Should this be renamed to object_reconstruction_callback instead since it's only called on a task retry where object_reconstruction is true?

israbbani · 2025-05-19T21:04:21Z

src/ray/core_worker/test/task_manager_test.cc

    ASSERT_TRUE(manager_.IsTaskPending(spec.TaskId()));
    ASSERT_TRUE(reference_counter_->IsObjectPendingCreation(return_id));
    ASSERT_EQ(reference_counter_->NumObjectIDsInScope(), 3);
    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);


I don't follow this. Isn't the number of attempts incremented in RetryTask and not retry_task_execution?

github-actions · 2025-06-09T00:35:06Z

This pull request has been automatically marked as stale because it has not had
any activity for 14 days. It will be closed in another 14 days if no further activity occurs.
Thank you for your contributions.

You can always ask for help on our discussion forum or Ray's public slack channel.

If you'd like to keep this open, just leave any comment, and the stale label will be removed.

kevin85421 · 2025-06-09T03:02:26Z

not stale

github-actions · 2025-06-25T00:32:49Z

This pull request has been automatically marked as stale because it has not had
any activity for 14 days. It will be closed in another 14 days if no further activity occurs.
Thank you for your contributions.

You can always ask for help on our discussion forum or Ray's public slack channel.

If you'd like to keep this open, just leave any comment, and the stale label will be removed.

github-actions · 2025-07-09T00:40:46Z

This pull request has been automatically closed because there has been no more activity in the 14 days
since being marked stale.

Please feel free to reopen or open a new pull request if you'd still like this to be addressed.

Again, you can always ask for help on our discussion forum or Ray's public slack channel.

Thanks again for your contribution!

kevin85421 added 5 commits May 5, 2025 02:16

refactor

640cf4f

Signed-off-by: kaihsun <[email protected]>

refactor

a12d6dd

Signed-off-by: kaihsun <[email protected]>

refactor

69e0dfc

Signed-off-by: kaihsun <[email protected]>

add unit tests

72ab5ca

Signed-off-by: kaihsun <[email protected]>

fix tests

097fb78

Signed-off-by: kaihsun <[email protected]>

kevin85421 commented May 5, 2025

View reviewed changes

kevin85421 added 2 commits May 15, 2025 19:13

Merge remote-tracking branch 'upstream/master' into 20250504-devbox1-…

969cb2c

…tmux3-ray2 Signed-off-by: kaihsun <[email protected]>

update

0c282df

Signed-off-by: kaihsun <[email protected]>

kevin85421 commented May 15, 2025

View reviewed changes

update

2bec81c

Signed-off-by: kaihsun <[email protected]>

kevin85421 added the go add ONLY when ready to merge, run all tests label May 15, 2025

kevin85421 marked this pull request as ready for review May 15, 2025 23:36

kevin85421 assigned jjyao and edoakes May 16, 2025

kevin85421 changed the title ~~[core] Move to_resubmit_ from CoreWorker to TaskManager to avoid an abstraction leak~~ [core][refactor] Move to_resubmit_ from CoreWorker to TaskManager to avoid an abstraction leak May 16, 2025

edoakes requested a review from a team May 16, 2025 21:19

israbbani self-assigned this May 19, 2025

israbbani reviewed May 19, 2025

View reviewed changes

hainesmichaelc added community-backlog and removed community-backlog labels May 22, 2025

github-actions bot added the stale The issue is stale. It will be closed within 7 days unless there are further conversation label Jun 9, 2025

github-actions bot removed the stale The issue is stale. It will be closed within 7 days unless there are further conversation label Jun 10, 2025

github-actions bot added the stale The issue is stale. It will be closed within 7 days unless there are further conversation label Jun 25, 2025

github-actions bot closed this Jul 9, 2025

		// Retry after a delay to emulate the existing Raylet reconstruction
		// behaviour. TODO(ekl) backoff exponentially.

[core][refactor] Move to_resubmit_ from CoreWorker to TaskManager to avoid an abstraction leak #52779

[core][refactor] Move to_resubmit_ from CoreWorker to TaskManager to avoid an abstraction leak #52779

Uh oh!

Conversation

kevin85421 commented May 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Why are these changes needed?

Related issue number

Checks

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kevin85421 commented May 16, 2025

Uh oh!

edoakes commented May 16, 2025

Uh oh!

israbbani left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented Jun 9, 2025

Uh oh!

kevin85421 commented Jun 9, 2025

Uh oh!

github-actions bot commented Jun 25, 2025

Uh oh!

github-actions bot commented Jul 9, 2025

Uh oh!

Uh oh!

[core][refactor] Move `to_resubmit_` from CoreWorker to TaskManager to avoid an abstraction leak #52779

[core][refactor] Move `to_resubmit_` from CoreWorker to TaskManager to avoid an abstraction leak #52779

kevin85421 commented May 5, 2025 •

edited

Loading