Method test_pool_context_sync

dali/test/python/test_pool.py:279–308

(self, start_method)

    @params(*start_methods)
    def test_pool_context_sync(self, start_method):
        callbacks = [simple_callback, another_callback]
        groups = [MockGroup.from_callback(cb, prefetch_queue_depth=3) for cb in callbacks]
        with create_pool(
            groups, keep_alive_queue_size=1, num_workers=4, start_method=start_method
        ) as pool:
            capture_processes(pool)
            for i in range(4):
                work_batch = TaskArgs.make_sample(SampleRange(0, 10 * (i + 1), 0, 0))
                pool.schedule_batch(context_i=0, work_batch=work_batch)
                pool.schedule_batch(context_i=1, work_batch=work_batch)
            assert_scheduled_num(pool.contexts[0], 4)
            assert_scheduled_num(pool.contexts[1], 4)
            # pool after a reset should discard all previously scheduled tasks
            # (and sync workers to avoid race on writing to results buffer)
            pool.reset()
            tasks = [(SampleInfo(1000 + j, j, 0, 1),) for j in range(5)]
            work_batch = TaskArgs.make_sample(SampleRange(1000, 1005, 0, 1))
            pool.schedule_batch(context_i=0, work_batch=work_batch)
            pool.schedule_batch(context_i=1, work_batch=work_batch)
            assert_scheduled_num(pool.contexts[0], 1)
            assert_scheduled_num(pool.contexts[1], 1)
            batch_0 = pool.receive_batch(context_i=0)
            batch_1 = pool.receive_batch(context_i=1)
            assert len(batch_0) == len(tasks)
            assert len(batch_1) == len(tasks)
            for task, sample in zip(tasks, batch_0):
                np.testing.assert_array_equal(answer(-1, *task)[1:], sample[1:])
            for task, sample in zip(tasks, batch_1):
                np.testing.assert_array_equal(answer(-1, *task)[1:] + 100, sample[1:])

    @params(1, 4)
    def test_multiple_stateful_sources_single_worker(self, num_workers):