Function run_parallel

dali/test/python/experimental_mode/test_multithreading.py:39–68

(function: Callable[[int], T], num_threads: int | None = None)

def run_parallel(function: Callable[[int], T], num_threads: int | None = None) -> dict[int, T]:
    num_threads = get_num_threads(num_threads)

    barrier = threading.Barrier(num_threads)
    results = {}
    errors = {}

    def wrapper(thread_id: int):
        try:
            barrier.wait()
            results[thread_id] = function(thread_id)
        except Exception as exception:
            errors[thread_id] = exception

    threads = []
    for i in range(num_threads):
        thread = threading.Thread(target=wrapper, args=(i,))
        thread.start()
        threads.append(thread)

    for thread in threads:
        thread.join()

    if errors:
        # Raise only the first error
        thread_id, error = next(iter(errors.items()))
        error.args = (f"Error on thread {thread_id}: {error}", *error.args[1:])
        raise error

    return results


@params("cpu", "gpu")