Add profiling script for GPU bottleneck analysis in EasyCom extraction. Introduce command-line options for GPU ID, log directory, and sampling period. Implement stage profiling to measure extraction performance and synchronize CUDA operations. Enhance CausalSourceAV class with new parameters for decoder threads and open readers, optimizing video decoding process.

2026-08-18 10:30:19 +08:00
parent 24cb92f1f4
commit d610305355
3 changed files with 638 additions and 8 deletions
@@ -230,6 +230,8 @@ def _decode_audio_exact(path: Path, duration: float) -> np.ndarray:
            "-nostdin",
            "-v",
            "error",
+            "-threads",
+            "1",
            "-i",
            str(path),
            "-map",
@@ -275,6 +277,8 @@ class CausalSourceAV:
        frame_cache_size: int = 8,
        decoder_backend: str = "decord",
        decoder_gpu_id: int = 0,
+        decoder_threads: int = 1,
+        max_open_readers: int = 2,
    ) -> None:
        if not specs:
            raise ValueError("At least one source clip is required")
@@ -288,10 +292,15 @@ class CausalSourceAV:
            raise ValueError(f"Unsupported decoder backend: {decoder_backend}")
        if decoder_gpu_id < 0:
            raise ValueError("decoder_gpu_id must be non-negative")
+        if decoder_threads <= 0:
+            raise ValueError("decoder_threads must be positive")
+        if max_open_readers <= 0:
+            raise ValueError("max_open_readers must be positive")
        if shutil.which("ffmpeg") is None:
            raise RuntimeError("ffmpeg is required to decode source audio")

        decord = None
+        self._decord = None
        self._nvc = None
        if decoder_backend == "decord":
            try:
@@ -300,6 +309,7 @@ class CausalSourceAV:
                raise RuntimeError(
                    "The decord backend requires qwen-omni-utils[decord]"
                ) from exc
+            self._decord = decord
        else:
            try:
                import PyNvVideoCodec as nvc
@@ -315,6 +325,11 @@ class CausalSourceAV:
        self.frame_cache_size = int(frame_cache_size)
        self.decoder_backend = decoder_backend
        self.decoder_gpu_id = int(decoder_gpu_id)
+        self.decoder_threads = int(decoder_threads)
+        self.max_open_readers = int(max_open_readers)
+        self._decord_readers: OrderedDict[int, Any] = OrderedDict()
+        self._decord_reader_opens = 0
+        self._decord_reader_evictions = 0
        self._nvc_decoder = None
        self._nvc_source_path: Path | None = None
        self._nvc_decoder_reconfigurations = 0
@@ -340,17 +355,23 @@ class CausalSourceAV:
            )
            if decoder_backend == "decord":
                assert decord is not None
-                reader = decord.VideoReader(str(path))
-                if not len(reader):
+                index_reader = decord.VideoReader(
+                    str(path), num_threads=self.decoder_threads
+                )
+                if not len(index_reader):
                    raise RuntimeError(f"No video frames decoded from {path}")
-                starts, ends = _timestamp_array(reader)
-                source_fps = float(reader.get_avg_fps())
-                first_shape = tuple(int(value) for value in reader[0].shape)
+                starts, ends = _timestamp_array(index_reader)
+                source_fps = float(index_reader.get_avg_fps())
+                first_shape = tuple(int(value) for value in index_reader[0].shape)
                if len(first_shape) != 3:
                    raise RuntimeError(
                        f"Unexpected decoded frame shape {first_shape}: {path}"
                    )
                height, width = first_shape[:2]
+                # VideoReader.__del__ releases its native handle. Do not retain
+                # one reader/thread pool for every one-minute EasyCom clip.
+                del index_reader
+                reader = None
            else:
                reader = None
                starts, ends, source_fps, width, height = _ffprobe_video_index(path)
@@ -487,9 +508,8 @@ class CausalSourceAV:

        clip = self.clips[clip_index]
        if self.decoder_backend == "decord":
