From 6e957a082b9ecc707a7d9bcd4a113c509042454d Mon Sep 17 00:00:00 2001 From: pinjie Date: Sun, 9 Aug 2026 12:51:10 +0000 Subject: [PATCH] rebuild thread pool Signed-off-by: pinjie --- .gitmodules | 3 + .../cpp_unit_tests.yaml | 6 + .../ext_impl/CMakeLists.txt | 11 + .../ext_impl/external/googletest | 1 + .../ext_impl/src/CMakeLists.txt | 8 - .../inc/PyNvBatchAsyncStreamReader.hpp | 1 + .../inc/PyNvGopDecoder.hpp | 17 +- .../inc/PyNvSampleReader.hpp | 1 + .../PyNvOnDemandDecoder/inc/ThreadPool.hpp | 154 ++++++++++ .../src/PyNvBatchAsyncStreamReader.cpp | 55 +--- .../src/PyNvGopDecoder_common.cpp | 53 +--- .../src/PyNvGopDecoder_constructors.cpp | 80 ++---- .../src/PyNvGopDecoder_random_decoder.cpp | 76 +++-- .../src/PyNvGopDecoder_separate_decoder.cpp | 240 ++++++---------- .../src/PyNvSampleReader.cpp | 69 ++--- .../ext_impl/utest/CMakeLists.txt | 66 +++++ .../ext_impl/utest/thread_pool_test.cpp | 263 ++++++++++++++++++ 17 files changed, 699 insertions(+), 405 deletions(-) create mode 100644 packages/on_demand_video_decoder/cpp_unit_tests.yaml create mode 160000 packages/on_demand_video_decoder/ext_impl/external/googletest create mode 100644 packages/on_demand_video_decoder/ext_impl/utest/CMakeLists.txt create mode 100644 packages/on_demand_video_decoder/ext_impl/utest/thread_pool_test.cpp diff --git a/.gitmodules b/.gitmodules index 21da87d6..9fbddbba 100644 --- a/.gitmodules +++ b/.gitmodules @@ -20,3 +20,6 @@ [submodule "packages/example_skbuild_package/ext_impl/external/googletest"] path = packages/example_skbuild_package/ext_impl/external/googletest url = https://github.com/google/googletest.git +[submodule "packages/on_demand_video_decoder/ext_impl/external/googletest"] + path = packages/on_demand_video_decoder/ext_impl/external/googletest + url = https://github.com/google/googletest.git diff --git a/packages/on_demand_video_decoder/cpp_unit_tests.yaml b/packages/on_demand_video_decoder/cpp_unit_tests.yaml new file mode 100644 index 00000000..fa69c9ce --- /dev/null +++ b/packages/on_demand_video_decoder/cpp_unit_tests.yaml @@ -0,0 +1,6 @@ +# Native C++ unit test target for on_demand_video_decoder. + +cmake_source_dir: ext_impl +cuda_arch_strategy: cmake +test_option: ACCVLAB_ON_DEMAND_VIDEO_DECODER_BUILD_CPP_TESTS +test_target: accvlab_on_demand_video_decoder_run_cpp_tests diff --git a/packages/on_demand_video_decoder/ext_impl/CMakeLists.txt b/packages/on_demand_video_decoder/ext_impl/CMakeLists.txt index 80b9191e..6ccc4fe9 100644 --- a/packages/on_demand_video_decoder/ext_impl/CMakeLists.txt +++ b/packages/on_demand_video_decoder/ext_impl/CMakeLists.txt @@ -93,6 +93,17 @@ endif() add_subdirectory(src) +option( + ACCVLAB_ON_DEMAND_VIDEO_DECODER_BUILD_CPP_TESTS + "Build on_demand_video_decoder native C++ tests" + OFF +) + +if(ACCVLAB_ON_DEMAND_VIDEO_DECODER_BUILD_CPP_TESTS) + enable_testing() + add_subdirectory(utest) +endif() + #If we need to run with santizer, the way is ASAN_OPTIONS=protect_shadow_gap=0:replace_intrin=0:detect_leaks=1 ./cpp_samples/run_decoding #set (CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -fno-omit-frame-pointer -fsanitize=address") #set (CMAKE_LINKER_FLAGS_DEBUG "${CMAKE_LINKER_FLAGS_DEBUG} -fno-omit-frame-pointer -fsanitize=address") diff --git a/packages/on_demand_video_decoder/ext_impl/external/googletest b/packages/on_demand_video_decoder/ext_impl/external/googletest new file mode 160000 index 00000000..52eb8108 --- /dev/null +++ b/packages/on_demand_video_decoder/ext_impl/external/googletest @@ -0,0 +1 @@ +Subproject commit 52eb8108c5bdec04579160ae17225d66034bd723 diff --git a/packages/on_demand_video_decoder/ext_impl/src/CMakeLists.txt b/packages/on_demand_video_decoder/ext_impl/src/CMakeLists.txt index 158e0500..f4306b29 100644 --- a/packages/on_demand_video_decoder/ext_impl/src/CMakeLists.txt +++ b/packages/on_demand_video_decoder/ext_impl/src/CMakeLists.txt @@ -25,14 +25,6 @@ if(WIN32) add_definitions(-D_CRT_SECURE_NO_WARNINGS) endif(WIN32) -option(PROCESS_SYNC "Run the demuxer and decoder in synchornous" OFF) -set(PROCESS_SYNC $ENV{PROCESS_SYNC}) -if (PROCESS_SYNC) - add_compile_definitions(PROCESS_SYNC) -endif() - -message("PROCESS_SYNC environment variable: $ENV{PROCESS_SYNC}") - option(USE_NVTX "enable nvtx support" FALSE) set(USE_NVTX $ENV{USE_NVTX}) diff --git a/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/inc/PyNvBatchAsyncStreamReader.hpp b/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/inc/PyNvBatchAsyncStreamReader.hpp index d8a7ca4b..bd0ead4f 100644 --- a/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/inc/PyNvBatchAsyncStreamReader.hpp +++ b/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/inc/PyNvBatchAsyncStreamReader.hpp @@ -154,6 +154,7 @@ class PyNvBatchAsyncStreamReader { int max_frames_per_decode_call = 0; std::vector VideoReaderMap; + ThreadPool frame_pool; // 2D-specific aggregator pools, one per video slot. Each pool holds the // F frames decoded for that slot in a single Decode() call. Per-video diff --git a/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/inc/PyNvGopDecoder.hpp b/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/inc/PyNvGopDecoder.hpp index 4b734120..73f0d147 100644 --- a/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/inc/PyNvGopDecoder.hpp +++ b/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/inc/PyNvGopDecoder.hpp @@ -84,7 +84,7 @@ class PyNvGopDecoder { ~PyNvGopDecoder(); /** - * Force all thread runners to synchronously terminate their tasks + * Force all thread pools to synchronously terminate their tasks * * This method provides a way to forcefully stop all running threads in case of * exceptions or when immediate termination is needed. It clears all pending tasks @@ -182,8 +182,8 @@ class PyNvGopDecoder { * - Frame data blocks follow the header * * Performance: - * - Files are read in parallel using internal thread pool (merge_runners) - * - Number of threads = min(file_paths.size(), merge_runners.size()) + * - Files are read in parallel using the internal thread pool + * - Worker count is managed by the internal thread pool * - Each file is validated for correct GOP format * * @param file_paths Vector of file paths to GOP binary files @@ -714,16 +714,13 @@ class PyNvGopDecoder { GPUMemoryPool gpu_mem_pool; - // Thread runners for reuse - std::vector demux_runners; - std::vector decode_runners; - std::vector merge_runners; + // Thread pools for reuse + ThreadPool demux_pool; + ThreadPool decode_pool; + ThreadPool parallel_pool; // Lazy loading functions void ensureCudaContextInitialized(); - void ensureDemuxRunnersInitialized(size_t required_count); - void ensureDecodeRunnersInitialized(); - void ensureMergeRunnersInitialized(); /** * Internal implementation for GOP extraction diff --git a/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/inc/PyNvSampleReader.hpp b/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/inc/PyNvSampleReader.hpp index f5fd09b8..388207e7 100644 --- a/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/inc/PyNvSampleReader.hpp +++ b/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/inc/PyNvSampleReader.hpp @@ -174,6 +174,7 @@ class PyNvSampleReader { int num_of_set = 0; std::vector VideoReaderMap; + ThreadPool frame_pool; // Async decode related members ConcurrentQueue decode_result_queue; // Buffer size = 1 diff --git a/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/inc/ThreadPool.hpp b/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/inc/ThreadPool.hpp index 0c478ac8..6cb76e10 100644 --- a/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/inc/ThreadPool.hpp +++ b/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/inc/ThreadPool.hpp @@ -14,6 +14,8 @@ * limitations under the License. */ +#include +#include #include #include #include @@ -21,6 +23,15 @@ #include #include #include +#include +#include +#include +#include +#include + +#if defined(__linux__) +#include +#endif class ThreadRunner { public: @@ -166,3 +177,146 @@ class ThreadRunner { bool hasException; std::exception_ptr exceptionPtr; // Store original exception for rethrow }; + +class ThreadPool { + public: + ThreadPool() : ThreadPool(available_cpu_count()) {} + explicit ThreadPool(size_t max_worker_count) : max_worker_count(std::max(1, max_worker_count)) {} + + static size_t available_cpu_count() { +#if defined(__linux__) + cpu_set_t cpu_set; + CPU_ZERO(&cpu_set); + if (sched_getaffinity(0, sizeof(cpu_set), &cpu_set) == 0) { + const size_t cpu_count = CPU_COUNT(&cpu_set); + if (cpu_count > 0) { + return cpu_count; + } + } +#endif + const size_t cpu_count = std::thread::hardware_concurrency(); + return cpu_count == 0 ? 1 : cpu_count; + } + + ThreadPool(const ThreadPool&) = delete; + ThreadPool& operator=(const ThreadPool&) = delete; + ThreadPool(ThreadPool&&) = delete; + ThreadPool& operator=(ThreadPool&&) = delete; + + template + void submit_indexed(size_t task_count, Func&& task) { + wait_all(); + if (task_count == 0) { + return; + } + + const size_t worker_count = std::min(task_count, max_worker_count); + while (workers.size() < worker_count) { + workers.emplace_back(std::make_unique()); + } + + using Task = typename std::decay::type; + auto task_ptr = std::make_shared(std::forward(task)); + auto next_index = std::make_shared>(0); + active_state = std::make_shared(task_count); + active_worker_count = 0; + + try { + for (; active_worker_count < worker_count; ++active_worker_count) { + workers[active_worker_count]->start( + [task_ptr, next_index, state = active_state, task_count]() { + while (true) { + const size_t index = next_index->fetch_add(1); + if (index >= task_count) { + return; + } + try { + (*task_ptr)(index); + } catch (...) { + state->exceptions[index] = std::current_exception(); + } + } + }); + } + } catch (...) { + active_state->submission_exception = std::current_exception(); + wait_all(); + } + } + + template + void run_indexed(size_t task_count, Func&& task) { + submit_indexed(task_count, std::forward(task)); + wait_all(); + } + + void wait_all() { + std::exception_ptr first_exception; + for (size_t index = 0; index < active_worker_count; ++index) { + try { + workers[index]->join(); + } catch (...) { + if (!first_exception) { + first_exception = std::current_exception(); + } + } + } + + auto state = std::move(active_state); + active_worker_count = 0; + if (state) { + if (!first_exception) { + first_exception = state->submission_exception; + } + for (const auto& exception : state->exceptions) { + if (!first_exception && exception) { + first_exception = exception; + } + } + } + + if (first_exception) { + std::rethrow_exception(first_exception); + } + } + + void force_join() { + for (auto& worker : workers) { + worker->force_join(); + } + active_worker_count = 0; + active_state.reset(); + } + + private: + struct TaskState { + explicit TaskState(size_t task_count) : exceptions(task_count) {} + + std::vector exceptions; + std::exception_ptr submission_exception; + }; + + const size_t max_worker_count; + std::vector> workers; + size_t active_worker_count = 0; + std::shared_ptr active_state; +}; + +inline void wait_all(ThreadPool& first, ThreadPool& second) { + std::exception_ptr first_exception; + try { + first.wait_all(); + } catch (...) { + first_exception = std::current_exception(); + } + try { + second.wait_all(); + } catch (...) { + if (!first_exception) { + first_exception = std::current_exception(); + } + } + if (first_exception) { + std::rethrow_exception(first_exception); + } +} diff --git a/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/src/PyNvBatchAsyncStreamReader.cpp b/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/src/PyNvBatchAsyncStreamReader.cpp index c88bea31..adfbd7c9 100644 --- a/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/src/PyNvBatchAsyncStreamReader.cpp +++ b/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/src/PyNvBatchAsyncStreamReader.cpp @@ -25,7 +25,6 @@ #include #include #include -#include #include #include @@ -36,44 +35,6 @@ namespace py = pybind11; -namespace { -// Parallel per-file fanout, mirroring PyNvSampleReader.cpp's local helper. -// Each filepath/frame_id pair is processed in its own thread; first exception -// captured rethrows after join. -template -std::vector process_frames_in_parallel(const std::vector& filepaths, - const std::vector& frame_ids, - const std::vector& video_readers, - Func process_frame) { - nvtxRangePushA("Process Frames in Parallel (2D worker)"); - std::vector res(filepaths.size()); - std::exception_ptr eptr = nullptr; - std::mutex mutex; - - std::vector threads; - threads.reserve(filepaths.size()); - - for (size_t i = 0; i < filepaths.size(); ++i) { - threads.emplace_back([&, i]() { - try { - res[i] = process_frame(video_readers[i], frame_ids[i]); - } catch (const std::exception&) { - std::lock_guard lock(mutex); - if (!eptr) eptr = std::current_exception(); - } - }); - } - for (auto& t : threads) t.join(); - - if (eptr) { - nvtxRangePop(); - std::rethrow_exception(eptr); - } - nvtxRangePop(); - return res; -} -} // namespace - namespace frame_output = accvlab::on_demand_video_decoder::internal::frame_output; PyNvBatchAsyncStreamReader::PyNvBatchAsyncStreamReader(int num_of_set, int num_of_file, @@ -319,10 +280,18 @@ std::vector PyNvBatchAsyncStreamReader::run_rgb_out_1d(const std::vect } nvtxRangePop(); - return process_frames_in_parallel(filepaths, frame_ids, video_readers, - [as_bgr](PyNvVideoReader* reader, int frame_id) { - return reader->run_single_rgb_out(frame_id, as_bgr); - }); + std::vector result(filepaths.size()); + nvtxRangePushA("Process Frames in Parallel (2D worker)"); + try { + frame_pool.run_indexed(filepaths.size(), [&](size_t index) { + result[index] = video_readers[index]->run_single_rgb_out(frame_ids[index], as_bgr); + }); + } catch (...) { + nvtxRangePop(); + throw; + } + nvtxRangePop(); + return result; } void PyNvBatchAsyncStreamReader::Decode(const std::vector& filepaths, diff --git a/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/src/PyNvGopDecoder_common.cpp b/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/src/PyNvGopDecoder_common.cpp index 687686fc..d6d9f550 100644 --- a/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/src/PyNvGopDecoder_common.cpp +++ b/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/src/PyNvGopDecoder_common.cpp @@ -399,53 +399,25 @@ int PyNvGopDecoder::InitializeDemuxers(const std::vector& filepaths int num_of_files = filepaths.size(); demuxers.resize(num_of_files); -#ifdef PROCESS_SYNC - for (int i = 0; i < num_of_files; ++i) { - nvtxRangePushA((std::string("Demuxer creation : ") + std::to_string(i)).c_str()); - if (fastStreamInfos) { - CreateDemuxer(demuxers[i], filepaths[i], fastStreamInfos + i); - } else { - CreateDemuxer(demuxers[i], filepaths[i], nullptr); - } - if (!demuxers[i]->IsValid()) { - LOG(ERROR) << "create demuxer failed with video files: " << filepaths[i]; - nvtxRangePop(); // Demuxer creation - nvtxRangePop(); // Initialize Demuxers - return -1; - } - nvtxRangePop(); // Demuxer creation - } -#endif - -#ifndef PROCESS_SYNC - for (int i = 0; i < num_of_files; ++i) { - nvtxRangePushA((std::string("Demuxer creation thread start: ") + std::to_string(i)).c_str()); - demux_runners[i].join(); - if (fastStreamInfos) { - demux_runners[i].start(PyNvGopDecoder::CreateDemuxer, std::ref(demuxers[i]), filepaths[i], - fastStreamInfos + i); - } else { - demux_runners[i].start(PyNvGopDecoder::CreateDemuxer, std::ref(demuxers[i]), filepaths[i], - nullptr); + demux_pool.run_indexed(num_of_files, [&](size_t index) { + nvtxRangePushA((std::string("Demuxer creation: ") + std::to_string(index)).c_str()); + try { + const FastStreamInfo* fast_stream_info = fastStreamInfos ? fastStreamInfos + index : nullptr; + CreateDemuxer(demuxers[index], filepaths[index], fast_stream_info); + } catch (...) { + nvtxRangePop(); + throw; } - nvtxRangePop(); // Demuxer creation thread start - } + nvtxRangePop(); + }); for (int i = 0; i < num_of_files; ++i) { - nvtxRangePushA((std::string("Demuxer creation thread join: ") + std::to_string(i)).c_str()); - demux_runners[i].join(); - nvtxRangePop(); if (!demuxers[i]->IsValid()) { - for (int index = i; index < num_of_files; index++) { - demux_runners[index].join(); - } LOG(ERROR) << "create demuxer failed with video files " << filepaths[i]; - nvtxRangePop(); // Demuxer creation thread join nvtxRangePop(); // Initialize Demuxers return -1; } } -#endif // check decoder and demuxer must have the same resolution for (int i = 0; i < num_of_files; ++i) { @@ -498,8 +470,6 @@ int PyNvGopDecoder::InitializeDecoders(const std::vector& codec_ids) { const int num_of_files = static_cast(codec_ids.size()); ensureCudaContextInitialized(); - ensureDecodeRunnersInitialized(); - // Temporarily push context for decoder creation ck(cuCtxPushCurrent(this->cu_context)); @@ -548,8 +518,6 @@ int PyNvGopDecoder::AssignGroupedDecoderSlots(const std::vector 0 ? static_cast(max_num_files) : 1; + const size_t cpu_budget = std::max(1, ThreadPool::available_cpu_count() / budget_divisor); + return std::min(file_limit, cpu_budget); +} + +} // namespace + std::vector GetFastInitInfo(const std::vector& filepaths) { std::vector fast_stream_infos; fast_stream_infos.reserve(filepaths.size()); @@ -139,46 +149,14 @@ void PyNvGopDecoder::ensureCudaContextInitialized() { } } -void PyNvGopDecoder::ensureDemuxRunnersInitialized(size_t required_count) { - if (required_count > static_cast(max_num_files)) { - throw std::invalid_argument("required demux runners exceed max_num_files"); - } - - // max_num_files is a request-capacity limit, not an eager thread count. - demux_runners.reserve(max_num_files); - while (demux_runners.size() < required_count) { - demux_runners.emplace_back(); - } -} - -void PyNvGopDecoder::ensureDecodeRunnersInitialized() { - if (!decode_runners.empty()) { - return; // Already initialized - } - - decode_runners.reserve(max_num_files); - for (size_t i = 0; i < max_num_files; ++i) { - decode_runners.emplace_back(); - } -} - -void PyNvGopDecoder::ensureMergeRunnersInitialized() { - if (!merge_runners.empty()) { - return; // Already initialized - } - - // Initialize merge thread pool with max_num_files threads for parallel file processing - merge_runners.reserve(max_num_files); - for (size_t i = 0; i < max_num_files; ++i) { - merge_runners.emplace_back(); - } -} - PyNvGopDecoder::PyNvGopDecoder(int iMaxFileNum, int iGpu, bool bSuppressNoColorRangeWarning, CUstream external_stream) : max_num_files(iMaxFileNum), gpu_id(iGpu), - suppress_no_color_range_given_warning(bSuppressNoColorRangeWarning) { + suppress_no_color_range_given_warning(bSuppressNoColorRangeWarning), + demux_pool(WorkerCountForCpuBudget(iMaxFileNum, 2)), + decode_pool(WorkerCountForCpuBudget(iMaxFileNum, 2)), + parallel_pool(WorkerCountForCpuBudget(iMaxFileNum, 1)) { #ifdef IS_DEBUG_BUILD std::cout << "New PyNvGopDecoder object" << std::endl; #endif @@ -193,20 +171,9 @@ PyNvGopDecoder::PyNvGopDecoder(int iMaxFileNum, int iGpu, bool bSuppressNoColorR } void PyNvGopDecoder::force_join_all() { - // Force join all demux runners - for (auto& runner : demux_runners) { - runner.force_join(); - } - - // Force join all decode runners - for (auto& runner : decode_runners) { - runner.force_join(); - } - - // Force join all merge runners - for (auto& runner : merge_runners) { - runner.force_join(); - } + demux_pool.force_join(); + decode_pool.force_join(); + parallel_pool.force_join(); } PyNvGopDecoder::~PyNvGopDecoder() { @@ -214,6 +181,8 @@ PyNvGopDecoder::~PyNvGopDecoder() { std::cout << "Delete PyNvGopDecoder object" << std::endl; #endif + force_join_all(); + // Temporarily push context for GPU resource cleanup. // This ensures the destructor works correctly on any thread. if (this->cu_context) { @@ -241,17 +210,6 @@ PyNvGopDecoder::~PyNvGopDecoder() { // No need to pop - we use temporary push/pop pattern instead. ck(cuDevicePrimaryCtxRelease(this->gpu_id)); } - - // Clean up thread runners - for (auto& runner : demux_runners) { - runner.join(); - } - for (auto& runner : decode_runners) { - runner.join(); - } - for (auto& runner : merge_runners) { - runner.join(); - } } void Init_PyNvGopDecoder(py::module& m) { diff --git a/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/src/PyNvGopDecoder_random_decoder.cpp b/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/src/PyNvGopDecoder_random_decoder.cpp index 80e44b27..767fd9db 100644 --- a/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/src/PyNvGopDecoder_random_decoder.cpp +++ b/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/src/PyNvGopDecoder_random_decoder.cpp @@ -56,6 +56,9 @@ void PyNvGopDecoder::decode_from_video(const std::vector& filepaths if (filepaths.size() != frame_ids.size()) { throw std::invalid_argument("[ERROR] filepaths and frame_ids must have the same length"); } + if (filepaths.size() > max_num_files) { + throw std::invalid_argument("[ERROR] filepaths size is greater than max_num_files"); + } nvtxRangePushA("Decode"); const size_t total_frames = frame_ids.size(); @@ -73,8 +76,6 @@ void PyNvGopDecoder::decode_from_video(const std::vector& filepaths // lazy loading ensureCudaContextInitialized(); - ensureDemuxRunnersInitialized(total_frames); - ensureDecodeRunnersInitialized(); // reset last decoded frame infos reset_last_decoded_frame_infos(this->last_decoded_frame_infos); @@ -132,6 +133,8 @@ void PyNvGopDecoder::decode_from_video(const std::vector& filepaths std::vector> all_first_frame_ids; all_gop_lens.resize(total_frames); all_first_frame_ids.resize(total_frames); + std::vector> all_sorted_frame_ids(total_frames); + std::vector process_frame(total_frames, true); nvtxRangePushA("Frame processing"); for (int i = 0; i < total_frames; ++i) { @@ -155,44 +158,41 @@ void PyNvGopDecoder::decode_from_video(const std::vector& filepaths decodedFrames[i].reserve(1); } -#ifdef PROCESS_SYNC - DemuxGopProc(demuxers[i].get(), vpacket_queue[i].get(), sorted_frame_ids, first_frame_ids, - gop_length, vpacket_array[i], false); - if (convert_to_rgb) { - DecProc(demuxers[i]->GetColorRange(), this->vdec[i].get(), rgb_frames[i], - per_file_frame_buffers[i], vpacket_queue[i].get(), sorted_frame_ids, as_bgr, - filepaths[i], this->last_decoded_frame_infos[i]); - } else { - DecProc(demuxers[i]->GetColorRange(), this->vdec[i].get(), decodedFrames[i], - per_file_frame_buffers[i], vpacket_queue[i].get(), sorted_frame_ids, - false, filepaths[i], this->last_decoded_frame_infos[i]); - } -#else - demux_runners[i].join(); - demux_runners[i].start(PyNvGopDecoder::DemuxGopProc, demuxers[i].get(), vpacket_queue[i].get(), - sorted_frame_ids, std::ref(first_frame_ids), std::ref(gop_length), - std::ref(vpacket_array[i]), false); - - if (convert_to_rgb) { - decode_runners[i].join(); - decode_runners[i].start(PyNvGopDecoder::DecProc, demuxers[i]->GetColorRange(), - this->vdec[i].get(), std::ref(rgb_frames[i]), - per_file_frame_buffers[i], vpacket_queue[i].get(), sorted_frame_ids, - as_bgr, filepaths[i], std::ref(this->last_decoded_frame_infos[i])); - } else { - decode_runners[i].join(); - decode_runners[i].start(PyNvGopDecoder::DecProc, - demuxers[i]->GetColorRange(), this->vdec[i].get(), - std::ref(decodedFrames[i]), per_file_frame_buffers[i], - vpacket_queue[i].get(), sorted_frame_ids, false, filepaths[i], - std::ref(this->last_decoded_frame_infos[i])); - } -#endif + all_sorted_frame_ids[i] = std::move(sorted_frame_ids); } catch (const std::exception& e) { + process_frame[i] = false; this->force_join_all(); std::cerr << "[ERROR] " << e.what() << std::endl; } } + try { + demux_pool.submit_indexed(total_frames, [&](size_t index) { + if (!process_frame[index]) { + return; + } + DemuxGopProc(demuxers[index].get(), vpacket_queue[index].get(), all_sorted_frame_ids[index], + all_first_frame_ids[index], all_gop_lens[index], vpacket_array[index], false); + }); + decode_pool.submit_indexed(total_frames, [&](size_t index) { + if (!process_frame[index]) { + return; + } + if (convert_to_rgb) { + DecProc(demuxers[index]->GetColorRange(), this->vdec[index].get(), + rgb_frames[index], per_file_frame_buffers[index], + vpacket_queue[index].get(), all_sorted_frame_ids[index], as_bgr, + filepaths[index], this->last_decoded_frame_infos[index]); + } else { + DecProc(demuxers[index]->GetColorRange(), this->vdec[index].get(), + decodedFrames[index], per_file_frame_buffers[index], + vpacket_queue[index].get(), all_sorted_frame_ids[index], false, + filepaths[index], this->last_decoded_frame_infos[index]); + } + }); + } catch (...) { + this->force_join_all(); + throw; + } nvtxRangePop(); //Frame processing nvtxRangePushA("Demux & decode thread join"); @@ -202,12 +202,8 @@ void PyNvGopDecoder::decode_from_video(const std::vector& filepaths out_if_no_color_conversion->resize(total_frames); } try { + wait_all(demux_pool, decode_pool); for (int i = 0; i < total_frames; ++i) { -#ifndef PROCESS_SYNC - demux_runners[i].join(); - decode_runners[i].join(); -#endif - if (convert_to_rgb) { if (rgb_frames[i].empty()) { this->force_join_all(); diff --git a/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/src/PyNvGopDecoder_separate_decoder.cpp b/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/src/PyNvGopDecoder_separate_decoder.cpp index 361abf69..88787bee 100644 --- a/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/src/PyNvGopDecoder_separate_decoder.cpp +++ b/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/src/PyNvGopDecoder_separate_decoder.cpp @@ -61,9 +61,6 @@ void PyNvGopDecoder::get_gop_internal( vpacket_array.emplace_back(); } - // lazy loading - ensureDemuxRunnersInitialized(total_frames); - // Initialize demuxers st = InitializeDemuxers(filepaths, demuxers, fastStreamInfos); if (st != 0) { @@ -73,6 +70,7 @@ void PyNvGopDecoder::get_gop_internal( // Initialize GOP metadata containers all_gop_lens.resize(total_frames); all_first_frame_ids.resize(total_frames); + std::vector> all_sorted_frame_ids(total_frames); // Extract packets for each video nvtxRangePushA("Packet extraction"); @@ -88,31 +86,29 @@ void PyNvGopDecoder::get_gop_internal( filepaths[i]); } } -#ifdef PROCESS_SYNC - DemuxGopProc(demuxers[i].get(), vpacket_queue[i].get(), sorted_frame_ids, all_first_frame_ids[i], - all_gop_lens[i], vpacket_array[i], true); -#else - demux_runners[i].join(); - demux_runners[i].start(PyNvGopDecoder::DemuxGopProc, demuxers[i].get(), vpacket_queue[i].get(), - sorted_frame_ids, std::ref(all_first_frame_ids[i]), - std::ref(all_gop_lens[i]), std::ref(vpacket_array[i]), true); -#endif + all_sorted_frame_ids[i] = std::move(sorted_frame_ids); } catch (const std::exception& e) { this->force_join_all(); LOG(ERROR) << "Packet extraction failed: " << e.what(); throw std::runtime_error(e.what()); } } + try { + demux_pool.submit_indexed(total_frames, [&](size_t index) { + DemuxGopProc(demuxers[index].get(), vpacket_queue[index].get(), all_sorted_frame_ids[index], + all_first_frame_ids[index], all_gop_lens[index], vpacket_array[index], true); + }); + } catch (const std::exception& e) { + this->force_join_all(); + LOG(ERROR) << "Packet extraction failed: " << e.what(); + throw std::runtime_error(e.what()); + } nvtxRangePop(); // Packet extraction // Wait for all demux threads to complete nvtxRangePushA("Demux thread join"); try { - for (int i = 0; i < total_frames; ++i) { -#ifndef PROCESS_SYNC - demux_runners[i].join(); -#endif - } + demux_pool.wait_all(); } catch (const std::exception& e) { this->force_join_all(); throw std::runtime_error(e.what()); @@ -321,8 +317,6 @@ void PyNvGopDecoder::decode_from_packet_list(std::vector> packe std::vector dummp; ensureCudaContextInitialized(); - // ensureDemuxRunnersInitialized(); - ensureDecodeRunnersInitialized(); st = InitGpuMemPool(heights, widths, dummp, true); if (st != 0) { @@ -759,7 +753,6 @@ int PyNvGopDecoder::main_decode_groups( std::vector>>>& vpacket_queue, std::vector& output) { ensureCudaContextInitialized(); - ensureDecodeRunnersInitialized(); const size_t num_groups = frame_id_groups.size(); if (decoder_configs.size() != num_groups || decoder_slots.size() != num_groups) { @@ -810,23 +803,9 @@ int PyNvGopDecoder::main_decode_groups( for (size_t group_idx = 0; group_idx < num_groups; ++group_idx) { try { - const size_t slot_idx = decoder_slots[group_idx]; const AVColorRange color_range = static_cast(color_ranges[group_idx]); WarnIfColorRangeUnspecified(color_range, source_names[group_idx]); rgb_frames[group_idx].reserve(frame_id_groups[group_idx].size()); -#ifdef PROCESS_SYNC - DecProc(color_range, vdec[slot_idx].get(), rgb_frames[group_idx], - group_frame_buffers[group_idx], vpacket_queue[group_idx].get(), - frame_id_groups[group_idx], as_bgr, source_names[group_idx], - last_decoded_frame_infos[slot_idx]); -#else - decode_runners[slot_idx].join(); - decode_runners[slot_idx].start(PyNvGopDecoder::DecProc, color_range, - vdec[slot_idx].get(), std::ref(rgb_frames[group_idx]), - group_frame_buffers[group_idx], vpacket_queue[group_idx].get(), - frame_id_groups[group_idx], as_bgr, source_names[group_idx], - std::ref(last_decoded_frame_infos[slot_idx])); -#endif } catch (const std::exception& error) { force_join_all(); LOG(ERROR) << "Grouped DecProc failed: " << error.what(); @@ -834,17 +813,20 @@ int PyNvGopDecoder::main_decode_groups( } } -#ifndef PROCESS_SYNC try { - for (size_t group_idx = 0; group_idx < num_groups; ++group_idx) { - decode_runners[decoder_slots[group_idx]].join(); - } + decode_pool.run_indexed(num_groups, [&](size_t group_idx) { + const size_t slot_idx = decoder_slots[group_idx]; + const AVColorRange color_range = static_cast(color_ranges[group_idx]); + DecProc(color_range, vdec[slot_idx].get(), rgb_frames[group_idx], + group_frame_buffers[group_idx], vpacket_queue[group_idx].get(), + frame_id_groups[group_idx], as_bgr, source_names[group_idx], + last_decoded_frame_infos[slot_idx]); + }); } catch (const std::exception& error) { force_join_all(); - LOG(ERROR) << "Grouped decode thread join failed: " << error.what(); + LOG(ERROR) << "Grouped DecProc failed: " << error.what(); return -1; } -#endif output.clear(); output.reserve(total_output_frames); @@ -861,7 +843,7 @@ int PyNvGopDecoder::main_decode_groups( } } - // DecProc queues color conversion on cu_stream. Joining its CPU runner only + // DecProc queues color conversion on cu_stream. Waiting for its CPU worker only // guarantees that the work was submitted; the output frames are ready when // this public synchronous API returns only after the stream is synchronized. CUDA_DRVAPI_CALL(cuStreamSynchronize(this->cu_stream)); @@ -880,7 +862,6 @@ int PyNvGopDecoder::main_decode( // lazy loading ensureCudaContextInitialized(); - ensureDecodeRunnersInitialized(); st = InitGpuMemPool(heights, widths, frame_sizes, convert_to_rgb); if (st != 0) { @@ -923,34 +904,6 @@ int PyNvGopDecoder::main_decode( } else { decodedFrames[i].reserve(sorted_frame_ids.size()); } - auto& packet_queue = vpacket_queue[i]; - auto& frame_buffers = per_file_frame_buffers[i]; - AVColorRange color_range = static_cast(color_ranges[i]); -#ifdef PROCESS_SYNC - if (convert_to_rgb) { - DecProc(color_range, this->vdec[i].get(), rgb_frames[i], frame_buffers, - packet_queue.get(), sorted_frame_ids, as_bgr, filepaths[i], - this->last_decoded_frame_infos[i]); - } else { - DecProc(color_range, this->vdec[i].get(), decodedFrames[i], frame_buffers, - packet_queue.get(), sorted_frame_ids, false, filepaths[i], - this->last_decoded_frame_infos[i]); - } -#else - if (convert_to_rgb) { - decode_runners[i].join(); - decode_runners[i].start(PyNvGopDecoder::DecProc, color_range, this->vdec[i].get(), - std::ref(rgb_frames[i]), frame_buffers, packet_queue.get(), - sorted_frame_ids, as_bgr, filepaths[i], - std::ref(this->last_decoded_frame_infos[i])); - } else { - decode_runners[i].join(); - decode_runners[i].start(PyNvGopDecoder::DecProc, color_range, - this->vdec[i].get(), std::ref(decodedFrames[i]), frame_buffers, - packet_queue.get(), sorted_frame_ids, false, filepaths[i], - std::ref(this->last_decoded_frame_infos[i])); - } -#endif } catch (const std::exception& e) { this->force_join_all(); LOG(ERROR) << "DecProc failed: " << e.what(); @@ -958,18 +911,26 @@ int PyNvGopDecoder::main_decode( } } -#ifndef PROCESS_SYNC - // Join all runners first; exceptions trigger force_join_all() before propagating. - for (int j = 0; j < total_frames; ++j) { - try { - decode_runners[j].join(); - } catch (const std::exception& e) { - this->force_join_all(); - LOG(ERROR) << "decode_runners[" << j << "].join() failed: " << e.what(); - return -1; - } + try { + decode_pool.run_indexed(total_frames, [&](size_t index) { + std::vector sorted_frame_ids = {frame_ids[index]}; + AVColorRange color_range = static_cast(color_ranges[index]); + if (convert_to_rgb) { + DecProc(color_range, this->vdec[index].get(), rgb_frames[index], + per_file_frame_buffers[index], vpacket_queue[index].get(), sorted_frame_ids, + as_bgr, filepaths[index], this->last_decoded_frame_infos[index]); + } else { + DecProc(color_range, this->vdec[index].get(), decodedFrames[index], + per_file_frame_buffers[index], vpacket_queue[index].get(), + sorted_frame_ids, false, filepaths[index], + this->last_decoded_frame_infos[index]); + } + }); + } catch (const std::exception& e) { + this->force_join_all(); + LOG(ERROR) << "DecProc failed: " << e.what(); + return -1; } -#endif for (int i = 0; i < total_frames; ++i) { if (convert_to_rgb) { @@ -1475,90 +1436,59 @@ void PyNvGopDecoder::LoadGOPFromFiles(const std::vector& file_paths throw std::invalid_argument("[ERROR] file_paths is empty"); } - // Ensure merge thread pool is initialized - ensureMergeRunnersInitialized(); - - // Calculate number of threads to use - size_t num_threads = std::min(file_paths.size(), merge_runners.size()); + // Read all binary files in parallel + file_data_buffers.resize(file_paths.size()); - // Helper lambda for parallel execution with exception handling - auto executeInParallel = [&](const std::function& task_func) { - std::vector exceptions(num_threads); + parallel_pool.run_indexed(file_paths.size(), [&](size_t file_idx) { + const auto& file_path = file_paths[file_idx]; - // Start parallel tasks - for (size_t i = 0; i < num_threads; ++i) { - merge_runners[i].start([&, i]() { - try { - task_func(i); - } catch (...) { - exceptions[i] = std::current_exception(); - } - }); + // Check if file exists + if (!std::filesystem::exists(file_path)) { + throw std::runtime_error("[ERROR] File does not exist: " + file_path); } - // Wait for all tasks to complete - for (size_t i = 0; i < num_threads; ++i) { - merge_runners[i].join(); + // Read entire file into memory + std::ifstream file(file_path, std::ios::binary | std::ios::ate); + if (!file.is_open()) { + throw std::runtime_error("[ERROR] Failed to open file: " + file_path); } - // Check for exceptions - for (auto& ex : exceptions) { - if (ex) { - std::rethrow_exception(ex); - } + const std::streampos end_position = file.tellg(); + if (end_position < 0) { + throw std::runtime_error("[ERROR] Failed to determine file size: " + file_path); } - }; - - // Read all binary files in parallel - file_data_buffers.resize(file_paths.size()); - - executeInParallel([&](size_t thread_id) { - // Process files assigned to this thread - for (size_t file_idx = thread_id; file_idx < file_paths.size(); file_idx += num_threads) { - const auto& file_path = file_paths[file_idx]; - - // Check if file exists - if (!std::filesystem::exists(file_path)) { - throw std::runtime_error("[ERROR] File does not exist: " + file_path); - } - - // Read entire file into memory - std::ifstream file(file_path, std::ios::binary | std::ios::ate); - if (!file.is_open()) { - throw std::runtime_error("[ERROR] Failed to open file: " + file_path); - } - - size_t file_size = file.tellg(); - file.seekg(0, std::ios::beg); - - std::vector file_buffer(file_size); - file.read(reinterpret_cast(file_buffer.data()), file_size); - if (file.fail()) { - throw std::runtime_error("[ERROR] Failed to read file: " + file_path); - } - file.close(); - - // Validate file header - if (file_size < sizeof(uint32_t)) { - throw std::invalid_argument("[ERROR] File too small: " + file_path); - } + const size_t file_size = static_cast(end_position); + file.seekg(0, std::ios::beg); - const uint8_t* data_ptr = file_buffer.data(); - uint32_t frame_count = *reinterpret_cast(data_ptr); - - if (frame_count == 0) { - throw std::invalid_argument("[ERROR] File contains no frames: " + file_path); - } - - // Validate header size - size_t expected_header_size = sizeof(uint32_t) + frame_count * sizeof(size_t); - if (file_size < expected_header_size) { - throw std::invalid_argument("[ERROR] File header invalid: " + file_path); - } + std::vector file_buffer(file_size); + file.read(reinterpret_cast(file_buffer.data()), file_size); + if (file.fail()) { + throw std::runtime_error("[ERROR] Failed to read file: " + file_path); + } + file.close(); - // Store file data - file_data_buffers[file_idx] = std::move(file_buffer); + // Validate the complete bundle before exposing it to Python. + std::vector color_ranges; + std::vector codec_ids; + std::vector widths; + std::vector heights; + std::vector frame_sizes; + std::vector gop_lens; + std::vector first_frame_ids; + std::vector> packets_bytes; + std::vector> decode_idxs; + std::vector packet_binary_data_ptrs; + std::vector packet_binary_data_sizes; + const uint32_t frame_count = + parseSerializedPacketData(file_buffer.data(), file_buffer.size(), color_ranges, codec_ids, widths, + heights, frame_sizes, gop_lens, first_frame_ids, packets_bytes, + decode_idxs, packet_binary_data_ptrs, packet_binary_data_sizes); + if (frame_count == 0) { + throw std::invalid_argument("[ERROR] File contains no frames: " + file_path); } + + // Store file data + file_data_buffers[file_idx] = std::move(file_buffer); }); nvtxRangePop(); diff --git a/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/src/PyNvSampleReader.cpp b/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/src/PyNvSampleReader.cpp index 2ea2ec56..d75b477e 100644 --- a/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/src/PyNvSampleReader.cpp +++ b/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/src/PyNvSampleReader.cpp @@ -21,7 +21,6 @@ #include #include #include -#include #include #include #include @@ -126,43 +125,6 @@ void PyNvSampleReader::clearAllReaders() { } } -// Helper function to process video frames in parallel -template -std::vector process_frames_in_parallel(const std::vector& filepaths, - const std::vector& frame_ids, - const std::vector& video_readers, - Func process_frame) { - nvtxRangePushA("Process Frames in Parallel"); - std::vector res(filepaths.size()); - std::exception_ptr eptr = nullptr; - std::mutex mutex; - - std::vector threads; - threads.reserve(filepaths.size()); - - for (int i = 0; i < filepaths.size(); i++) { - threads.emplace_back([&, i]() { - try { - res[i] = process_frame(video_readers[i], frame_ids[i]); - } catch (const std::exception& e) { - std::lock_guard lock(mutex); - eptr = std::current_exception(); - } - }); - } - - for (auto& thread : threads) { - thread.join(); - } - - if (eptr) { - nvtxRangePop(); - std::rethrow_exception(eptr); - } - nvtxRangePop(); - return res; -} - std::vector PyNvSampleReader::run_rgb_out(const std::vector& filepaths, const std::vector frame_ids, bool as_bgr) { // NOTE: Do NOT call waitForPendingAsyncTask() here! @@ -203,10 +165,18 @@ std::vector PyNvSampleReader::run_rgb_out(const std::vector(filepaths, frame_ids, video_readers, - [as_bgr](PyNvVideoReader* reader, int frame_id) { - return reader->run_single_rgb_out(frame_id, as_bgr); - }); + std::vector result(filepaths.size()); + nvtxRangePushA("Process Frames in Parallel"); + try { + frame_pool.run_indexed(filepaths.size(), [&](size_t index) { + result[index] = video_readers[index]->run_single_rgb_out(frame_ids[index], as_bgr); + }); + } catch (...) { + nvtxRangePop(); + throw; + } + nvtxRangePop(); + return result; } std::vector PyNvSampleReader::run(const std::vector& filepaths, @@ -246,9 +216,18 @@ std::vector PyNvSampleReader::run(const std::vector( - filepaths, frame_ids, video_readers, - [](PyNvVideoReader* reader, int frame_id) { return reader->run_single(frame_id); }); + std::vector result(filepaths.size()); + nvtxRangePushA("Process Frames in Parallel"); + try { + frame_pool.run_indexed(filepaths.size(), [&](size_t index) { + result[index] = video_readers[index]->run_single(frame_ids[index]); + }); + } catch (...) { + nvtxRangePop(); + throw; + } + nvtxRangePop(); + return result; } void Init_PyNvSampleReader(py::module& m) { diff --git a/packages/on_demand_video_decoder/ext_impl/utest/CMakeLists.txt b/packages/on_demand_video_decoder/ext_impl/utest/CMakeLists.txt new file mode 100644 index 00000000..8434ff04 --- /dev/null +++ b/packages/on_demand_video_decoder/ext_impl/utest/CMakeLists.txt @@ -0,0 +1,66 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +set(ACCVLAB_ON_DEMAND_VIDEO_DECODER_GTEST_SOURCE_DIR + "${CMAKE_CURRENT_LIST_DIR}/../external/googletest" +) + +if(NOT EXISTS "${ACCVLAB_ON_DEMAND_VIDEO_DECODER_GTEST_SOURCE_DIR}/CMakeLists.txt") + message(FATAL_ERROR "GoogleTest submodule not found. Run: git submodule update --init --recursive") +endif() + +add_subdirectory( + "${ACCVLAB_ON_DEMAND_VIDEO_DECODER_GTEST_SOURCE_DIR}" + "${CMAKE_CURRENT_BINARY_DIR}/googletest" + EXCLUDE_FROM_ALL +) + +find_package(Threads REQUIRED) + +add_executable(on_demand_video_decoder_thread_pool_test EXCLUDE_FROM_ALL + thread_pool_test.cpp +) + +target_include_directories(on_demand_video_decoder_thread_pool_test PRIVATE + "${CMAKE_CURRENT_LIST_DIR}/../src/PyNvOnDemandDecoder/inc" +) + +target_link_libraries(on_demand_video_decoder_thread_pool_test PRIVATE + GTest::gtest_main + Threads::Threads +) + +add_test( + NAME on_demand_video_decoder_thread_pool_test + COMMAND on_demand_video_decoder_thread_pool_test +) +set_tests_properties(on_demand_video_decoder_thread_pool_test PROPERTIES + LABELS "on_demand_video_decoder_native" +) + +set(ACCVLAB_ON_DEMAND_VIDEO_DECODER_CTEST_CONFIG_ARGS) +if(CMAKE_CONFIGURATION_TYPES) + set(ACCVLAB_ON_DEMAND_VIDEO_DECODER_CTEST_CONFIG_ARGS -C $) +endif() + +add_custom_target(accvlab_on_demand_video_decoder_run_cpp_tests + COMMAND ${CMAKE_CTEST_COMMAND} + --output-on-failure + ${ACCVLAB_ON_DEMAND_VIDEO_DECODER_CTEST_CONFIG_ARGS} + -L on_demand_video_decoder_native + DEPENDS on_demand_video_decoder_thread_pool_test + WORKING_DIRECTORY ${CMAKE_BINARY_DIR} + COMMENT "Running on_demand_video_decoder native tests" + COMMAND_EXPAND_LISTS +) diff --git a/packages/on_demand_video_decoder/ext_impl/utest/thread_pool_test.cpp b/packages/on_demand_video_decoder/ext_impl/utest/thread_pool_test.cpp new file mode 100644 index 00000000..466a3e3f --- /dev/null +++ b/packages/on_demand_video_decoder/ext_impl/utest/thread_pool_test.cpp @@ -0,0 +1,263 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "ThreadPool.hpp" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if defined(__linux__) +#include +#endif + +namespace { + +class ThreadBarrier { + public: + explicit ThreadBarrier(size_t participant_count) : participant_count_(participant_count) {} + + void arrive_and_wait() { + std::unique_lock lock(mutex_); + ++arrived_count_; + if (arrived_count_ == participant_count_) { + condition_.notify_all(); + return; + } + if (!condition_.wait_for(lock, std::chrono::seconds(5), + [this]() { return arrived_count_ == participant_count_; })) { + throw std::runtime_error("Timed out waiting for thread-pool workers"); + } + } + + private: + const size_t participant_count_; + size_t arrived_count_ = 0; + std::mutex mutex_; + std::condition_variable condition_; +}; + +void update_max(std::atomic& maximum, size_t value) { + size_t current = maximum.load(); + while (current < value && !maximum.compare_exchange_weak(current, value)) { + } +} + +TEST(ThreadPoolTest, ZeroTasksDoNotInvokeCallable) { + ThreadPool pool(4); + std::atomic call_count{0}; + + pool.run_indexed(0, [&call_count](size_t) { ++call_count; }); + + EXPECT_EQ(call_count.load(), 0U); +} + +TEST(ThreadPoolTest, RunsEveryIndexExactlyOnce) { + ThreadPool pool(4); + std::vector> visit_counts(257); + for (auto& visit_count : visit_counts) { + visit_count.store(0); + } + + pool.run_indexed(visit_counts.size(), [&visit_counts](size_t index) { ++visit_counts[index]; }); + + for (size_t index = 0; index < visit_counts.size(); ++index) { + EXPECT_EQ(visit_counts[index].load(), 1U) << "index " << index; + } +} + +TEST(ThreadPoolTest, DoesNotExceedWorkerLimit) { + constexpr size_t worker_count = 3; + ThreadPool pool(worker_count); + ThreadBarrier barrier(worker_count); + std::atomic active_count{0}; + std::atomic maximum_active_count{0}; + + pool.run_indexed(12, [&](size_t index) { + const size_t active = ++active_count; + update_max(maximum_active_count, active); + if (index < worker_count) { + barrier.arrive_and_wait(); + } + --active_count; + }); + + EXPECT_EQ(maximum_active_count.load(), worker_count); +} + +TEST(ThreadPoolTest, ExpandsForLargerBatches) { + constexpr size_t worker_count = 4; + ThreadPool pool(worker_count); + std::atomic first_batch_count{0}; + + pool.run_indexed(1, [&first_batch_count](size_t) { ++first_batch_count; }); + + ThreadBarrier barrier(worker_count); + std::atomic active_count{0}; + std::atomic maximum_active_count{0}; + pool.run_indexed(worker_count, [&](size_t) { + const size_t active = ++active_count; + update_max(maximum_active_count, active); + barrier.arrive_and_wait(); + --active_count; + }); + + EXPECT_EQ(first_batch_count.load(), 1U); + EXPECT_EQ(maximum_active_count.load(), worker_count); +} + +TEST(ThreadPoolTest, SubmitIndexedReturnsBeforeBlockedTasksFinish) { + ThreadPool pool(3); + std::promise release_promise; + std::shared_future release = release_promise.get_future().share(); + std::atomic completed_count{0}; + + pool.submit_indexed(15, [&release, &completed_count](size_t) { + release.wait(); + ++completed_count; + }); + + EXPECT_EQ(completed_count.load(), 0U); + release_promise.set_value(); + pool.wait_all(); + EXPECT_EQ(completed_count.load(), 15U); +} + +TEST(ThreadPoolTest, WaitsForAllTasksBeforeRethrowingEarliestIndexedException) { + ThreadPool pool(4); + std::vector> visit_counts(32); + for (auto& visit_count : visit_counts) { + visit_count.store(0); + } + + try { + pool.run_indexed(visit_counts.size(), [&visit_counts](size_t index) { + ++visit_counts[index]; + if (index == 2) { + throw std::logic_error("exception at index 2"); + } + if (index == 7) { + throw std::runtime_error("exception at index 7"); + } + }); + FAIL() << "run_indexed did not rethrow a task exception"; + } catch (const std::logic_error& error) { + EXPECT_EQ(std::string(error.what()), "exception at index 2"); + } catch (...) { + FAIL() << "run_indexed did not preserve the earliest exception type"; + } + + for (size_t index = 0; index < visit_counts.size(); ++index) { + EXPECT_EQ(visit_counts[index].load(), 1U) << "index " << index; + } +} + +TEST(ThreadPoolTest, CanBeReusedAfterTaskException) { + ThreadPool pool(2); + EXPECT_THROW(pool.run_indexed(1, [](size_t) { throw std::runtime_error("expected failure"); }), + std::runtime_error); + + std::atomic completed_count{0}; + EXPECT_NO_THROW(pool.run_indexed(11, [&completed_count](size_t) { ++completed_count; })); + EXPECT_EQ(completed_count.load(), 11U); +} + +TEST(ThreadPoolTest, SupportsMoveOnlyCallable) { + ThreadPool pool(2); + std::atomic completed_count{0}; + auto marker = std::make_unique(17); + + pool.run_indexed(9, [marker = std::move(marker), &completed_count](size_t) { + if (*marker != 17) { + throw std::runtime_error("move-only callable state was not preserved"); + } + ++completed_count; + }); + + EXPECT_EQ(completed_count.load(), 9U); +} + +TEST(ThreadPoolTest, DualPoolWaitCompletesBothPoolsBeforeRethrowing) { + ThreadPool first(2); + ThreadPool second(2); + std::atomic second_pool_completed_count{0}; + + first.submit_indexed(1, [](size_t) { throw std::runtime_error("first pool failure"); }); + second.submit_indexed(23, [&second_pool_completed_count](size_t) { ++second_pool_completed_count; }); + + try { + wait_all(first, second); + FAIL() << "wait_all did not rethrow the first pool exception"; + } catch (const std::runtime_error& error) { + EXPECT_EQ(std::string(error.what()), "first pool failure"); + } catch (...) { + FAIL() << "wait_all did not preserve the first pool exception type"; + } + EXPECT_EQ(second_pool_completed_count.load(), 23U); +} + +#if defined(__linux__) +class AffinityGuard { + public: + explicit AffinityGuard(const cpu_set_t& affinity) : affinity_(affinity) {} + + ~AffinityGuard() { sched_setaffinity(0, sizeof(affinity_), &affinity_); } + + AffinityGuard(const AffinityGuard&) = delete; + AffinityGuard& operator=(const AffinityGuard&) = delete; + + private: + cpu_set_t affinity_; +}; + +TEST(ThreadPoolTest, AvailableCpuCountRespectsProcessAffinity) { + cpu_set_t original_affinity; + CPU_ZERO(&original_affinity); + if (sched_getaffinity(0, sizeof(original_affinity), &original_affinity) != 0) { + GTEST_SKIP() << "sched_getaffinity is not available"; + } + AffinityGuard affinity_guard(original_affinity); + + int selected_cpu = -1; + for (int cpu = 0; cpu < CPU_SETSIZE; ++cpu) { + if (CPU_ISSET(cpu, &original_affinity)) { + selected_cpu = cpu; + break; + } + } + ASSERT_GE(selected_cpu, 0); + + cpu_set_t single_cpu_affinity; + CPU_ZERO(&single_cpu_affinity); + CPU_SET(selected_cpu, &single_cpu_affinity); + if (sched_setaffinity(0, sizeof(single_cpu_affinity), &single_cpu_affinity) != 0) { + GTEST_SKIP() << "sched_setaffinity is not permitted"; + } + + EXPECT_EQ(ThreadPool::available_cpu_count(), 1U); +} +#endif + +} // namespace