Rough implementation of %task kill #154

vatlab · Dec 6, 2018 · 3ae715d · 3ae715d
1 parent 74305a5
commit 3ae715d
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 16 deletions.
diff --git 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
@@ -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)