Cleaning up

python/ray/_raylet.pyx

@@ -559,14 +559,13 @@ cdef void async_plasma_callback(CObjectID object_id,
             event_handler._loop.call_soon_threadsafe(
                 event_handler._complete_future, obj_id)
 
-cdef c_bool kill_main_task() nogil:
+cdef void kill_main_task() nogil:
     with gil:
         # This prevents this interrupt from being called to early. It may still
         # be called later, but will not cancel another task because this is
         # called with a C++ mutex lock.
         if setproctitle.getproctitle() != "ray::IDLE":
             _thread.interrupt_main()
-        return True
 
 
 cdef CRayStatus check_signals() nogil:

python/ray/includes/libcoreworker.pxd

@@ -214,8 +214,7 @@ cdef extern from "ray/core_worker/core_worker.h" nogil:
         c_bool ref_counting_enabled
         c_bool is_local_mode
         int num_workers
-        (c_bool() nogil) kill_main
-        # ABCDEFG
+        (void() nogil) kill_main
         CCoreWorkerOptions()
 
     cdef cppclass CCoreWorkerProcess "ray::CoreWorkerProcess":

src/ray/core_worker/core_worker.cc

@@ -1326,6 +1326,7 @@ Status CoreWorker::ExecuteTask(const TaskSpecification &task_spec,
   {
     absl::MutexLock lock(&mutex_);
     current_task_ = task_spec;
+    // Try to force kill if Cancellation RPC arrived before execution started.
     TryForceKillTask();
   }
 
@@ -1668,7 +1669,7 @@ void CoreWorker::TryForceKillTask() {
     return;
   }
   if (main_thread_task_id_ == task_to_kill_) {
-    RAY_LOG(INFO) << "Force killing worker running: " << main_thread_task_id_;
+    RAY_LOG(INFO) << "Force killing a worker running " << main_thread_task_id_;
     RAY_IGNORE_EXPR(local_raylet_client_->Disconnect());
     if (options_.log_dir != "") {
       RayLog::ShutDownRayLog();
@@ -1687,7 +1688,7 @@ void CoreWorker::HandleKillTask(const rpc::KillTaskRequest &request,
   absl::MutexLock lock(&mutex_);
   if (!request.force_kill()) {
     if (main_thread_task_id_ == task_id) {
-      RAY_LOG(INFO) << "Soft killing worker running: " << main_thread_task_id_;
+      RAY_LOG(INFO) << "Interrupting a running task " << main_thread_task_id_;
       options_.kill_main();
     }
     return;

src/ray/core_worker/core_worker.h

@@ -104,7 +104,7 @@ struct CoreWorkerOptions {
   /// Language worker callback to get the current call stack.
   std::function<void(std::string *)> get_lang_stack;
   // Function that tries to interrupt the currently running Python thread.
-  std::function<bool()> kill_main;
+  std::function<void()> kill_main;
   /// Whether to enable object ref counting.
   bool ref_counting_enabled;
   /// Is local mode being used.
@@ -589,6 +589,7 @@ class CoreWorker : public rpc::CoreWorkerServiceHandler {
   /// Stops the task associated with the given Object ID.
   ///
   /// \param[in] object_id of the task to kill (must be a Non-Actor task)
+  /// \param[in] force_kill Whether to force kill a task by killing the worker.
   /// \param[out] Status
   Status KillTask(const ObjectID &object_id, bool force_kill);
   /// Decrease the reference count for this actor. Should be called by the

src/ray/core_worker/reference_count.cc

@@ -298,7 +298,7 @@ void ReferenceCounter::RemoveSubmittedTaskReferences(
                        << argument_id;
       return;
     }
-    RAY_CHECK(it->second.submitted_task_ref_count > 0) << argument_id;
+    RAY_CHECK(it->second.submitted_task_ref_count > 0);
     it->second.submitted_task_ref_count--;
     if (release_lineage) {
       if (it->second.lineage_ref_count > 0) {

src/ray/core_worker/task_manager.cc

@@ -202,7 +202,7 @@ void TaskManager::CompletePendingTask(const TaskID &task_id,
   ShutdownIfNeeded();
 }
 
-void TaskManager::MarkTaskCancelled(const TaskID &task_id, bool store_output) {
+void TaskManager::MarkTaskCancelled(const TaskID &task_id) {
   {
     absl::MutexLock lock(&mu_);
     auto it = submissible_tasks_.find(task_id);
@@ -212,9 +212,7 @@ void TaskManager::MarkTaskCancelled(const TaskID &task_id, bool store_output) {
     it->second.num_retries_left = 0;
     it->second.cancelled = true;
   }
-  if (store_output) {
-    MarkPendingTaskFailed(task_id, GetTaskSpec(task_id), rpc::ErrorType::TASK_CANCELLED);
-  }
+  MarkPendingTaskFailed(task_id, GetTaskSpec(task_id), rpc::ErrorType::TASK_CANCELLED);
 }
 
 void TaskManager::PendingTaskFailed(const TaskID &task_id, rpc::ErrorType error_type,
@@ -379,8 +377,8 @@ void TaskManager::RemoveLineageReference(const ObjectID &object_id,
 void TaskManager::MarkPendingTaskFailed(const TaskID &task_id,
                                         const TaskSpecification &spec,
                                         rpc::ErrorType error_type) {
-  RAY_LOG(ERROR) << "Treat task as failed. task_id: " << task_id
-                 << ", error_type: " << ErrorType_Name(error_type);
+  RAY_LOG(INFO) << "Treat task as failed. task_id: " << task_id
+                << ", error_type: " << ErrorType_Name(error_type);
   int64_t num_returns = spec.NumReturns();
   for (int i = 0; i < num_returns; i++) {
     const auto object_id = ObjectID::ForTaskReturn(

src/ray/core_worker/task_manager.h

@@ -35,7 +35,7 @@ class TaskFinisherInterface {
   virtual void PendingTaskFailed(const TaskID &task_id, rpc::ErrorType error_type,
                                  Status *status = nullptr) = 0;
 
-  virtual void MarkTaskCancelled(const TaskID &task_id, bool store_output) = 0;
+  virtual void MarkTaskCancelled(const TaskID &task_id) = 0;
 
   virtual void OnTaskDependenciesInlined(
       const std::vector<ObjectID> &inlined_dependency_ids,
@@ -113,7 +113,7 @@ class TaskManager : public TaskFinisherInterface {
   /// Cancels a task by treating it as failed and removes any remaining retries
   ///
   /// \param[in] task_id The TaskId that is being cancelled
-  void MarkTaskCancelled(const TaskID &task_id, bool store_output);
+  void MarkTaskCancelled(const TaskID &task_id);
 
   /// Return the spec for a pending task.
   TaskSpecification GetTaskSpec(const TaskID &task_id) const;

src/ray/core_worker/transport/direct_task_transport.cc

@@ -242,7 +242,7 @@ Status CoreWorkerDirectTaskSubmitter::KillTask(TaskSpecification task_spec,
       task_spec.GetSchedulingClass(), task_spec.GetDependencies(),
       task_spec.IsActorCreationTask() ? task_spec.ActorCreationId() : ActorID::Nil());
   std::shared_ptr<rpc::CoreWorkerClientInterface> client = nullptr;
-  bool full_kill = false;
+  bool is_pending = false;
   {
     absl::MutexLock lock(&mu_);
     auto scheduled_tasks = task_queues_.find(scheduling_key);
@@ -253,8 +253,8 @@ Status CoreWorkerDirectTaskSubmitter::KillTask(TaskSpecification task_spec,
            spec != scheduled_tasks->second.end(); spec++) {
         if (spec->TaskId() == task_spec.TaskId()) {
           scheduled_tasks->second.erase(spec);
-          full_kill = true;
-          // Erase an empty queue
+          is_pending = true;
+
           if (scheduled_tasks->second.empty()) {
             task_queues_.erase(scheduling_key);
           }
@@ -273,16 +273,15 @@ Status CoreWorkerDirectTaskSubmitter::KillTask(TaskSpecification task_spec,
       }
     }
   }
+  task_finisher_->MarkTaskCancelled(task_spec.TaskId());
+  if (is_pending) {
+    task_finisher_->PendingTaskFailed(task_spec.TaskId(), rpc::ErrorType::TASK_CANCELLED);
+  }
+
   if (client == nullptr) {
-    task_finisher_->MarkTaskCancelled(task_spec.TaskId(), true);
-    if (full_kill) {
-      task_finisher_->PendingTaskFailed(task_spec.TaskId(),
-                                        rpc::ErrorType::TASK_CANCELLED);
-    }
     return Status::OK();
   }
 
-  task_finisher_->MarkTaskCancelled(task_spec.TaskId(), false);
   auto request = rpc::KillTaskRequest();
   request.set_intended_task_id(task_spec.TaskId().Binary());
   request.set_force_kill(force_kill);