[core] Remove worker_context_ dependency from the task receiver (#52740)

This was only used for one very minor thing: setting the current actor
ID. Better to lift that logic into `core_worker.cc` and avoid this
dependency entirely.

---------

Signed-off-by: Edward Oakes <ed.nmi.oakes@gmail.com>

Author: edoakes

src/ray/core_worker/core_worker.cc

@@ -432,7 +432,6 @@ CoreWorker::CoreWorker(CoreWorkerOptions options, const WorkerID &worker_id)
                                   std::placeholders::_7,
                                   std::placeholders::_8);
     task_receiver_ = std::make_unique<TaskReceiver>(
-        worker_context_,
         task_execution_service_,
         *task_event_buffer_,
         execute_task,
@@ -3835,12 +3834,31 @@ void CoreWorker::HandlePushTask(rpc::PushTaskRequest request,
                            send_reply_callback)) {
     return;
   }
+
+  // Set actor info in the worker context.
+  if (request.task_spec().type() == TaskType::ACTOR_CREATION_TASK) {
+    auto actor_id =
+        ActorID::FromBinary(request.task_spec().actor_creation_task_spec().actor_id());
+
+    // Handle duplicate actor creation tasks that might be sent from the GCS on restart.
+    // Ignore the message and reply OK.
+    if (worker_context_.GetCurrentActorID() == actor_id) {
+      RAY_LOG(INFO) << "Ignoring duplicate actor creation task for actor " << actor_id
+                    << ". This is likely due to a GCS server restart.";
+      send_reply_callback(Status::OK(), nullptr, nullptr);
+      return;
+    }
+    worker_context_.SetCurrentActorId(actor_id);
+  }
+
+  // Set job info in the worker context.
   if (request.task_spec().type() == TaskType::ACTOR_CREATION_TASK ||
       request.task_spec().type() == TaskType::NORMAL_TASK) {
     auto job_id = JobID::FromBinary(request.task_spec().job_id());
     worker_context_.MaybeInitializeJobInfo(job_id, request.task_spec().job_config());
     task_counter_.SetJobId(job_id);
   }
+
   // Increment the task_queue_length and per function counter.
   task_queue_length_ += 1;
   std::string func_name =

src/ray/core_worker/test/direct_actor_transport_test.cc

@@ -772,14 +772,6 @@ class MockDependencyWaiter : public DependencyWaiter {
   virtual ~MockDependencyWaiter() {}
 };
 