-            if clip.reader is None:
-                raise RuntimeError(f"Missing decord reader for {clip.path}")
-            array = clip.reader[frame_index].asnumpy()
+            reader = self._decord_reader_for_clip(clip_index)
+            array = reader[frame_index].asnumpy()
        else:
            decoder = self._pynv_decoder_for_clip(clip)
            decoded_frame = decoder[frame_index]
@@ -516,6 +536,35 @@ class CausalSourceAV:
                self._frame_cache.popitem(last=False)
        return image.copy()

+    def _decord_reader_for_clip(self, clip_index: int) -> Any:
+        """Return an LRU-managed CPU reader instead of retaining every clip."""
+
+        reader = self._decord_readers.get(clip_index)
+        if reader is not None:
+            self._decord_readers.move_to_end(clip_index)
+            return reader
+        if self._decord is None:
+            raise RuntimeError("Decord was not initialized")
+
+        clip = self.clips[clip_index]
+        reader = self._decord.VideoReader(
+            str(clip.path), num_threads=self.decoder_threads
+        )
+        if len(reader) != clip.num_frames:
+            raise RuntimeError(
+                f"Decord frame count changed for {clip.path}: "
+                f"runtime={len(reader)}, indexed={clip.num_frames}"
+            )
+        self._decord_readers[clip_index] = reader
+        self._decord_readers.move_to_end(clip_index)
+        self._decord_reader_opens += 1
+
+        while len(self._decord_readers) > self.max_open_readers:
+            _, evicted_reader = self._decord_readers.popitem(last=False)
+            del evicted_reader
+            self._decord_reader_evictions += 1
+        return reader
+
    def _pynv_decoder_for_clip(self, clip: IndexedSourceClip) -> Any:
        """Return one reusable NVDEC decoder configured for ``clip``."""

@@ -613,6 +662,12 @@ class CausalSourceAV:
            "decoder_gpu_id": (
                self.decoder_gpu_id if self.decoder_backend == "pynvcodec" else None
            ),
+            "decoder_threads": (
+                self.decoder_threads if self.decoder_backend == "decord" else None
+            ),
+            "max_open_readers": (
+                self.max_open_readers if self.decoder_backend == "decord" else None
+            ),
            "decoder_version": (
                getattr(self._nvc, "__version__", None)
                if self.decoder_backend == "pynvcodec"
@@ -650,6 +705,8 @@ class CausalSourceAV:
                "short_interval_fallbacks": self._short_interval_fallbacks,
                "tail_frame_reuses": self._tail_frame_reuses,
                "decoder_reconfigurations": self._nvc_decoder_reconfigurations,
+                "decord_reader_opens": self._decord_reader_opens,
+                "decord_reader_evictions": self._decord_reader_evictions,
            },
        }

@@ -788,6 +845,18 @@ def parse_args() -> argparse.Namespace:
        default=0,
        help="GPU visible to this process used by PyNvVideoCodec",
    )
+    parser.add_argument(
+        "--decoder-threads",
+        type=int,
+        default=1,
+        help="CPU decode threads per Decord VideoReader",
+    )
+    parser.add_argument(
+        "--max-open-readers",
+        type=int,
+        default=2,
+        help="Maximum simultaneously open Decord VideoReaders per process",
+    )
    parser.add_argument("--verify-every", type=int, default=0)
    parser.add_argument("--verify-continuous-every", type=int, default=0)
    parser.add_argument("--no-progress", action="store_true")
@@ -923,6 +992,8 @@ def main() -> None:
        frame_cache_size=args.frame_cache_size,
        decoder_backend=args.decoder,
        decoder_gpu_id=args.decoder_gpu_id,
+        decoder_threads=args.decoder_threads,
+        max_open_readers=args.max_open_readers,
    )
    is_easycom = (
        args.easycom_session_dir is not None