Pinjie/rebuild pool

.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

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

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")

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})
 

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<FixedSizeVideoReaderMap> 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

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<ThreadRunner> demux_runners;
-    std::vector<ThreadRunner> decode_runners;
-    std::vector<ThreadRunner> 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

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<FixedSizeVideoReaderMap> VideoReaderMap;
+    ThreadPool frame_pool;
 
     // Async decode related members
     ConcurrentQueue<DecodeResult> decode_result_queue;  // Buffer size = 1

packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/inc/ThreadPool.hpp

@@ -14,13 +14,24 @@
  * limitations under the License.
  */
 
+#include <algorithm>
+#include <atomic>
 #include <thread>
 #include <mutex>
 #include <condition_variable>
 #include <queue>
 #include <functional>
 #include <exception>
 #include <iostream>
+#include <memory>
+#include <stdexcept>
+#include <type_traits>
+#include <utility>
+#include <vector>
+
+#if defined(__linux__)
+#include <sched.h>
+#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<size_t>(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 <typename Func>
+    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<ThreadRunner>());
+        }
+
+        using Task = typename std::decay<Func>::type;
+        auto task_ptr = std::make_shared<Task>(std::forward<Func>(task));
+        auto next_index = std::make_shared<std::atomic<size_t>>(0);
+        active_state = std::make_shared<TaskState>(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 <typename Func>
+    void run_indexed(size_t task_count, Func&& task) {
+        submit_indexed(task_count, std::forward<Func>(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<std::exception_ptr> exceptions;
+        std::exception_ptr submission_exception;
+    };
+
+    const size_t max_worker_count;
+    std::vector<std::unique_ptr<ThreadRunner>> workers;
+    size_t active_worker_count = 0;
+    std::shared_ptr<TaskState> 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);
+    }
+}

packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/src/PyNvBatchAsyncStreamReader.cpp

@@ -25,7 +25,6 @@
 #include <sstream>
 #include <stdexcept>
 #include <string>
-#include <thread>
 #include <tuple>
 #include <vector>
 
@@ -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 <typename T, typename Func>
-std::vector<T> process_frames_in_parallel(const std::vector<std::string>& filepaths,
-                                          const std::vector<int>& frame_ids,
-                                          const std::vector<PyNvVideoReader*>& video_readers,
-                                          Func process_frame) {
-    nvtxRangePushA("Process Frames in Parallel (2D worker)");
-    std::vector<T> res(filepaths.size());
-    std::exception_ptr eptr = nullptr;
-    std::mutex mutex;
-
-    std::vector<std::thread> 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<std::mutex> 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<RGBFrame> PyNvBatchAsyncStreamReader::run_rgb_out_1d(const std::vect
     }
     nvtxRangePop();
 
-    return process_frames_in_parallel<RGBFrame>(filepaths, frame_ids, video_readers,
-                                                [as_bgr](PyNvVideoReader* reader, int frame_id) {
-                                                    return reader->run_single_rgb_out(frame_id, as_bgr);
-                                                });
+    std::vector<RGBFrame> 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<std::string>& filepaths,