def setUp(self):
self.worker = Worker()
self.worker.stop = Mock()
self.worker_scheduler_factory = WorkerSchedulerFactory()
self.worker_scheduler_factory.create_worker = Mock(
return_value=self.worker)
self.worker_scheduler_factory.create_local_scheduler = Mock()
class NoOpTask(luigi.Task):
param = luigi.Parameter()
self.task_a = NoOpTask("a")
self.task_b = NoOpTask("b")
def test_purge_multiple_workers(self):
w = Worker(worker_processes=2, wait_interval=0.01)
t1 = SuicidalWorker(signal.SIGTERM)
t2 = SuicidalWorker(signal.SIGKILL)
w.add(t1)
w.add(t2)
w._run_task(t1.task_id)
w._run_task(t2.task_id)
time.sleep(1.0)
w._handle_next_task()
w._handle_next_task()
w._handle_next_task()
def test_wait_jitter(self, mock_sleep, mock_random):
""" verify configured jitter amount """
mock_random.return_value = 1.0
w = Worker()
x = w._sleeper()
six.next(x)
mock_random.assert_called_with(0, 10.0)
mock_sleep.assert_called_with(2.0)
mock_random.return_value = 2.0
six.next(x)
mock_random.assert_called_with(0, 10.0)
mock_sleep.assert_called_with(3.0)
def test_task_limit_exceeded(self):
w = Worker()
t = ForkBombTask(3, 2)
w.add(t)
w.run()
self.assertFalse(t.complete())
leaf_tasks = [
ForkBombTask(3, 2, branch)
for branch in [(0, 0, 0), (0, 0, 1), (0, 1, 0), (0, 1, 1)]
]
self.assertEqual(
3, sum(t.complete() for t in leaf_tasks),
"should have gracefully completed as much as possible even though the single last leaf didn't get scheduled"
)
def test_purge_hung_worker_override_timeout_time(self, mock_time):
w = Worker(worker_processes=2, wait_interval=0.01, timeout=5)
mock_time.time.return_value = 0
task = HungWorker(worker_timeout=10)
w.add(task)
w._run_task(task.task_id)
mock_time.time.return_value = 10
w._handle_next_task()
self.assertEqual(1, len(w._running_tasks))
mock_time.time.return_value = 11
w._handle_next_task()
self.assertEqual(0, len(w._running_tasks))
def test_purge_hung_worker_default_timeout_time(self, mock_time):
w = Worker(worker_processes=2, wait_interval=0.01, timeout=5)
mock_time.time.return_value = 0
task = HangTheWorkerTask()
w.add(task)
w._run_task(task.task_id)
mock_time.time.return_value = 5
w._handle_next_task()
self.assertEqual(1, len(w._running_tasks))
mock_time.time.return_value = 6
w._handle_next_task()
self.assertEqual(0, len(w._running_tasks))
def test_interleaved_workers3(self):
class A(DummyTask):
def run(self):
logging.debug('running A')
time.sleep(0.1)
super(A, self).run()
a = A()
class B(DummyTask):
def requires(self):
return a
def run(self):
logging.debug('running B')
super(B, self).run()
b = B()
sch = CentralPlannerScheduler(retry_delay=100, remove_delay=1000, worker_disconnect_delay=10)
w = Worker(scheduler=sch, worker_id='X', keep_alive=True, count_uniques=True)
w2 = Worker(scheduler=sch, worker_id='Y', keep_alive=True, count_uniques=True, wait_interval=0.1)
self.assertTrue(w.add(a))
self.assertTrue(w2.add(b))
threading.Thread(target=w.run).start()
self.assertTrue(w2.run())
self.assertTrue(a.complete())
self.assertTrue(b.complete())
w.stop()
w2.stop()
def test_interleaved_workers(self):
class A(DummyTask):
pass
a = A()
class B(DummyTask):
def requires(self):
return a
class ExternalB(ExternalTask):
task_family = "B"
def complete(self):
return False
b = B()
eb = ExternalB()
self.assertEqual(eb.task_id, "B()")
sch = CentralPlannerScheduler(retry_delay=100, remove_delay=1000, worker_disconnect_delay=10)
w = Worker(scheduler=sch, worker_id='X')
w2 = Worker(scheduler=sch, worker_id='Y')
self.assertTrue(w.add(b))
self.assertTrue(w2.add(eb))
logging.debug("RUNNING BROKEN WORKER")
self.assertTrue(w2.run())
self.assertFalse(a.complete())
self.assertFalse(b.complete())
logging.debug("RUNNING FUNCTIONAL WORKER")
self.assertTrue(w.run())
self.assertTrue(a.complete())
self.assertTrue(b.complete())
w.stop()
w2.stop()
def test_stop_worker_kills_subprocesses(self):
with Worker(worker_processes=2) as w:
hung_task = HungWorker()
w.add(hung_task)
w._run_task(hung_task.task_id)
pids = [p.pid for p in w._running_tasks.values()]
self.assertEqual(1, len(pids))
pid = pids[0]
def is_running():
return pid in {p.pid for p in psutil.Process().children()}
self.assertTrue(is_running())
self.assertFalse(is_running())
def test_die_for_non_unique_pending(self):
class A(DummyTask):
def run(self):
logging.debug('running A')
time.sleep(0.1)
super(A, self).run()
a = A()
class B(DummyTask):
def requires(self):
return a
def run(self):
logging.debug('running B')
super(B, self).run()
b = B()
sch = CentralPlannerScheduler(retry_delay=100, remove_delay=1000, worker_disconnect_delay=10)
w = Worker(scheduler=sch, worker_id='X', keep_alive=True, count_uniques=True)
w2 = Worker(scheduler=sch, worker_id='Y', keep_alive=True, count_uniques=True, wait_interval=0.1)
self.assertTrue(w.add(b))
self.assertTrue(w2.add(b))
self.assertEqual(w._get_work()[0], 'A()')
self.assertTrue(w2.run())
self.assertFalse(a.complete())
self.assertFalse(b.complete())
w2.stop()
def test_timeout_handler(self, mock_time):
result = []
@HangTheWorkerTask.event_handler(Event.TIMEOUT)
def store_task(t, error_msg):
self.assertTrue(error_msg)
result.append(t)
w = Worker(worker_processes=2, wait_interval=0.01, timeout=5)
mock_time.time.return_value = 0
task = HangTheWorkerTask(worker_timeout=1)
w.add(task)
w._run_task(task.task_id)
mock_time.time.return_value = 3
w._handle_next_task()
self.assertEqual(result, [task])
def test_disabled_shutdown_hook(self):
w = Worker(scheduler=self.sch,
keep_alive=True,
no_install_shutdown_handler=True)
with w:
try:
# try to kill the worker!
os.kill(os.getpid(), signal.SIGUSR1)
except AttributeError:
raise unittest.SkipTest(
'signal.SIGUSR1 not found on this system')
# try to kill the worker... AGAIN!
t = SuicidalWorker(signal.SIGUSR1)
w.add(t)
w.run()
# task should have stepped away from the ledge, and completed successfully despite all the SIGUSR1 signals
self.assertEqual(list(self.sch.task_list('DONE', '').keys()),
[t.task_id])
def _test_context_manager(self, force_multiprocessing):
CONTEXT_MANAGER_MODULE = b'''
class MyContextManager(object):
def __init__(self, task_process):
self.task = task_process.task
def __enter__(self):
assert not self.task.run_event.is_set(), "the task should not have run yet"
self.task.enter_event.set()
return self
def __exit__(self, exc_type=None, exc_value=None, traceback=None):
assert self.task.run_event.is_set(), "the task should have run"
self.task.exit_event.set()
'''
class DummyEventRecordingTask(luigi.Task):
def __init__(self, *args, **kwargs):
self.enter_event = multiprocessing.Event()
self.exit_event = multiprocessing.Event()
self.run_event = multiprocessing.Event()
super(DummyEventRecordingTask, self).__init__(*args, **kwargs)
def run(self):
assert self.enter_event.is_set(
), "the context manager should have been entered"
assert not self.exit_event.is_set(
), "the context manager should not have been exited yet"
assert not self.run_event.is_set(
), "the task should not have run yet"
self.run_event.set()
def complete(self):
return self.run_event.is_set()
with temporary_unloaded_module(CONTEXT_MANAGER_MODULE) as module_name:
t = DummyEventRecordingTask()
w = Worker(task_process_context=module_name + '.MyContextManager',
force_multiprocessing=force_multiprocessing)
w.add(t)
self.assertTrue(w.run())
self.assertTrue(t.complete())
self.assertTrue(t.enter_event.is_set())
self.assertTrue(t.exit_event.is_set())
def test_process_killed_handler(self, task_proc):
result = []
@HangTheWorkerTask.event_handler(Event.PROCESS_FAILURE)
def store_task(t, error_msg):
self.assertTrue(error_msg)
result.append(t)
w = Worker()
task = HangTheWorkerTask()
task_process = mock.MagicMock(is_alive=lambda: False,
exitcode=-14,
task=task)
task_proc.return_value = task_process
w.add(task)
w._run_task(task.task_id)
w._handle_next_task()
self.assertEqual(result, [task])
def test_connection_error(self):
sch = RemoteScheduler(host="this_host_doesnt_exist", port=1337)
worker = Worker(scheduler=sch)
self.waits = 0
def dummy_wait():
self.waits += 1
sch._wait = dummy_wait
class A(DummyTask):
pass
a = A()
self.assertEquals(self.last_email, None)
worker.add(a)
self.assertEquals(self.waits, 2) # should attempt to add it 3 times
self.assertNotEquals(self.last_email, None)
self.assertEquals(self.last_email[0], "Luigi: Framework error while scheduling %s" % (a,))
worker.stop()
def test_connection_error(self, emails):
sch = RemoteScheduler('http://tld.invalid:1337', connect_timeout=1)
worker = Worker(scheduler=sch)
self.waits = 0
def dummy_wait():
self.waits += 1
sch._wait = dummy_wait
class A(DummyTask):
pass
a = A()
self.assertEqual(emails, [])
worker.add(a)
self.assertEqual(self.waits, 2) # should attempt to add it 3 times
self.assertNotEquals(emails, [])
self.assertTrue(emails[0].find("Luigi: Framework error while scheduling %s" % (a,)) != -1)
worker.stop()
def setUp(self):
self.scheduler = RemoteScheduler()
self.scheduler.add_worker = Mock()
self.scheduler.add_task = Mock()
self.worker = Worker(scheduler=self.scheduler, worker_id='X', worker_processes=2)
def test_ping_thread_shutdown(self):
with Worker(ping_interval=0.01) as w:
self.assertTrue(w._keep_alive_thread.is_alive())
self.assertFalse(w._keep_alive_thread.is_alive())
def test_fails_registering_signal(self):
with mock.patch('luigi.worker.signal', spec=['signal']):
# mock will raise an attribute error getting signal.SIGUSR1
Worker()
def run(self, result=None):
super(WorkerEmailTest, self).setUp()
sch = CentralPlannerScheduler(retry_delay=100, remove_delay=1000, worker_disconnect_delay=10)
with Worker(scheduler=sch, worker_id="foo") as self.worker:
super(WorkerEmailTest, self).run(result)
def test_asserts_for_worker(self):
"""
Test that Worker() asserts that it's sanely configured
"""
Worker(wait_interval=1) # This shouldn't raise
self.assertRaises(AssertionError, Worker, wait_interval=0)
def test_no_task_limit(self):
w = Worker()
t = ForkBombTask(4, 2)
w.add(t)
w.run()
self.assertTrue(t.complete())
def test_task_limit_not_exceeded(self):
w = Worker()
t = ForkBombTask(3, 2)
w.add(t)
w.run()
self.assertTrue(t.complete())
def setUp(self):
super(WorkerEmailTest, self).setUp()
sch = CentralPlannerScheduler(retry_delay=100,
remove_delay=1000,
worker_disconnect_delay=10)
self.worker = Worker(scheduler=sch, worker_id="foo")
def test_ping_thread_shutdown(self):
w = Worker(ping_interval=0.01)
self.assertTrue(w._keep_alive_thread.is_alive())
w.stop() # should stop within 0.01 s
self.assertFalse(w._keep_alive_thread.is_alive())
def setUp(self):
self.sch = mock.Mock()
self.w = Worker(scheduler=self.sch, worker_id='x')
def setUp(self):
# InstanceCache.disable()
self.sch = CentralPlannerScheduler(retry_delay=100, remove_delay=1000, worker_disconnect_delay=10)
self.w = Worker(scheduler=self.sch, worker_id='X')
self.w2 = Worker(scheduler=self.sch, worker_id='Y')
self.time = time.time
def setUp(self):
self.sch = CentralPlannerScheduler(retry_delay=100, remove_delay=1000, worker_disconnect_delay=10)
self.w = Worker(scheduler=self.sch, worker_id='X')
self.assistant = Worker(scheduler=self.sch, worker_id='Y', assistant=True)
def run(self, result=None):
with Worker() as w:
self.w = w
super(Dependency, self).run(result)
def setUp(self):
self.sch = CentralPlannerScheduler()
self.w = Worker(scheduler=self.sch)
