def test_is_finished_task_catch_os_error(self):
t = Task.objects.create(
worker=self._worker.worker,
arch=self._arch,
channel=self._channel,
owner=self._user,
method='DummyForkTask',
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()
err = OSError()
err.errno = errno.ECHILD
with patch('kobo.worker.taskmanager.os',
waitpid=Mock(side_effect=err)) as os_mock:
self.assertFalse(tm.is_finished_task(t.id))
os_mock.waitpid.assert_called_once()
def test_take_task_set_load_and_ready(self):
t = Task.objects.create(
worker=self._worker.worker,
arch=self._arch,
channel=self._channel,
owner=self._user,
method='DummyHeavyTask',
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()
self.assertEqual(tm.worker_info['current_load'], 100)
self.assertFalse(tm.worker_info['ready'])
# reload task info
t = Task.objects.get(id=t.id)
self.assertEqual(t.state, TASK_STATES['OPEN'])
def test_fork_task_runs_task_if_cant_fork(self):
t = Task.objects.create(
worker=self._worker.worker,
arch=self._arch,
channel=self._channel,
owner=self._user,
method='DummyForegroundTask',
state=TASK_STATES['OPEN'],
)
self.assertEqual(t.state, TASK_STATES['OPEN'])
tm = TaskManager(conf={'worker': self._worker})
task_info = t.export(False)
with patch('kobo.worker.taskmanager.os', fork=Mock(return_value=0)) as os_mock:
os_mock.devnull = os.devnull
with patch('kobo.worker.taskmanager.signal') as signal_mock:
tm.fork_task(task_info)
os_mock.fork.assert_called_once()
os_mock.setpgrp.assert_called_once()
os_mock._exit.assert_called_once()
signal_mock.signal.assert_called()
t = Task.objects.get(id=t.id)
self.assertEqual(t.state, TASK_STATES['CLOSED'])
def test_take_task_if_wrong_channel(self):
channel = Channel.objects.create(name='channel_x86')
t = Task.objects.create(
worker=self._worker.worker,
arch=self._arch,
channel=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_not_called()
# reload task info
t = Task.objects.get(id=t.id)
self.assertEqual(t.state, TASK_STATES['FREE'])
def test_take_task_catch_errors_if_cant_open_task(self):
t = Task.objects.create(
worker=self._worker.worker,
arch=self._arch,
channel=self._channel,
owner=self._user,
method='DummyForegroundTask',
state=TASK_STATES['FREE'],
)
self.assertEqual(t.state, TASK_STATES['FREE'])
tm = TaskManager(conf={'worker': self._worker})
tm.hub.worker.open_task = Mock(side_effect=ValueError)
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_not_called()
# reload task info
t = Task.objects.get(id=t.id)
self.assertEqual(t.state, TASK_STATES['FREE'])
def test_take_task_runs_task_if_foreground(self):
t = Task.objects.create(
worker=self._worker.worker,
arch=self._arch,
channel=self._channel,
owner=self._user,
method='DummyForegroundTask',
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:
os_mock.devnull = os.devnull
tm.take_task(task_info)
os_mock.fork.assert_not_called()
# reload task info
t = Task.objects.get(id=t.id)
self.assertEqual(t.state, TASK_STATES['CLOSED'])
def test_fork_task_logs_exceptions(self):
"""Exceptions from the child within fork_task are logged."""
t = Task.objects.create(
worker=self._worker.worker,
arch=self._arch,
channel=self._channel,
owner=self._user,
method='DummyForegroundTask',
state=TASK_STATES['OPEN'],
)
logger = Mock()
with patch('kobo.worker.taskmanager.HubProxy') as hub_mock:
# Arrange for close_task call to fail (at end of task)
hub_mock.return_value.worker.close_task.side_effect = RuntimeError("simulated error")
tm = TaskManager(conf={'worker': self._worker}, logger=logger)
task_info = t.export(False)
with patch('kobo.worker.taskmanager.os', fork=Mock(return_value=0)) as os_mock:
os_mock.devnull = os.devnull
tm.fork_task(task_info)
# It should have logged something about the failure to close the task.
logger.log.assert_called_with(logging.CRITICAL, 'Error running forked task', exc_info=1)
def test_finish_task(self):
tm = TaskManager(conf={'worker': self._worker})
mock = Mock()
# inject mock plugin
tm.task_container.plugins['Mock'] = mock
tm.finish_task({'method': 'Mock'})
mock.cleanup.assert_called_once()
mock.notification.assert_called_once()
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_wakeup_task_catch_os_error(self):
t = Task.objects.create(
worker=self._worker.worker,
arch=self._arch,
channel=self._channel,
owner=self._user,
method='DummyForkTask',
state=TASK_STATES['FREE'],
)
tm = TaskManager(conf={'worker': self._worker})
task_info = t.export(False)
task_info['alert'] = True
with patch('kobo.worker.taskmanager.os', fork=Mock(return_value=9999)) as os_mock:
tm.take_task(task_info)
with patch('kobo.worker.taskmanager.os', kill=Mock(side_effect=OSError)) as os_mock:
tm.wakeup_task(task_info)
os_mock.kill.assert_called_once_with(9999, signal.SIGUSR2)
with self.assertRaises(ValueError):
with patch('kobo.worker.taskmanager.os', kill=Mock(side_effect=ValueError)) as os_mock:
tm.wakeup_task(task_info)
os_mock.kill.assert_called_once_with(9999, signal.SIGUSR2)
def test_get_next_task_runs_free_task(self):
t = Task.objects.create(
worker=self._worker.worker,
arch=self._arch,
channel=self._channel,
owner=self._user,
method='DummyForegroundTask',
state=TASK_STATES['FREE'],
)
self.assertEqual(t.state, TASK_STATES['FREE'])
tm = TaskManager(conf={'worker': self._worker})
tm.get_next_task()
# reload task info
t = Task.objects.get(id=t.id)
self.assertEqual(t.state, TASK_STATES['CLOSED'])
def test_fork_task(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.fork_task(task_info)
os_mock.fork.assert_called_once()
def test_run_task_mark_task_as_failed_if_generic_exception(self):
t = Task.objects.create(
worker=self._worker.worker,
arch=self._arch,
channel=self._channel,
owner=self._user,
method='DummyExceptionTask',
state=TASK_STATES['OPEN'],
)
self.assertEqual(t.state, TASK_STATES['OPEN'])
tm = TaskManager(conf={'worker': self._worker})
task_info = t.export(False)
tm.run_task(task_info)
t = Task.objects.get(id=t.id)
self.assertEqual(t.state, TASK_STATES['FAILED'])
def test_run_task_runs_fork_task(self):
t = Task.objects.create(
worker=self._worker.worker,
arch=self._arch,
channel=self._channel,
owner=self._user,
method='DummyForkTask',
state=TASK_STATES['OPEN'],
)
self.assertEqual(t.state, TASK_STATES['OPEN'])
tm = TaskManager(conf={'worker': self._worker})
task_info = t.export(False)
tm.run_task(task_info)
t = Task.objects.get(id=t.id)
self.assertEqual(t.state, TASK_STATES['CLOSED'])
def test_get_next_task_dont_run_tasks_if_not_ready(self):
t = Task.objects.create(
worker=self._worker.worker,
arch=self._arch,
channel=self._channel,
owner=self._user,
method='DummyForegroundTask',
state=TASK_STATES['FREE'],
)
self.assertEqual(t.state, TASK_STATES['FREE'])
tm = TaskManager(conf={'worker': self._worker})
tm.worker_info['ready'] = False
tm.get_next_task()
# reload task info
t = Task.objects.get(id=t.id)
self.assertEqual(t.state, TASK_STATES['FREE'])
def test_update_tasks_interrupt_taks_if_not_on_pid_list(self):
t = Task.objects.create(
worker=self._worker.worker,
arch=self._arch,
channel=self._channel,
owner=self._user,
method='DummyForegroundTask',
state=TASK_STATES['OPEN'],
)
self.assertEqual(t.state, TASK_STATES['OPEN'])
tm = TaskManager(conf={'worker': self._worker})
self.assertFalse(t.id in tm.pid_dict)
tm.update_tasks()
self.assertFalse(t.id in tm.pid_dict)
# reload task info
t = Task.objects.get(id=t.id)
self.assertEqual(t.state, TASK_STATES['INTERRUPTED'])
def test_update_tasks_timeout_task_if_running_for_to_long(self):
t = Task.objects.create(
worker=self._worker.worker,
arch=self._arch,
channel=self._channel,
owner=self._user,
method='DummyForkTask',
timeout=0,
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()
# reload task info
t = Task.objects.get(id=t.id)
self.assertEqual(t.state, TASK_STATES['OPEN'])
self.assertTrue(t.id in tm.pid_dict)
tm.update_tasks()
self.assertFalse(t.id in tm.pid_dict)
# reload task info
t = Task.objects.get(id=t.id)
self.assertEqual(t.state, TASK_STATES['TIMEOUT'])
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 test_update_worker_info_catch_protocol_error(self):
tm = TaskManager(conf={'worker': self._worker})
tm.hub.worker.update_worker = Mock(side_effect=ProtocolError(None, None, None, None))
tm.update_worker_info()
with self.assertRaises(ValueError):
tm.hub.worker.update_worker = Mock(side_effect=ValueError)
tm.update_worker_info()
def test_get_next_task_raise_shutdown_exception_if_locked_and_no_tasks_running(self):
tm = TaskManager(conf={'worker': self._worker})
tm.lock()
# ensure no tasks are running
self.assertEqual(len(self._worker.worker.running_tasks()), 0)
with self.assertRaises(ShutdownException):
tm.get_next_task()
def test_update_worker_info(self):
tm = TaskManager(conf={'worker': self._worker})
self.assertTrue(tm.worker_info['enabled'])
self.assertTrue(tm.worker_info['ready'])
tm.worker_info['enabled'] = False
tm.worker_info['ready'] = False
tm.update_worker_info()
self.assertFalse(tm.worker_info['enabled'])
self.assertFalse(tm.worker_info['ready'])
def test_update_tasks_cleanup_finished_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()
# reload task info
t = Task.objects.get(id=t.id)
self.assertEqual(t.state, TASK_STATES['OPEN'])
tm.run_task(task_info)
# reload task info
t = Task.objects.get(id=t.id)
self.assertEqual(t.state, TASK_STATES['CLOSED'])
with patch('kobo.worker.taskmanager.os',
waitpid=Mock(return_value=(123, 0))) as os_mock:
self.assertTrue(t.id in tm.pid_dict)
tm.update_tasks()
self.assertFalse(t.id in tm.pid_dict)
os_mock.waitpid.assert_called_once()
# reload task info
t = Task.objects.get(id=t.id)
self.assertEqual(t.state, TASK_STATES['CLOSED'])
def test_get_next_task_runs_assigened_awaited_task_if_locked(self):
t1 = Task.objects.create(
worker=self._worker.worker,
arch=self._arch,
channel=self._channel,
owner=self._user,
method='DummyForegroundTask',
state=TASK_STATES['OPEN'],
)
t2 = Task.objects.create(
worker=self._worker.worker,
arch=self._arch,
channel=self._channel,
owner=self._user,
method='DummyForegroundTask',
state=TASK_STATES['ASSIGNED'],
awaited=True,
parent=t1,
)
t3 = Task.objects.create(
worker=self._worker.worker,
arch=self._arch,
channel=self._channel,
owner=self._user,
method='DummyForegroundTask',
waiting=True,
parent=t2,
state=TASK_STATES['FREE'],
)
self.assertEqual(t1.state, TASK_STATES['OPEN'])
self.assertEqual(t2.state, TASK_STATES['ASSIGNED'])
self.assertEqual(t3.state, TASK_STATES['FREE'])
tm = TaskManager(conf={'worker': self._worker})
tm.lock()
tm.get_next_task()
# reload task info
t1 = Task.objects.get(id=t1.id)
self.assertEqual(t1.state, TASK_STATES['OPEN'])
t2 = Task.objects.get(id=t2.id)
self.assertEqual(t2.state, TASK_STATES['CLOSED'])
t3 = Task.objects.get(id=t3.id)
self.assertEqual(t3.state, TASK_STATES['FREE'])
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, ex:
raise
sys.stderr.write("Error initializing TaskManager: %s\n" % ex)
sys.exit(1)
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 test_lock(self):
tm = TaskManager(conf={'worker': self._worker})
self.assertFalse(tm.locked)
tm.lock()
self.assertTrue(tm.locked)
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()
def test_update_tasks_without_tasks(self):
tm = TaskManager(conf={'worker': self._worker})
tm.update_tasks()
self.assertEqual(len(tm.pid_dict.keys()), 0)
self.assertEqual(len(tm.task_dict.keys()), 0)