-class MockWorkerContext : public WorkerContext {
- public:
-  MockWorkerContext(WorkerType worker_type, const JobID &job_id)
-      : WorkerContext(worker_type, WorkerID::FromRandom(), job_id) {
-    current_actor_is_direct_call_ = true;
-  }
-};
-
 class MockTaskEventBuffer : public worker::TaskEventBuffer {
  public:
   void AddTaskEvent(std::unique_ptr<worker::TaskEvent> task_event) override {}
@@ -797,14 +789,12 @@ class MockTaskEventBuffer : public worker::TaskEventBuffer {
 
 class MockTaskReceiver : public TaskReceiver {
  public:
-  MockTaskReceiver(WorkerContext &worker_context,
-                   instrumented_io_context &task_execution_service,
+  MockTaskReceiver(instrumented_io_context &task_execution_service,
                    worker::TaskEventBuffer &task_event_buffer,
                    const TaskHandler &task_handler,
                    std::function<std::function<void()>()> initialize_thread_callback,
                    const OnActorCreationTaskDone &actor_creation_task_done_)
-      : TaskReceiver(worker_context,
-                     task_execution_service,
+      : TaskReceiver(task_execution_service,
                      task_event_buffer,
                      task_handler,
                      initialize_thread_callback,
@@ -819,8 +809,7 @@ class MockTaskReceiver : public TaskReceiver {
 class TaskReceiverTest : public ::testing::Test {
  public:
   TaskReceiverTest()
-      : worker_context_(WorkerType::WORKER, JobID::FromInt(0)),
-        worker_client_(std::make_shared<MockWorkerClient>()),
+      : worker_client_(std::make_shared<MockWorkerClient>()),
         dependency_waiter_(std::make_unique<MockDependencyWaiter>()) {
     auto execute_task = std::bind(&TaskReceiverTest::MockExecuteTask,
                                   this,
@@ -831,7 +820,6 @@ class TaskReceiverTest : public ::testing::Test {
                                   std::placeholders::_5,
                                   std::placeholders::_6);
     receiver_ = std::make_unique<MockTaskReceiver>(
-        worker_context_,
         task_execution_service_,
         task_event_buffer_,
         execute_task,
@@ -867,7 +855,6 @@ class TaskReceiverTest : public ::testing::Test {
 
  private:
   rpc::Address rpc_address_;
-  MockWorkerContext worker_context_;
   instrumented_io_context task_execution_service_;
   MockTaskEventBuffer task_event_buffer_;
   std::shared_ptr<MockWorkerClient> worker_client_;

src/ray/core_worker/transport/task_receiver.cc

@@ -38,20 +38,7 @@ void TaskReceiver::HandleTask(rpc::PushTaskRequest request,
   RAY_CHECK(waiter_ != nullptr) << "Must call init() prior to use";
   TaskSpecification task_spec(std::move(*request.mutable_task_spec()));
 
-  // If GCS server is restarted after sending an actor creation task to this core worker,
-  // the restarted GCS server will send the same actor creation task to the core worker
-  // again. We just need to ignore it and reply ok.
-  if (task_spec.IsActorCreationTask() &&
-      worker_context_.GetCurrentActorID() == task_spec.ActorCreationId()) {
-    send_reply_callback(Status::OK(), nullptr, nullptr);
-    RAY_LOG(INFO) << "Ignoring duplicate actor creation task for actor "
-                  << task_spec.ActorCreationId()
-                  << ". This is likely due to a GCS server restart.";
-    return;
-  }
-
   if (task_spec.IsActorCreationTask()) {
-    worker_context_.SetCurrentActorId(task_spec.ActorCreationId());
     SetupActor(task_spec.IsAsyncioActor(),
                task_spec.MaxActorConcurrency(),
                task_spec.ExecuteOutOfOrder());

src/ray/core_worker/transport/task_receiver.h

@@ -32,7 +32,6 @@
 #include "ray/core_worker/actor_creator.h"
 #include "ray/core_worker/actor_handle.h"
 #include "ray/core_worker/common.h"
-#include "ray/core_worker/context.h"
 #include "ray/core_worker/fiber.h"
 #include "ray/core_worker/store_provider/memory_store/memory_store.h"
 #include "ray/core_worker/transport/actor_scheduling_queue.h"
@@ -63,14 +62,12 @@ class TaskReceiver {
 
   using OnActorCreationTaskDone = std::function<Status()>;
 
-  TaskReceiver(WorkerContext &worker_context,
-               instrumented_io_context &task_execution_service,
+  TaskReceiver(instrumented_io_context &task_execution_service,
                worker::TaskEventBuffer &task_event_buffer,
                TaskHandler task_handler,
                std::function<std::function<void()>()> initialize_thread_callback,
                const OnActorCreationTaskDone &actor_creation_task_done)
-      : worker_context_(worker_context),
-        task_handler_(std::move(task_handler)),
+      : task_handler_(std::move(task_handler)),
         task_execution_service_(task_execution_service),
         task_event_buffer_(task_event_buffer),
         initialize_thread_callback_(std::move(initialize_thread_callback)),
@@ -124,8 +121,6 @@ class TaskReceiver {
   absl::flat_hash_map<ActorID, std::vector<ConcurrencyGroup>> concurrency_groups_cache_;
 
  private:
-  // Worker context.
-  WorkerContext &worker_context_;
   /// The callback function to process a task.
   TaskHandler task_handler_;
   /// The event loop for running tasks on.