TaskCancellation

python/ray/_raylet.pyx

@@ -17,6 +17,8 @@ import logging
 import os
 import pickle
 import sys
+import _thread
+import setproctitle
 
 from libc.stdint cimport (
     int32_t,
@@ -90,6 +92,7 @@ from ray.exceptions import (
     RayTaskError,
     ObjectStoreFullError,
     RayTimeoutError,
+    RayCancellationError
 )
 from ray.utils import decode
 
@@ -449,8 +452,11 @@ cdef execute_task(
             with ray.worker._changeproctitle(title, next_title):
                 with core_worker.profile_event(b"task:execute"):
                     task_exception = True
-                    outputs = function_executor(*args, **kwargs)
-                    task_exception = False
+                    try:
+                        outputs = function_executor(*args, **kwargs)
+                        task_exception = False
+                    except KeyboardInterrupt as e:
+                        raise RayCancellationError("")
                     if c_return_ids.size() == 1:
                         outputs = (outputs,)
             # Store the outputs in the object store.
@@ -466,10 +472,12 @@ cdef execute_task(
             if isinstance(error, RayTaskError):
                 # Avoid recursive nesting of RayTaskError.
                 failure_object = RayTaskError(function_name, backtrace,
-                                              error.cause_cls)
+                                              error.cause_cls, proctitle=title)
+            elif isinstance(error, RayCancellationError):
+                failure_object = RayCancellationError("")
             else:
                 failure_object = RayTaskError(function_name, backtrace,
-                                              error.__class__)
+                                              error.__class__, proctitle=title)
             errors = []
             for _ in range(c_return_ids.size()):
                 errors.append(failure_object)
@@ -545,6 +553,12 @@ cdef void async_plasma_callback(CObjectID object_id,
             event_handler._loop.call_soon_threadsafe(
                 event_handler._complete_future, obj_id)
 
+cdef void kill_main_task() nogil:
+    with gil:
+        if setproctitle.getproctitle() != "ray::IDLE":
+            _thread.interrupt_main()
+
+
 cdef CRayStatus check_signals() nogil:
     with gil:
         try:
@@ -656,7 +670,7 @@ cdef class CoreWorker:
             gcs_options.native()[0], log_dir.encode("utf-8"),
             node_ip_address.encode("utf-8"), node_manager_port,
             task_execution_handler, check_signals, gc_collect,
-            get_py_stack, True))
+            get_py_stack, True, kill_main_task))
 
     def run_task_loop(self):
         with nogil:
@@ -924,6 +938,13 @@ cdef class CoreWorker:
             check_status(self.core_worker.get().KillActor(
                   c_actor_id, True, no_reconstruction))
 
