From 3ae715d92c3f606c89ff93cbb03880e4f2052cba Mon Sep 17 00:00:00 2001
From: Bo Peng <ben.bog@gmail.com>
Date: Thu, 6 Dec 2018 17:37:29 -0600
Subject: [PATCH] Rough implementation of %task kill #154

---
 src/sos_notebook/kernel.py | 10 ----------
 src/sos_notebook/magics.py | 19 +++++++++++++------
 2 files changed, 13 insertions(+), 16 deletions(-)

diff --git a/src/sos_notebook/kernel.py b/src/sos_notebook/kernel.py
index 76d0a98..a046417 100644
--- a/src/sos_notebook/kernel.py
+++ b/src/sos_notebook/kernel.py
@@ -496,16 +496,6 @@ def handle_frontend_msg(msg):
                     self.subkernels.notify_frontend()
                 elif k == 'set-editor-kernel':
                     self.editor_kernel = v
-                elif k == 'kill-task':
-                    # kill specified task
-                    from sos.hosts import Host
-                    Host(v[1])._task_engine.kill_tasks([v[0]])
-                    self.send_frontend_msg('task_status',
-                                           {
-                                               'task_id': v[0],
-                                               'queue': v[1],
-                                               'status': 'abort'
-                                           })
                 elif k == 'cancel-workflow':
                     from .workflow_executor import cancel_workflow
                     cancel_workflow(v[0], self)
diff --git a/src/sos_notebook/magics.py b/src/sos_notebook/magics.py
index 0b410d9..3507147 100644
--- a/src/sos_notebook/magics.py
+++ b/src/sos_notebook/magics.py
@@ -2033,16 +2033,15 @@ def status(self, args):
                 try:
                     # return creation time, start time, and duration
                     tid, tags, _, tst = line.split('\t')
-                    env.tapping_listener_socket.send_pyobj({
-                        'msg_type': 'task_status',
-                        'data': {
+                    self.sos_kernel.send_frontend_msg('task_status',
+                        {
                             'update_only': True,
                             'queue': args.queue,
                             'task_id': tid,
                             'status': tst,
                             'tags': tags
                         }
-                    })
+                    )
                 except Exception as e:
                     env.logger.warning(
                         f'Unrecognized response "{line}" ({e.__class__.__name__}): {e}')
@@ -2052,8 +2051,16 @@ def resume(self, args):
         return
 
     def kill(self, args):
-        self.sos_kernel.warn(args)
-        return
+        # kill specified task
+        from sos.hosts import Host
+        Host(args.queue)._task_engine.kill_tasks(args.tasks)
+        for tid in args.tasks:
+            self.sos_kernel.send_frontend_msg('task_status',
+               {
+                   'task_id': tid,
+                   'queue': args.queue,
+                   'status': 'abort'
+               })
 
     def purge(self, args):
         self.sos_kernel.warn(args)