import ctypes import multiprocessing as mp import queue import threading import unittest import warnings from collections import deque from datetime import datetime from itertools import product from unittest.mock import patch from mpire.comms import MAIN_PROCESS, NEW_MAP_PARAMS_PILL, NON_LETHAL_POISON_PILL, POISON_PILL, WorkerComms from mpire.context import DEFAULT_START_METHOD, FORK_AVAILABLE, MP_CONTEXTS from mpire.params import WorkerMapParams def _f1(): pass def _f2(): pass class WorkerCommsTest(unittest.TestCase): def test_init_comms(self): """ Test if initializing/resetting the comms is done properly """ test_ctx = [MP_CONTEXTS['mp_dill']['spawn'], MP_CONTEXTS['threading']] if FORK_AVAILABLE: test_ctx.extend([MP_CONTEXTS['mp_dill']['fork'], MP_CONTEXTS['mp']['forkserver']]) for ctx, n_jobs, order_tasks in product(test_ctx, [1, 2, 4], [False, True]): comms = WorkerComms(ctx, n_jobs, order_tasks) with self.subTest('__init__ called', ctx=ctx, n_jobs=n_jobs, order_tasks=order_tasks): condition_type = type(ctx.Condition(ctx.Lock())) event_type = type(ctx.Event()) lock_type = type(ctx.Lock()) value_type = type(ctx.Value('i', 0, lock=True)) self.assertEqual(comms.ctx, ctx) self.assertEqual(comms.n_jobs, n_jobs) self.assertEqual(comms.order_tasks, order_tasks) self.assertFalse(comms.is_initialized()) self.assertIsInstance(comms._keep_order, value_type) self.assertFalse(comms._keep_order.value) self.assertIsInstance(comms._task_queues, list) self.assertEqual(len(comms._task_queues), 0) self.assertIsNone(comms._task_idx) self.assertIsInstance(comms._worker_running_task, list) self.assertEqual(len(comms._worker_running_task), 0) self.assertIsInstance(comms._last_completed_task_worker_id, deque) self.assertEqual(len(comms._last_completed_task_worker_id), 0) self.assertIsNone(comms._worker_working_on_job) self.assertIsNone(comms._results_queue) self.assertIsInstance(comms._results_added, list) self.assertEqual(len(comms._results_added), 0) self.assertIsNone(comms._results_received) self.assertIsNone(comms._worker_restart_array) self.assertIsInstance(comms._worker_restart_condition, condition_type) self.assertIsNone(comms._workers_dead) self.assertIsNone(comms._workers_time_task_started) self.assertIsInstance(comms.exception_lock, lock_type) self.assertIsInstance(comms._exception_thrown, event_type) self.assertIsNone(comms._exception_job_id) self.assertIsInstance(comms._kill_signal_received, value_type) self.assertFalse(comms._exception_thrown.is_set()) self.assertFalse(comms._kill_signal_received.value) self.assertIsNone(comms._tasks_completed_array) self.assertIsNone(comms._progress_bar_last_updated) self.assertIsNone(comms._progress_bar_shutdown) self.assertIsNone(comms._progress_bar_complete) with self.subTest('without initial values', ctx=ctx, n_jobs=n_jobs, order_tasks=order_tasks), \ patch('time.time', return_value=0.0): comms.init_comms() self._check_comms_are_initialized(comms, n_jobs) # Set values so we can test if the containers will be properly resetted comms._task_idx = 4 comms._last_completed_task_worker_id.append(2) comms._last_completed_task_worker_id.append(1) comms._last_completed_task_worker_id.append(2) for i in range(n_jobs): comms._worker_running_task[i].value = i % 2 == 0 for i in range(n_jobs): comms._worker_working_on_job[i] = i + 1 comms._results_added = [i + 1 for i in range(n_jobs)] for i in range(n_jobs): comms._results_received[i] = i + 1 for i in range(n_jobs): comms._worker_restart_array[i] = i % 2 == 0 comms._workers_dead[:] = [False] * n_jobs for i in range(n_jobs): for j in range(3): comms._workers_time_task_started[i * 3 + j] = i + j + 1 comms._exception_thrown.set() comms._exception_job_id = 89 comms._kill_signal_received.value = True for i in range(n_jobs): comms._tasks_completed_array[i] = i + 1 comms._progress_bar_last_updated = 3 comms._progress_bar_shutdown.value = True comms._progress_bar_complete.set() comms.reset() with self.subTest('with initial values', ctx=ctx, n_jobs=n_jobs, order_tasks=order_tasks), \ patch('time.time', return_value=0.0): comms.init_comms() self._check_comms_are_initialized(comms, n_jobs) # Reset initialized flag comms.reset() self.assertFalse(comms.is_initialized()) def _check_comms_are_initialized(self, comms: WorkerComms, n_jobs: int) -> None: """ Checks if the WorkerComms have been properly initialized :param comms: WorkerComms object :param n_jobs: Number of jobs """ array_type = type(comms.ctx.Array('i', n_jobs, lock=True)) event_type = type(comms.ctx.Event()) joinable_queue_type = type(comms.ctx.JoinableQueue()) rlock_type = type(comms.ctx.RLock()) value_type = type(comms.ctx.Value('i', 0, lock=True)) self.assertEqual(len(comms._task_queues), n_jobs) for q in comms._task_queues: self.assertIsInstance(q, joinable_queue_type) self.assertIsInstance(comms._last_completed_task_worker_id, deque) self.assertEqual(len(comms._last_completed_task_worker_id), 0) self.assertEqual(len(comms._worker_running_task), n_jobs) for v in comms._worker_running_task: self.assertIsInstance(v, value_type) self.assertIsInstance(v.get_lock(), rlock_type) self.assertIsInstance(comms._worker_working_on_job, array_type) self.assertEqual(len(comms._worker_working_on_job), n_jobs) self.assertIsInstance(comms._results_queue, joinable_queue_type) self.assertIsInstance(comms._results_added, list) self.assertEqual(len(comms._results_added), n_jobs) self.assertIsInstance(comms._results_received, array_type) self.assertEqual(len(comms._results_received), n_jobs) self.assertIsInstance(comms._worker_restart_array, array_type) self.assertEqual(len(comms._worker_restart_array), n_jobs) self.assertIsInstance(comms._workers_dead, array_type) self.assertEqual(len(comms._workers_dead), n_jobs) self.assertIsInstance(comms._workers_time_task_started, array_type) self.assertEqual(len(comms._workers_time_task_started), n_jobs * 3) self.assertFalse(comms._exception_thrown.is_set()) self.assertIsInstance(comms._exception_job_id, value_type) self.assertEqual(comms._exception_job_id.value, 0) self.assertFalse(comms._kill_signal_received.value) self.assertIsInstance(comms._tasks_completed_array, array_type) self.assertEqual(len(comms._tasks_completed_array), n_jobs) self.assertEqual(comms._progress_bar_last_updated, 0.0) self.assertIsInstance(comms._progress_bar_shutdown, value_type) self.assertFalse(comms._progress_bar_shutdown.value) self.assertIsInstance(comms._progress_bar_complete, event_type) self.assertFalse(comms._progress_bar_complete.is_set()) self.assertTrue(comms.is_initialized()) # Basic sanity checks for the values self.assertEqual(comms._task_idx, 0) self.assertEqual([v.value for v in comms._worker_running_task], [False] * n_jobs) self.assertEqual(comms._worker_working_on_job[:], [0] * n_jobs) self.assertEqual(comms._results_received[:], [0] * n_jobs) self.assertEqual(comms._worker_restart_array[:], [False] * n_jobs) self.assertEqual(comms._workers_dead[:], [True] * n_jobs) self.assertEqual(comms._workers_time_task_started[:], [0.0] * n_jobs * 3) self.assertEqual(comms._tasks_completed_array[:], [0] * n_jobs) def test_progress_bar(self): """ Test progress bar related functions """ comms = WorkerComms(MP_CONTEXTS['mp'][DEFAULT_START_METHOD], 2, False) comms.init_comms() # Nothing available yet self.assertEqual(sum(comms._tasks_completed_array), 0) # 3 task done, but not enough time has passed to send the update last_updated = 0.0 n_tasks_completed = 0 with patch('time.time', return_value=0.0): for n in range(1, 4): last_updated, n_tasks_completed = comms.task_completed_progress_bar(0, last_updated, n_tasks_completed, force_update=False) self.assertEqual(n_tasks_completed, n) self.assertEqual(sum(comms._tasks_completed_array), 0) # Not enough time has passed, but we'll force the update. Number of tasks done should still be 3 with patch('time.time', return_value=0.0): last_updated, n_tasks_completed = comms.task_completed_progress_bar(0, last_updated, n_tasks_completed, force_update=True) self.assertEqual(comms.get_tasks_completed_progress_bar(), 3) self.assertEqual(last_updated, 0.0) self.assertEqual(n_tasks_completed, 0) # 4 tasks already done and another 4 tasks done. Enough time should've passed for each update call, except the # second. In total we have 3 (from above) + 4 + 4 = 11 tasks done last_updated = 0.0 n_tasks_completed = 4 with patch('time.time', side_effect=[1.0, 1.0, 3.0, 4.0]): for expected_last_updated in [1.0, 1.0, 3.0, 4.0]: last_updated, n_tasks_completed = comms.task_completed_progress_bar(0, last_updated, n_tasks_completed, force_update=False) self.assertEqual(last_updated, expected_last_updated) self.assertEqual(comms.get_tasks_completed_progress_bar(), 11) self.assertEqual(last_updated, 4.0) self.assertEqual(n_tasks_completed, 0) # Signal shutdown comms.signal_progress_bar_shutdown() self.assertEqual(comms.get_tasks_completed_progress_bar(), POISON_PILL) comms._progress_bar_shutdown.value = False # Set exception comms.signal_exception_thrown(MAIN_PROCESS) self.assertEqual(comms.get_tasks_completed_progress_bar(), POISON_PILL) comms._exception_thrown.clear() # Set kill signal received comms.signal_kill_signal_received() self.assertEqual(comms.get_tasks_completed_progress_bar(), POISON_PILL) comms._kill_signal_received.value = False def test_progress_bar_shutdown(self): """ Test progress bar complete related functions """ comms = WorkerComms(MP_CONTEXTS['mp'][DEFAULT_START_METHOD], 2, False) comms.init_comms() self.assertFalse(comms._progress_bar_shutdown.value) comms.signal_progress_bar_shutdown() self.assertTrue(comms._progress_bar_shutdown.value) def test_progress_bar_complete(self): """ Test progress bar complete related functions """ comms = WorkerComms(MP_CONTEXTS['mp'][DEFAULT_START_METHOD], 2, False) comms.init_comms() self.assertFalse(comms._progress_bar_complete.is_set()) comms.signal_progress_bar_complete() self.assertTrue(comms._progress_bar_complete.is_set()) with patch.object(comms._progress_bar_complete, 'wait') as p: comms.wait_until_progress_bar_is_complete() self.assertEqual(p.call_count, 1) def test_keep_order(self): """ Test keep_order related functions """ comms = WorkerComms(MP_CONTEXTS['mp'][DEFAULT_START_METHOD], 2, False) self.assertFalse(comms.keep_order()) comms.signal_keep_order() self.assertTrue(comms.keep_order()) comms.clear_keep_order() self.assertFalse(comms.keep_order()) def test_tasks(self): """ Test task related functions """ comms = WorkerComms(MP_CONTEXTS['mp'][DEFAULT_START_METHOD], 3, False) comms.init_comms() # Nothing available yet for worker_id in range(3): with self.assertRaises(queue.Empty): comms._task_queues[worker_id].get(block=False) # Add a few tasks. As no worker has completed tasks yet, it should give the task to workers in order job_id = 0 comms.add_task(job_id, 12) comms.add_task(job_id, 'hello world') comms.add_task(job_id, {'foo': 'bar'}) comms.add_task(job_id, {'foo': 'baz'}) comms.add_task(job_id, 34.43) comms.add_task(job_id, datetime(2000, 1, 1, 1, 2, 3)) tasks = [] for worker_id in [0, 1, 2, 0, 1, 2]: tasks.append(comms.get_task(worker_id)) comms.task_done(worker_id) self.assertListEqual(tasks, [(job_id, task) for task in [12, 'hello world', {'foo': 'bar'}, {'foo': 'baz'}, 34.43, datetime(2000, 1, 1, 1, 2, 3)]]) # When order_tasks is set it should give the tasks to the workers in the order they were added, independent of # workers that have completed tasks. If set to False and workers have completed tasks, it depends on the last # one who gets the new task. After giving a task to that worker the worker ID that last completed a task is # reset again. So, it should continue with the normal order job_id = 1 for order_tasks, worker_order in [(False, [0, 1, 2, 2, 1, 0, 2, 0, 1]), (True, [0, 1, 2, 0, 1, 2, 0, 1, 2])]: with self.subTest(order_tasks=order_tasks): comms.order_tasks = order_tasks comms.init_comms() comms.add_task(job_id, 12) comms.add_task(job_id, 'hello world') comms.add_task(job_id, {'foo': 'bar'}) comms._last_completed_task_worker_id.append(2) comms.add_task(job_id, {'foo': 'baz'}) comms._last_completed_task_worker_id.append(1) comms.add_task(job_id, 34.43) comms._last_completed_task_worker_id.append(0) comms.add_task(job_id, datetime(2000, 1, 1, 1, 2, 3)) comms._last_completed_task_worker_id.append(2) comms.add_task(job_id, '123') comms.add_task(job_id, 123) comms.add_task(job_id, 1337) tasks = [] for worker_id in worker_order: tasks.append(comms.get_task(worker_id)) comms.task_done(worker_id) self.assertListEqual(tasks, [(job_id, task) for task in [12, 'hello world', {'foo': 'bar'}, {'foo': 'baz'}, 34.43, datetime(2000, 1, 1, 1, 2, 3), '123', 123, 1337]]) # Throw in an exception. Should return None comms.signal_exception_thrown(job_id) for worker_id in range(3): self.assertIsNone(comms.get_task(worker_id)) # Should be joinable comms.join_task_queues() def test_worker_running_task(self): """ Tests that the worker_running_task functions work as expected """ comms = WorkerComms(MP_CONTEXTS['mp'][DEFAULT_START_METHOD], 3, False) comms.init_comms() comms.set_worker_running_task(0, True) comms.set_worker_running_task(1, False) comms.set_worker_running_task(2, True) self.assertTrue(comms.get_worker_running_task(0)) self.assertFalse(comms.get_worker_running_task(1)) self.assertTrue(comms.get_worker_running_task(2)) comms.set_worker_running_task(0, False) comms.set_worker_running_task(1, True) self.assertFalse(comms.get_worker_running_task(0)) self.assertTrue(comms.get_worker_running_task(1)) self.assertTrue(comms.get_worker_running_task(2)) def test_worker_working_on_job(self): """ Tests that the worker_working_on_job function works as expected """ comms = WorkerComms(MP_CONTEXTS['mp'][DEFAULT_START_METHOD], 3, False) comms.init_comms() comms.signal_worker_working_on_job(0, 1) comms.signal_worker_working_on_job(1, 8) comms.signal_worker_working_on_job(2, 19) self.assertEqual(comms.get_worker_working_on_job(0), 1) self.assertEqual(comms.get_worker_working_on_job(1), 8) self.assertEqual(comms.get_worker_working_on_job(2), 19) comms.signal_worker_working_on_job(1, 3) comms.signal_worker_working_on_job(2, 0) comms.signal_worker_working_on_job(0, -1) self.assertEqual(comms.get_worker_working_on_job(0), -1) self.assertEqual(comms.get_worker_working_on_job(1), 3) self.assertEqual(comms.get_worker_working_on_job(2), 0) def test_results(self): """ Test results related functions """ comms = WorkerComms(MP_CONTEXTS['mp'][DEFAULT_START_METHOD], 2, False) comms.init_comms() # Nothing available yet with self.assertRaises(queue.Empty): comms._results_queue.get(block=False) # Add a few results. Note that `get_results` calls `task_done` comms.add_results(0, [(0, True, 12)]) comms.add_results(1, [(1, True, 'hello world')]) comms.add_results(1, [(1, False, {'foo': 'bar'})]) comms.add_results(0, [(2, True, '123')]) self.assertEqual(comms.get_results(), [(0, True, 12)]) self.assertEqual(comms._last_completed_task_worker_id.popleft(), 0) self.assertEqual(comms.get_results(), [(1, True, 'hello world')]) self.assertEqual(comms._last_completed_task_worker_id.popleft(), 1) self.assertEqual(comms.get_results(), [(1, False, {'foo': 'bar'})]) self.assertEqual(comms._last_completed_task_worker_id.popleft(), 1) self.assertEqual(comms.get_results(), [(2, True, '123')]) self.assertEqual(comms._last_completed_task_worker_id.popleft(), 0) # Should be joinable. When using keep_alive, it should still be open comms.join_results_queues(keep_alive=True) comms.add_results(0, [(2, True, 12)]) comms.get_results() comms.join_results_queues(keep_alive=False) # Depending on Python version it either throws AssertionError or ValueError with self.assertRaises((AssertionError, ValueError)): comms.add_results(0, [(2, True, 12)]) def test_add_new_map_params(self): """ Test new map parameters function """ comms = WorkerComms(MP_CONTEXTS['mp_dill'][DEFAULT_START_METHOD], 2, False) comms.init_comms() map_params = WorkerMapParams(_f1, None, None, 1) comms.add_new_map_params(map_params) with warnings.catch_warnings(): warnings.simplefilter("ignore") for worker_id in range(2): self.assertEqual(comms.get_task(worker_id), NEW_MAP_PARAMS_PILL) self.assertEqual(comms.get_task(worker_id), map_params) comms.task_done(worker_id) comms.task_done(worker_id) map_params = WorkerMapParams(_f1, _f2, None, None) comms.add_new_map_params(map_params) with warnings.catch_warnings(): warnings.simplefilter("ignore") for worker_id in range(2): self.assertEqual(comms.get_task(worker_id), NEW_MAP_PARAMS_PILL) self.assertEqual(comms.get_task(worker_id), map_params) comms.task_done(worker_id) comms.task_done(worker_id) def test_exception_thrown(self): """ Test exception thrown related functions """ comms = WorkerComms(MP_CONTEXTS['mp'][DEFAULT_START_METHOD], 2, False) comms.init_comms() self.assertFalse(comms.exception_thrown()) comms.signal_exception_thrown(0) self.assertTrue(comms.exception_thrown()) self.assertEqual(comms.get_exception_thrown_job_id(), 0) comms._exception_thrown.clear() self.assertFalse(comms.exception_thrown()) comms.signal_exception_thrown(13) self.assertTrue(comms.exception_thrown()) self.assertEqual(comms.get_exception_thrown_job_id(), 13) def test_kill_signal_received(self): """ Test kill signal received related functions """ comms = WorkerComms(MP_CONTEXTS['mp'][DEFAULT_START_METHOD], 2, False) self.assertFalse(comms.kill_signal_received()) comms.signal_kill_signal_received() self.assertTrue(comms.kill_signal_received()) comms._kill_signal_received.value = False self.assertFalse(comms.kill_signal_received()) def test_worker_poison_pill(self): """ Test that a poison pill is inserted for every worker """ for n_jobs in [1, 2, 4]: with self.subTest(n_jobs=n_jobs): comms = WorkerComms(MP_CONTEXTS['mp'][DEFAULT_START_METHOD], n_jobs, False) comms.init_comms() comms.insert_poison_pill() for worker_id in range(n_jobs): self.assertEqual(comms.get_task(worker_id), POISON_PILL) comms.task_done(worker_id) comms.join_task_queues() def test_worker_non_lethal_poison_pill(self): """ Test that a non-lethal poison pill is inserted for every worker """ # Shouldn't be the same thing self.assertNotEqual(POISON_PILL, NON_LETHAL_POISON_PILL) for n_jobs in [1, 2, 4]: with self.subTest(n_jobs=n_jobs): comms = WorkerComms(MP_CONTEXTS['mp'][DEFAULT_START_METHOD], n_jobs, False) comms.init_comms() comms.insert_non_lethal_poison_pill() for worker_id in range(n_jobs): self.assertEqual(comms.get_task(worker_id), NON_LETHAL_POISON_PILL) comms.task_done(worker_id) comms.join_task_queues() def test_worker_restart(self): """ Test worker restart related functions """ comms = WorkerComms(MP_CONTEXTS['mp'][DEFAULT_START_METHOD], 5, False) comms.init_comms() # No restarts yet with patch.object(comms._worker_restart_condition, 'wait'): self.assertListEqual(list(comms.get_worker_restarts()), []) # Signal some restarts comms.signal_worker_restart(0) comms.signal_worker_restart(2) comms.signal_worker_restart(3) # Restarts available self.assertListEqual(list(comms.get_worker_restarts()), [0, 2, 3]) # Reset some comms.reset_worker_restart(0) comms.reset_worker_restart(3) # Restarts available self.assertListEqual(list(comms.get_worker_restarts()), [2]) # Reset last one comms.reset_worker_restart(2) with patch.object(comms._worker_restart_condition, 'wait'): self.assertListEqual(list(comms.get_worker_restarts()), []) def test_worker_alive(self): """ Test worker alive related functions """ comms = WorkerComms(MP_CONTEXTS['mp'][DEFAULT_START_METHOD], 5, False) comms.init_comms() # Signal some workers are alive comms.signal_worker_alive(0) comms.signal_worker_alive(1) comms.signal_worker_dead(1) comms.signal_worker_alive(2) comms.signal_worker_alive(3) # Check alive status self.assertListEqual([comms.is_worker_alive(worker_id) for worker_id in range(5)], [True, False, True, True, False]) # Reset some comms.signal_worker_dead(0) comms.signal_worker_dead(3) # Check alive status self.assertListEqual([comms.is_worker_alive(worker_id) for worker_id in range(5)], [False, False, True, False, False]) def test_drain_result_queue_terminate_worker(self): """ get_results should be called once, get_exit_results should be called when exit function is defined """ comms = WorkerComms(MP_CONTEXTS['mp'][DEFAULT_START_METHOD], 5, False) comms.init_comms() dont_wait_event = threading.Event() dont_wait_event.set() with patch.object(comms, 'get_results', side_effect=comms.get_results) as p: comms.drain_results_queue_terminate_worker(dont_wait_event) self.assertEqual(p.call_count, 1) self.assertTrue(dont_wait_event.is_set()) def test_drain_queues(self): """ _drain_and_join_queue should be called for every queue that matters. There are as many tasks queues as workers and 1 results queue """ for n_jobs in [1, 2, 4]: comms = WorkerComms(MP_CONTEXTS['mp'][DEFAULT_START_METHOD], n_jobs, False) comms.init_comms() with self.subTest(n_jobs=n_jobs), patch.object(comms, 'drain_and_join_queue') as p: comms.drain_queues() self.assertEqual(p.call_count, n_jobs + 1) def test__drain_and_join_queue(self): """ Test draining queues """ comms = WorkerComms(MP_CONTEXTS['mp'][DEFAULT_START_METHOD], 2, False) # Create a custom queue with some data q = mp.JoinableQueue() q.put(1) q.put('hello') q.put('world') # Drain queue. It should now be closed comms._drain_and_join_queue(q) # Depending on Python version it either throws OSError or ValueError with self.assertRaises((OSError, ValueError)): q.get(block=False) def test_timeouts(self): """ Tests timeout related function """ comms = WorkerComms(MP_CONTEXTS['mp'][DEFAULT_START_METHOD], 5, False) comms.init_comms() # Signal workers started for worker_id in range(5): with self.subTest(worker_id=worker_id), patch('time.time', side_effect=[1000.0, 2000.0, 3000.0]): self.assertListEqual( comms._workers_time_task_started[worker_id * 3 : worker_id * 3 + 3], [0.0, 0.0, 0.0] ) comms.signal_worker_init_started(worker_id) comms.signal_worker_task_started(worker_id) comms.signal_worker_exit_started(worker_id) self.assertListEqual( comms._workers_time_task_started[worker_id * 3 : worker_id * 3 + 3], [1000.0, 2000.0, 3000.0] ) # Check timeouts for worker_id in range(5): # worker_init, times out at > 10 for timeout, has_timed_out in [(8, True), (9, True), (10, True), (11, False)]: with self.subTest(timeout=timeout, worker_id=worker_id), patch('time.time', return_value=1010.0): self.assertEqual(comms.has_worker_init_timed_out(worker_id, timeout), has_timed_out) # task, times out at > 9 for timeout, has_timed_out in [(8, True), (9, True), (10, False), (11, False)]: with self.subTest(timeout=timeout, worker_id=worker_id), patch('time.time', return_value=2009.0): self.assertEqual(comms.has_worker_task_timed_out(worker_id, timeout), has_timed_out) # worker_exit, times out at > 8 for timeout, has_timed_out in [(8, True), (9, False), (10, False), (11, False)]: with self.subTest(timeout=timeout, worker_id=worker_id), patch('time.time', return_value=3008.0): self.assertEqual(comms.has_worker_exit_timed_out(worker_id, timeout), has_timed_out) # Reset for worker_id in range(5): with self.subTest(worker_id=worker_id): comms.signal_worker_init_completed(worker_id) self.assertListEqual( comms._workers_time_task_started[worker_id * 3 : worker_id * 3 + 3], [0.0, 2000.0, 3000.0] ) comms.signal_worker_task_completed(worker_id) self.assertListEqual( comms._workers_time_task_started[worker_id * 3 : worker_id * 3 + 3], [0.0, 0.0, 3000.0] ) comms.signal_worker_exit_completed(worker_id) self.assertListEqual( comms._workers_time_task_started[worker_id * 3 : worker_id * 3 + 3], [0.0, 0.0, 0.0] ) # Check timeouts. Should be False because nothing started for worker_id in range(5): for timeout in [0, 0.1, 3]: with self.subTest(worker_id=worker_id, timeout=timeout): self.assertFalse(comms.has_worker_init_timed_out(worker_id, timeout))