def test_shutdown_with_running_tasks(self): t = Task.objects.create( worker=self._worker.worker, arch=self._arch, channel=self._channel, owner=self._user, method='DummyForkTask', state=TASK_STATES['FREE'], ) self.assertEqual(t.state, TASK_STATES['FREE']) tm = TaskManager(conf={'worker': self._worker}) task_info = t.export(False) with patch('kobo.worker.taskmanager.os', fork=Mock(return_value=9999)) as os_mock: tm.take_task(task_info) os_mock.fork.assert_called_once() tm.update_tasks() # reload task info t = Task.objects.get(id=t.id) self.assertEqual(t.state, TASK_STATES['OPEN']) tm.shutdown() # reload task info t = Task.objects.get(id=t.id) self.assertEqual(t.state, TASK_STATES['INTERRUPTED'])
def test_shutdown_without_running_tasks(self): tm = TaskManager(conf={'worker': self._worker}) tm.update_tasks() tm.shutdown() self.assertTrue(tm.worker_info['enabled']) self.assertTrue(tm.worker_info['ready'])
def main_loop(conf, foreground=False, task_manager_class=None): """infinite daemon loop""" # define custom signal handlers signal.signal(signal.SIGTERM, daemon_shutdown) # initialize TaskManager try: log_file = conf.get("LOG_FILE", None) logger = logging.Logger("TaskManager") logger.setLevel(logging.DEBUG) if log_file: log_level = logging.getLevelName(conf.get("LOG_LEVEL", "DEBUG").upper()) kobo.log.add_rotating_file_logger(logger, log_file, log_level=log_level) if not task_manager_class: tm = TaskManager(conf=conf, logger=logger) else: tm = task_manager_class(conf=conf, logger=logger) except Exception as ex: raise sys.stderr.write("Error initializing TaskManager: %s\n" % ex) sys.exit(1) if foreground and tm._logger is not None: kobo.log.add_stderr_logger(tm._logger) while 1: try: tm.log_debug(80 * '-') # poll hub for new tasks tm.hub._login() tm.update_worker_info() tm.update_tasks() tm.get_next_task() # write to stdout / stderr sys.stdout.flush() sys.stderr.flush() # sleep for some time tm.sleep() except (ShutdownException, KeyboardInterrupt): # ignore keyboard interrupts and sigterm signal.signal(signal.SIGINT, signal.SIG_IGN) signal.signal(signal.SIGTERM, signal.SIG_IGN) tm.log_info('Exiting...') tm.shutdown() tm.log_info('Cleanup done...') break except: # this is a little extreme: log the exception and continue traceback = Traceback() tm.log_error(traceback.get_traceback()) tm.sleep()
def main_loop(conf, foreground=False): """infinite daemon loop""" # define custom signal handlers signal.signal(signal.SIGTERM, daemon_shutdown) # initialize TaskManager try: log_file = conf.get("LOG_FILE", None) logger = logging.Logger("TaskManager") logger.setLevel(logging.DEBUG) if log_file: log_level = logging._levelNames.get(conf.get("LOG_LEVEL", "DEBUG").upper()) kobo.log.add_rotating_file_logger(logger, log_file, log_level=log_level) tm = TaskManager(conf=conf, logger=logger) except Exception as ex: raise sys.stderr.write("Error initializing TaskManager: %s\n" % ex) sys.exit(1) if foreground and tm._logger is not None: kobo.log.add_stderr_logger(tm._logger) while 1: try: tm.log_debug(80 * '-') # poll hub for new tasks tm.hub._login() tm.update_worker_info() tm.update_tasks() tm.get_next_task() # write to stdout / stderr sys.stdout.flush() sys.stderr.flush() # sleep for some time tm.sleep() except (ShutdownException, KeyboardInterrupt): # ignore keyboard interrupts and sigterm signal.signal(signal.SIGINT, signal.SIG_IGN) signal.signal(signal.SIGTERM, signal.SIG_IGN) tm.log_info('Exiting...') tm.shutdown() tm.log_info('Cleanup done...') break except: # this is a little extreme: log the exception and continue traceback = Traceback() tm.log_error(traceback.get_traceback()) tm.sleep()