Method test_pool_no_overwrite_batch

dali/test/python/test_pool.py:157–179

(self, start_method)

    # Test that we can safely hold as many results as the keep_alive_queue_size
    @params(*start_methods)
    def test_pool_no_overwrite_batch(self, start_method):
        groups = [MockGroup.from_callback(simple_callback, prefetch_queue_depth=0)]
        for depth in [1, 2, 4, 8]:
            with create_pool(
                groups, keep_alive_queue_size=depth, num_workers=1, start_method=start_method
            ) as pool:
                pids = get_pids(pool)
                pid = pids[0]
                work_batches = [
                    TaskArgs.make_sample(SampleRange(i, i + 1, i, 0)) for i in range(depth)
                ]
                task_list = [[(SampleInfo(i, 0, i, 0),)] for i in range(depth)]
                for i, work_batch in enumerate(work_batches):
                    pool.schedule_batch(context_i=0, work_batch=work_batch)
                assert_scheduled_num(pool.contexts[0], depth)
                batches = []
                for i in range(depth):
                    batches.append(pool.receive_batch(context_i=0))
                    assert_scheduled_num(pool.contexts[0], depth - 1 - i)
                tasks_batches = zip(task_list, batches)
                for tasks, batch in tasks_batches:
                    for task, sample in zip(tasks, batch):
                        np.testing.assert_array_equal(answer(pid, *task), sample)


# ################################################################################################ #