+    def kill_task(self, ObjectID object_id):
+        cdef:
+            CObjectID c_object_id = object_id.native()
+
+        with nogil:
+            check_status(self.core_worker.get().KillTask(c_object_id))
+
     def resource_ids(self):
         cdef:
             ResourceMappingType resource_mapping = (

python/ray/exceptions.py

@@ -16,6 +16,11 @@ class RayConnectionError(RayError):
     pass
 
 
+class RayCancellationError(RayError):
+    """Raised when this task or a dependency is cancelled"""
+    pass
+
+
 class RayTaskError(RayError):
     """Indicates that a task threw an exception during execution.
 

python/ray/includes/libcoreworker.pxd

@@ -99,7 +99,8 @@ cdef extern from "ray/core_worker/core_worker.h" nogil:
                     CRayStatus() nogil,
                     void() nogil,
                     void(c_string *stack_out) nogil,
-                    c_bool ref_counting_enabled)
+                    c_bool ref_counting_enabled,
+                    void() nogil)
         CWorkerType &GetWorkerType()
         CLanguage &GetLanguage()
 
@@ -120,6 +121,7 @@ cdef extern from "ray/core_worker/core_worker.h" nogil:
         CRayStatus KillActor(
             const CActorID &actor_id, c_bool force_kill,
             c_bool no_reconstruction)
+        CRayStatus KillTask(const CObjectID &object_id)
 
         unique_ptr[CProfileEvent] CreateProfileEvent(
             const c_string &event_type)

python/ray/includes/unique_ids.pxd

@@ -122,6 +122,8 @@ cdef extern from "ray/common/id.h" namespace "ray" nogil:
         CTaskID ForNormalTask(CJobID job_id, CTaskID parent_task_id,
                               int64_t parent_task_counter)
 
+        CActorID ActorId() const
+
     cdef cppclass CObjectID" ray::ObjectID"(CBaseID[CObjectID]):
 
         @staticmethod

python/ray/includes/unique_ids.pxi

@@ -219,6 +219,9 @@ cdef class TaskID(BaseID):
     def is_nil(self):
         return self.data.IsNil()
 
+    def actor_id(self):
+        return ActorID(self.data.ActorId().Binary())
+
     cdef size_t hash(self):
         return self.data.Hash()
 

python/ray/serialization.py

@@ -8,13 +8,9 @@
 import ray.utils
 from ray.utils import _random_string
 from ray.gcs_utils import ErrorType
-from ray.exceptions import (
-    PlasmaObjectNotAvailable,
-    RayTaskError,
-    RayActorError,
-    RayWorkerError,
-    UnreconstructableError,
-)
+from ray.exceptions import (PlasmaObjectNotAvailable, RayTaskError,
+                            RayActorError, RayWorkerError,
+                            UnreconstructableError, RayCancellationError)
 from ray._raylet import Pickle5Writer, unpack_pickle5_buffers
 
 logger = logging.getLogger(__name__)
@@ -288,6 +284,8 @@ def _deserialize_object(self, data, metadata, object_id):
                 return RayWorkerError()
             elif error_type == ErrorType.Value("ACTOR_DIED"):
                 return RayActorError()
+            elif error_type == ErrorType.Value("TASK_CANCELLED"):
+                return RayCancellationError()
             elif error_type == ErrorType.Value("OBJECT_UNRECONSTRUCTABLE"):
                 return UnreconstructableError(ray.ObjectID(object_id.binary()))
             else:

python/ray/tests/test_advanced.py

@@ -708,6 +708,27 @@ def method(self):
     assert ray.worker.global_worker.core_worker.object_exists(x_id)
 
 
+def test_task_cancellation(shutdown_only):
+    # import ray, time
+    ray.init(num_cpus=1)
+
+    @ray.remote
+    def micro_sleep_for(t):
+        for _ in range(t * 1000):
+            time.sleep(1 / 1000)
+        return t
+
+    obj1 = micro_sleep_for.remote(100)
+    ray.kill(obj1)
+    ray.get(obj1, 1)
+    obj2 = micro_sleep_for.remote(200)
+    obj3 = micro_sleep_for.remote(300)
+    ray.kill(obj3)
+    ray.get(obj3)
+    ray.kill(obj2)
+    ray.get(obj2)
+
+
 if __name__ == "__main__":
     import pytest
     sys.exit(pytest.main(["-v", __file__]))

python/ray/worker.py

@@ -1371,8 +1371,10 @@ def disconnect(exiting_interpreter=False):
 @contextmanager
 def _changeproctitle(title, next_title):
     setproctitle.setproctitle(title)
-    yield
-    setproctitle.setproctitle(next_title)
+    try:
+        yield
+    finally:
+        setproctitle.setproctitle(next_title)
 
 
 def register_custom_serializer(cls,
@@ -1655,7 +1657,7 @@ def wait(object_ids, num_returns=1, timeout=None):
         return ready_ids, remaining_ids
 
 
-def kill(actor):
+def kill(id):
     """Kill an actor forcefully.
 
     This will interrupt any running tasks on the actor, causing them to fail
@@ -1668,15 +1670,22 @@ def kill(actor):
     If this actor is reconstructable, it will be attempted to be reconstructed.
 
     Args:
-        actor (ActorHandle): Handle to the actor to kill.
+        id (ActorHandle or ObjectID): Handle for the actor to kill or ObjectID
+        of the task to kill.
     """
-    if not isinstance(actor, ray.actor.ActorHandle):
-        raise ValueError("ray.kill() only supported for actors. "
-                         "Got: {}.".format(type(actor)))
-
     worker = ray.worker.global_worker
     worker.check_connected()
-    worker.core_worker.kill_actor(actor._ray_actor_id, False)
+
+    if isinstance(id, ray.actor.ActorHandle):
+        worker.core_worker.kill_actor(id._ray_actor_id, False)
+    elif isinstance(id, ray.ObjectID):
+        if id.task_id().actor_id().hex() in ray.actors().keys():
+            raise ValueError(
+                "Please use ray.kill(ActorHandle) to kill an actor task")
+        worker.core_worker.kill_task(id)
+    else:
+        raise ValueError("ray.kill() only supported for actors and objects. "
+                         "Got: {}.".format(type(id)))
 
 
 def _mode(worker=global_worker):

src/ray/core_worker/core_worker.cc

@@ -84,12 +84,13 @@ CoreWorker::CoreWorker(const WorkerType worker_type, const Language language,
                        std::function<Status()> check_signals,
                        std::function<void()> gc_collect,
                        std::function<void(std::string *)> get_lang_stack,
-                       bool ref_counting_enabled)
+                       bool ref_counting_enabled, std::function<void()> kill_main)
     : worker_type_(worker_type),
       language_(language),
       log_dir_(log_dir),
       ref_counting_enabled_(ref_counting_enabled),
       check_signals_(check_signals),
+      kill_main_thread_(kill_main),
       gc_collect_(gc_collect),
       get_call_site_(RayConfig::instance().record_ref_creation_sites() ? get_lang_stack
                                                                        : nullptr),
@@ -356,10 +357,10 @@ void CoreWorker::RunIOService() {
 
 void CoreWorker::SetCurrentTaskId(const TaskID &task_id) {
   worker_context_.SetCurrentTaskId(task_id);
-  main_thread_task_id_ = task_id;
   bool not_actor_task = false;
   {
     absl::MutexLock lock(&mutex_);
+    main_thread_task_id_ = task_id;
     not_actor_task = actor_id_.IsNil();
   }
   if (not_actor_task && task_id.IsNil()) {
@@ -374,7 +375,7 @@ void CoreWorker::SetCurrentTaskId(const TaskID &task_id) {
 }
 
 void CoreWorker::CheckForRayletFailure() {
-// If the raylet fails, we will be reassigned to init (PID=1).
+  // If the raylet fails, we will be reassigned to init (PID=1).
   if (getppid() == 1) {
     RAY_LOG(ERROR) << "Raylet failed. Shutting down.";
     Shutdown();
@@ -896,6 +897,16 @@ Status CoreWorker::SubmitActorTask(const ActorID &actor_id, const RayFunction &f
   return status;
 }
 
+Status CoreWorker::KillTask(const ObjectID &object_id) {
+  auto task_id = object_id.TaskId();
+  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);
+  }
+  return Status::OK();
+}
+
 Status CoreWorker::KillActor(const ActorID &actor_id, bool force_kill,
                              bool no_reconstruction) {
   ActorHandle *actor_handle = nullptr;
@@ -1078,6 +1089,7 @@ Status CoreWorker::ExecuteTask(const TaskSpecification &task_spec,
   if (resource_ids != nullptr) {
     resource_ids_ = resource_ids;
   }
+  RAY_LOG(ERROR) << "Setting current task id: " << task_spec.TaskId();
   worker_context_.SetCurrentTask(task_spec);
   SetCurrentTaskId(task_spec.TaskId());
 
@@ -1389,6 +1401,23 @@ void CoreWorker::HandleWaitForRefRemoved(const rpc::WaitForRefRemovedRequest &re
   reference_counter_->SetRefRemovedCallback(object_id, contained_in_id, owner_id,
                                             owner_address, ref_removed_callback);
 }
+void CoreWorker::HandleKillTask(const rpc::KillTaskRequest &request,
+                                rpc::KillTaskReply *reply,
+                                rpc::SendReplyCallback send_reply_callback) {
+  TaskID intended_task_id = TaskID::FromBinary(request.intended_task_id());
+  for (int i = 0; i < 3; i++) {
+    {
+      absl::MutexLock lock(&mutex_);
+      if (main_thread_task_id_ == intended_task_id) {
+        RAY_LOG(ERROR) << "Calling Kill Main";
+        kill_main_thread_();
+        return;
+      }
+    }
+    RAY_LOG(ERROR) << "Failed: " << i;
+    usleep(5 * 1000);
+  }
+}
 
 void CoreWorker::HandleKillActor(const rpc::KillActorRequest &request,
                                  rpc::KillActorReply *reply,

src/ray/core_worker/core_worker.h

@@ -92,7 +92,8 @@ class CoreWorker : public rpc::CoreWorkerServiceHandler {
              std::function<Status()> check_signals = nullptr,
              std::function<void()> gc_collect = nullptr,
              std::function<void(std::string *)> get_lang_stack = nullptr,
-             bool ref_counting_enabled = false);
+             bool ref_counting_enabled = false,
+             std::function<void()> kill_main = nullptr);
 
   virtual ~CoreWorker();
 
@@ -402,6 +403,11 @@ class CoreWorker : public rpc::CoreWorkerServiceHandler {
   /// \param[out] Status
   Status KillActor(const ActorID &actor_id, bool force_kill, bool no_reconstruction);
 
+  /// 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[out] Status
+  Status KillTask(const ObjectID &object_id);
   /// Decrease the reference count for this actor. Should be called by the
   /// language frontend when a reference to the ActorHandle destroyed.
   ///
@@ -513,6 +519,10 @@ class CoreWorker : public rpc::CoreWorkerServiceHandler {
   void HandleKillActor(const rpc::KillActorRequest &request, rpc::KillActorReply *reply,
                        rpc::SendReplyCallback send_reply_callback) override;
 
+  /// Implements gRPC server handler.
+  void HandleKillTask(const rpc::KillTaskRequest &request, rpc::KillTaskReply *reply,
+                      rpc::SendReplyCallback send_reply_callback);
+
   /// Implements gRPC server handler.
   void HandlePlasmaObjectReady(const rpc::PlasmaObjectReadyRequest &request,
                                rpc::PlasmaObjectReadyReply *reply,
@@ -691,6 +701,8 @@ class CoreWorker : public rpc::CoreWorkerServiceHandler {
   /// 1s) during long-running operations.
   std::function<Status()> check_signals_;
 
+  std::function<void()> kill_main_thread_;
+
   /// Application-language callback to trigger garbage collection in the language
   /// runtime. This is required to free distributed references that may otherwise
   /// be held up in garbage objects.

src/ray/core_worker/task_manager.cc

@@ -202,6 +202,17 @@ void TaskManager::CompletePendingTask(const TaskID &task_id,
   ShutdownIfNeeded();
 }
 
+void TaskManager::CancelTask(const TaskID &task_id) {
+  {
+    absl::MutexLock lock(&mu_);
+    auto it = submissible_tasks_.find(task_id);
+    if (it != submissible_tasks_.end()) {
+      it->second.num_retries_left = 0;
+    }
+  }
+  PendingTaskFailed(task_id, rpc::ErrorType::TASK_CANCELLED);
+}
+
 void TaskManager::PendingTaskFailed(const TaskID &task_id, rpc::ErrorType error_type,
                                     Status *status) {
   // Note that this might be the __ray_terminate__ task, so we don't log

src/ray/core_worker/task_manager.h

@@ -35,6 +35,8 @@ 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 OnTaskDependenciesInlined(
       const std::vector<ObjectID> &inlined_dependency_ids,
       const std::vector<ObjectID> &contained_ids) = 0;
@@ -108,6 +110,8 @@ class TaskManager : public TaskFinisherInterface {
   void OnTaskDependenciesInlined(const std::vector<ObjectID> &inlined_dependency_ids,
                                  const std::vector<ObjectID> &contained_ids) override;
 
+  void CancelTask(const TaskID &task_id);
+
   /// Return the spec for a pending task.
   TaskSpecification GetTaskSpec(const TaskID &task_id) const;