Make a tool to stop schedulers gracefully

pipelet/launchers.py

@@ -558,6 +558,8 @@ if __name__ == "__main__":
                       help='worker log filename')
     parser.add_option('-N', '--maximum-number-of-tasks', dest='max_number_tasks', default=None,type=int,
                       help='Make worker exit gracefully after a fix number of jobs')
+    parser.add_option('-A', '--abort', dest='abort', default=False, action='store_true',
+                      help='Ask the scheduler to stop')
 #    parser.add_option('-L', '--log-level', metavar='level',
 #                      type=
 #                      help='worker log filename')
@@ -567,7 +569,11 @@ if __name__ == "__main__":
     # print options.secret
     name  = options.logfile.split("_")[-1]
     wl    = init_logger(name, options.logfile)
-    w = worker.ProcessWorker(address=(options.host, options.port),
-                             authkey=options.secret, logger=wl)
+    if options.abort:
+        w = worker.StopWorker(address=(options.host, options.port),
+                              authkey=options.secret, logger=wl)
+    else:
+        w = worker.ProcessWorker(address=(options.host, options.port),
+                                 authkey=options.secret, logger=wl)
     sys.argv = [sys.argv[0]]  # do not propage command line argument to script
     w.run(maximum_task_number=options.max_number_tasks)

pipelet/worker.py

@@ -339,3 +339,9 @@ class _FakeWorker(ProcessWorker):
     def execute_task(self, task):
         task.task_output = range(10)
         return task
+
+class StopWorker(ProcessWorker):
+    def run(self, **keys):
+        self.logger.info('Asking the scheduler to stop')
+        self.scheduler.abort_scheduling(None)
+        self.logger.info('Stop request send')