Working_new_impl

python/ray/_raylet.pyx

@@ -555,6 +555,8 @@ cdef void async_plasma_callback(CObjectID object_id,
 
 cdef c_bool kill_main_task() nogil:
     with gil:
+        x = asyncio.get_event_loop()
+        print(x)
         if setproctitle.getproctitle() == "ray::IDLE":
             return False
         _thread.interrupt_main()

src/ray/core_worker/core_worker.cc

@@ -99,7 +99,6 @@ CoreWorker::CoreWorker(const WorkerType worker_type, const Language language,
       client_call_manager_(new rpc::ClientCallManager(io_service_)),
       death_check_timer_(io_service_),
       internal_timer_(io_service_),
-      kill_retry_timer_(io_service_),
       core_worker_server_(WorkerTypeString(worker_type), 0 /* let grpc choose a port */),
       task_queue_length_(0),
       num_executed_tasks_(0),
@@ -903,7 +902,8 @@ Status CoreWorker::KillTask(const ObjectID &object_id) {
   if (task_manager_->IsTaskPending(task_id)) {
     auto task_spec = task_manager_->GetTaskSpec(object_id.TaskId());
     if (!task_spec.IsActorCreationTask())
-      return direct_task_submitter_->KillTask(task_spec);
+      RAY_RETURN_NOT_OK(direct_task_submitter_->KillTask(task_spec));
+    RAY_LOG(ERROR) << "YEET";
   }
   return Status::OK();
 }
@@ -1096,6 +1096,7 @@ Status CoreWorker::ExecuteTask(const TaskSpecification &task_spec,
   {
     absl::MutexLock lock(&mutex_);
     current_task_ = task_spec;
+    TryKillTask();
   }
 
   RayFunction func{task_spec.GetLanguage(), task_spec.FunctionDescriptor()};
@@ -1402,23 +1403,18 @@ void CoreWorker::HandleWaitForRefRemoved(const rpc::WaitForRefRemovedRequest &re
                                             owner_address, ref_removed_callback);
 }
 
-void CoreWorker::TryKillTask(const TaskID &task_id, int num_tries) {
-  if (num_tries == 0) {
-    RAY_LOG(WARNING) << "Failed to Kill Task: " << task_id;
+void CoreWorker::TryKillTask() {
+  if (task_to_kill_.IsNil()) {
     return;
   }
-  {
-    absl::MutexLock lock(&mutex_);
-    if (main_thread_task_id_ == task_id) {
-      RAY_LOG(INFO) << "Interrupting main: " << main_thread_task_id_;
-      if (kill_main_thread_()) {
-        return;
-      }
+  if (main_thread_task_id_ == task_to_kill_) {
+    RAY_LOG(INFO) << "Killing worker: " << main_thread_task_id_;
+    RAY_IGNORE_EXPR(local_raylet_client_->Disconnect());
+    if (log_dir_ != "") {
+      RayLog::ShutDownRayLog();
     }
+    exit(1);
   }
-  kill_retry_timer_.expires_from_now(boost::asio::chrono::milliseconds(500));
-  kill_retry_timer_.async_wait(
-      boost::bind(&CoreWorker::TryKillTask, this, task_id, num_tries - 1));
 }
 
 void CoreWorker::HandleKillTask(const rpc::KillTaskRequest &request,
@@ -1427,7 +1423,24 @@ void CoreWorker::HandleKillTask(const rpc::KillTaskRequest &request,
   TaskID task_id = TaskID::FromBinary(request.intended_task_id());
   // Try to kill the task immediately, if it fails, it will try again twice. This it to
   // avoid the situation where a cancellation RPC is processed before the execute RPC.
-  TryKillTask(task_id, 3);
+  absl::MutexLock lock(&mutex_);
+  task_to_kill_ = task_id;
+  TryKillTask();
+  // Exit(true);
+  // RAY_LOG(ERROR) << "Trying to kill";
+  // task_manager_->CancelTask(task_id);
+  // RAY_LOG(ERROR) << "Trying to disc";
+  // Disconnect();
+  // RAY_LOG(ERROR) << "Trying to shutdown";
+  // Shutdown();
+  // return;
+  // absl::MutexLock lock(&mutex_);
+  // if (main_thread_task_id_ == task_id) {
+  //   RAY_LOG(INFO) << "Interrupting main: " << main_thread_task_id_;
+  //   if (kill_main_thread_()) {
+  //     return;
+  //   }
+  // }
 }
 
 void CoreWorker::HandleKillActor(const rpc::KillActorRequest &request,

src/ray/core_worker/core_worker.h

@@ -685,7 +685,7 @@ class CoreWorker : public rpc::CoreWorkerServiceHandler {
   }
 
   /// Try to kill a task multiple times.
-  void TryKillTask(const TaskID &task_id, int num_tries);
+  void TryKillTask() EXCLUSIVE_LOCKS_REQUIRED(mutex_);
 
   /// Type of this worker (i.e., DRIVER or WORKER).
   const WorkerType worker_type_;
@@ -734,6 +734,10 @@ class CoreWorker : public rpc::CoreWorkerServiceHandler {
   /// worker context.
   TaskID main_thread_task_id_;
 
+  /// Task to kill (this handles the situation where the cancellation rpc
+  // reaches the remort worker before the task begins executing)
+  TaskID task_to_kill_ GUARDED_BY(mutex_) = TaskID::Nil();
+
   // Flag indicating whether this worker has been shut down.
   bool shutdown_ = false;
 
@@ -752,8 +756,6 @@ class CoreWorker : public rpc::CoreWorkerServiceHandler {
   /// Timer for internal book-keeping.
   boost::asio::steady_timer internal_timer_;
 
-  boost::asio::steady_timer kill_retry_timer_;
-
   /// RPC server used to receive tasks to execute.
   rpc::GrpcServer core_worker_server_;
 

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);
+    RAY_CHECK(it->second.submitted_task_ref_count > 0) << argument_id;
     it->second.submitted_task_ref_count--;
     if (release_lineage) {
       if (it->second.lineage_ref_count > 0) {

src/ray/core_worker/task_manager.cc

@@ -202,15 +202,19 @@ void TaskManager::CompletePendingTask(const TaskID &task_id,
   ShutdownIfNeeded();
 }
 
-void TaskManager::CancelTask(const TaskID &task_id) {
+void TaskManager::CancelTask(const TaskID &task_id, bool pending) {
   {
     absl::MutexLock lock(&mu_);
     auto it = submissible_tasks_.find(task_id);
-    if (it != submissible_tasks_.end()) {
-      it->second.num_retries_left = 0;
+    if (it == submissible_tasks_.end()) {
+      return;
     }
+    it->second.num_retries_left = 0;
+    it->second.canceled = true;
+  }
+  if (pending) {
+    MarkPendingTaskFailed(task_id, GetTaskSpec(task_id), rpc::ErrorType::TASK_CANCELLED);
   }
-  PendingTaskFailed(task_id, rpc::ErrorType::TASK_CANCELLED);
 }
 
 void TaskManager::PendingTaskFailed(const TaskID &task_id, rpc::ErrorType error_type,
@@ -234,6 +238,7 @@ void TaskManager::PendingTaskFailed(const TaskID &task_id, rpc::ErrorType error_
     if (num_retries_left == 0) {
       submissible_tasks_.erase(it);
       num_pending_tasks_--;
+      error_type = it->second.canceled ? rpc::ErrorType::TASK_CANCELLED : error_type;
     } else {
       RAY_CHECK(it->second.num_retries_left > 0);
       it->second.num_retries_left--;
@@ -374,7 +379,7 @@ void TaskManager::RemoveLineageReference(const ObjectID &object_id,
 void TaskManager::MarkPendingTaskFailed(const TaskID &task_id,
                                         const TaskSpecification &spec,
                                         rpc::ErrorType error_type) {
-  RAY_LOG(DEBUG) << "Treat task as failed. task_id: " << task_id
+  RAY_LOG(ERROR) << "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++) {

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 CancelTask(const TaskID &task_id) = 0;
+  virtual void CancelTask(const TaskID &task_id, bool pending) = 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 CancelTask(const TaskID &task_id);
+  void CancelTask(const TaskID &task_id, bool pending);
 
   /// Return the spec for a pending task.
   TaskSpecification GetTaskSpec(const TaskID &task_id) const;
@@ -178,6 +178,8 @@ class TaskManager : public TaskFinisherInterface {
     //    pending tasks and tasks that finished execution but that may be
     //    retried in the future.
     absl::flat_hash_set<ObjectID> reconstructable_return_ids;
+
+    bool canceled = false;
   };
 
   /// Remove a lineage reference to this object ID. This should be called

src/ray/core_worker/transport/direct_task_transport.cc

@@ -24,6 +24,12 @@ Status CoreWorkerDirectTaskSubmitter::SubmitTask(TaskSpecification task_spec) {
   resolver_.ResolveDependencies(task_spec, [this, task_spec]() {
     RAY_LOG(DEBUG) << "Task dependencies resolved " << task_spec.TaskId();
     absl::MutexLock lock(&mu_);
+    if (canceled_resolving_tasks_.find(task_spec.TaskId()) !=
+        canceled_resolving_tasks_.end()) {
+      canceled_resolving_tasks_.erase(task_spec.TaskId());
+
+      return;
+    }
     // Note that the dependencies in the task spec are mutated to only contain
     // plasma dependencies after ResolveDependencies finishes.
     const SchedulingKey scheduling_key(
@@ -234,37 +240,56 @@ Status CoreWorkerDirectTaskSubmitter::KillTask(TaskSpecification task_spec) {
   const SchedulingKey scheduling_key(
       task_spec.GetSchedulingClass(), task_spec.GetDependencies(),
       task_spec.IsActorCreationTask() ? task_spec.ActorCreationId() : ActorID::Nil());
-  absl::MutexLock lock(&mu_);
-  auto scheduled_tasks = task_queues_.find(scheduling_key);
-
-  // See if task has not been shipped yet
-  if (scheduled_tasks != task_queues_.end()) {
-    for (auto spec = scheduled_tasks->second.begin();
-         spec != scheduled_tasks->second.end(); spec++) {
-      if (spec->TaskId() == task_spec.TaskId()) {
-        scheduled_tasks->second.erase(spec);
-        task_finisher_->CancelTask(task_spec.TaskId());
-
-        // Erase an empty queue
-        if (scheduled_tasks->second.size() == 0) {
-          task_queues_.erase(scheduling_key);
+  std::shared_ptr<rpc::CoreWorkerClientInterface> client = nullptr;
+  bool full_kill = false;
+  {
+    absl::MutexLock lock(&mu_);
+    auto scheduled_tasks = task_queues_.find(scheduling_key);
+    RAY_LOG(ERROR) << "First place";
+    // See if task has not been shipped yet
+    if (scheduled_tasks != task_queues_.end()) {
+      for (auto spec = scheduled_tasks->second.begin();
+           spec != scheduled_tasks->second.end(); spec++) {
+        if (spec->TaskId() == task_spec.TaskId()) {
+          scheduled_tasks->second.erase(spec);
+          RAY_LOG(ERROR) << "Canceling task here";
+          full_kill = true;
+          // Erase an empty queue
+          if (scheduled_tasks->second.empty()) {
+            task_queues_.erase(scheduling_key);
+          }
+          break;
         }
-        return Status::OK();
+      }
+    } else {
+      auto rpc_client = sent_tasks_.find(task_spec.TaskId());
+      // No RPC handle for worker
+      if (rpc_client == sent_tasks_.end() ||
+          client_cache_.find(rpc_client->second) == client_cache_.end()) {
+        // Task has dependencies being resolved
+        canceled_resolving_tasks_.emplace(task_spec.TaskId());
+      } else {
+        client = client_cache_.find(rpc_client->second)->second;
       }
     }
   }
-
-  auto rpc_client = sent_tasks_.find(task_spec.TaskId());
-  // No RPC handle for worker
-  if (rpc_client == sent_tasks_.end() ||
-      client_cache_.find(rpc_client->second) == client_cache_.end()) {
+  if (client == nullptr) {
+    if (full_kill) {
+      task_finisher_->CancelTask(task_spec.TaskId(), false);
+      task_finisher_->PendingTaskFailed(task_spec.TaskId(),
+                                        rpc::ErrorType::TASK_CANCELLED);
+    } else {
+      task_finisher_->CancelTask(task_spec.TaskId(), true);
+    }
     return Status::OK();
   }
 
-  auto client = client_cache_.find(rpc_client->second);
+  RAY_LOG(ERROR) << "Second place";
+  task_finisher_->CancelTask(task_spec.TaskId(), false);
+  RAY_LOG(ERROR) << "third place";
   auto request = rpc::KillTaskRequest();
   request.set_intended_task_id(task_spec.TaskId().Binary());
-  return client->second->KillTask(request, nullptr);
+  return client->KillTask(request, nullptr);
 }
 
 };  // namespace ray

src/ray/core_worker/transport/direct_task_transport.h

@@ -174,6 +174,9 @@ class CoreWorkerDirectTaskSubmitter {
   absl::flat_hash_map<SchedulingKey, std::deque<TaskSpecification>> task_queues_
       GUARDED_BY(mu_);
 
+  // Tasks that were canceled while being resolved.
+  absl::flat_hash_set<TaskID> canceled_resolving_tasks_ GUARDED_BY(mu_);
+
   // Keeps track of where currently executing tasks are being run.
   absl::flat_hash_map<TaskID, rpc::WorkerAddress> sent_tasks_ GUARDED_BY(mu_);
 };