fix: use canceled_tasks_ set only

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

Author: machichima

src/ray/core_worker/context.cc

@@ -98,25 +98,16 @@ struct WorkerThreadContext {
     SetCurrentPlacementGroupId(task_spec.PlacementGroupBundleId().first);
     SetPlacementGroupCaptureChildTasks(task_spec.PlacementGroupCaptureChildTasks());
     current_task_ = std::make_shared<const TaskSpecification>(task_spec);
-
-    is_current_task_canceled_.store(false);
   }
 
   void ResetCurrentTask() {
     SetCurrentTaskId(TaskID::Nil(), /*attempt_number=*/0);
     task_index_ = 0;
     put_counter_ = 0;
-    is_current_task_canceled_.store(false);
   }
 
   uint32_t GetMaxNumGeneratorReturnIndex() const { return max_num_generator_returns_; }
 
-  void SetCurrentTaskCanceled(bool is_canceled) {
-    is_current_task_canceled_.store(is_canceled);
-  }
-
-  bool IsCurrentTaskCanceled() const { return is_current_task_canceled_.load(); }
-
  private:
   /// The task ID for current task.
   TaskID current_task_id_;
@@ -156,12 +147,6 @@ struct WorkerThreadContext {
 
   /// The maximum number of generator return values.
   uint32_t max_num_generator_returns_;
-
-  /// Whether the current task has been canceled via ray.cancel().
-  /// This is set to true when a CancelTask RPC is received for the currently executing task.
-  /// Using atomic because the cancellation request may come from a different thread
-  /// (e.g., via the gRPC handler thread) than the task execution thread.
-  std::atomic<bool> is_current_task_canceled_{false};
 };
 
 thread_local std::unique_ptr<WorkerThreadContext> WorkerContext::thread_context_ =
@@ -490,18 +475,5 @@ WorkerThreadContext &WorkerContext::GetThreadContext() const {
   return *thread_context_;
 }
 
-void WorkerContext::SetCurrentTaskCanceled(bool is_canceled) {
-  if (thread_context_ != nullptr) {
-    thread_context_->SetCurrentTaskCanceled(is_canceled);
-  }
-}
-
-bool WorkerContext::IsCurrentTaskCanceled() const {
-  if (thread_context_ != nullptr) {
-    return thread_context_->IsCurrentTaskCanceled();
-  }
-  return false;
-}
-
 }  // namespace core
 }  // namespace ray

src/ray/core_worker/context.h

@@ -136,9 +136,6 @@ class WorkerContext {
 
   int64_t GetTaskDepth() const;
 
-  void SetCurrentTaskCanceled(bool is_canceled);
-  bool IsCurrentTaskCanceled() const;
-
  private:
   const WorkerType worker_type_;
   const WorkerID worker_id_;

src/ray/core_worker/core_worker.cc

@@ -2495,24 +2495,10 @@ Status CoreWorker::CancelTask(const ObjectID &object_id,
 }
 
 bool CoreWorker::IsTaskCanceled(const TaskID &task_id) const {
-  auto current_task_id = worker_context_->GetCurrentTaskID();
-
-  if (current_task_id == task_id) {
-    // We're checking if the currently executing task is canceled (executor side).
-    // First check the thread-local state (for tasks canceled before execution started).
-    if (worker_context_->IsCurrentTaskCanceled()) {
-      return true;
-    }
-    // Then check the global canceled_tasks_ set (for tasks canceled during execution,
-    // possibly from a different thread via CancelTask RPC).
-    {
-      absl::MutexLock lock(&mutex_);
-      return canceled_tasks_.find(task_id) != canceled_tasks_.end();
-    }
-  }
-
-  // For submitted tasks (submitter side), check TaskManager.
-  return task_manager_->IsTaskCanceled(task_id);
+  // Check if the task is canceled on executor side. Check the canceled_tasks_ which is 
+  // populated when CancelTask RPC is received.
+  absl::MutexLock lock(&mutex_);
+  return canceled_tasks_.find(task_id) != canceled_tasks_.end();
 }
 
 Status CoreWorker::CancelChildren(const TaskID &task_id, bool force_kill) {

src/ray/core_worker/core_worker.h

@@ -1845,9 +1845,12 @@ class CoreWorker {
   /// contexts from GetCoreWorkerStats().
   absl::flat_hash_map<TaskID, TaskSpecification> running_tasks_ ABSL_GUARDED_BY(mutex_);
 
-  /// Tracks which tasks have been marked as canceled. This is needed because cancellation
-  /// requests may come from a different thread (e.g., via gRPC handler) than the thread
-  /// executing the task, so we can't rely solely on thread-local WorkerThreadContext.
+  /// Tracks which tasks have been marked as canceled. For single-threaded, non-async actors this will
+  /// contain at most one task ID.
+  ///
+  /// We have to track this separately because cancellation requests come from gRPC
+  /// thread than the thread executing the task, so we cannot get the cancellation status
+  /// from the thread-local WorkerThreadContext.
   absl::flat_hash_set<TaskID> canceled_tasks_ ABSL_GUARDED_BY(mutex_);
 
   /// Key value pairs to be displayed on Web UI.