| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import threading |
| 5 | from typing import Any, ClassVar |
| 6 | |
| 7 | import numpy.typing as npt |
| 8 | |
| 9 | from vllm.logger import init_logger |
| 10 | |
| 11 | from .base import VideoSourceMetadata, VideoTargetMetadata, check_frame_pixel_limit |
| 12 | |
| 13 | logger = init_logger(__name__) |
| 14 | |
| 15 | |
| 16 | def decode_deepstream( |
| 17 | loader_cls, |
| 18 | data: bytes, |
| 19 | target: VideoTargetMetadata, |
| 20 | sampling_kwargs: dict, |
| 21 | *, |
| 22 | pool_size: int | None = None, |
| 23 | timeout_sec: float = 120.0, |
| 24 | ) -> tuple[npt.NDArray, VideoSourceMetadata, list[int], list[int]]: |
| 25 | from nvidia.deepstream_videodecode import probe_metadata |
| 26 | |
| 27 | total_frames, original_fps, duration, _width, _height, codec = probe_metadata(data) |
| 28 | check_frame_pixel_limit(_width, _height) |
| 29 | source = loader_cls._prepare_source( |
| 30 | VideoSourceMetadata(total_frames, original_fps, duration) |
| 31 | ) |
| 32 | frame_idx = loader_cls.compute_frames_index_to_sample( |
| 33 | source=source, target=target, **sampling_kwargs |
| 34 | ) |
| 35 | frames, valid = DeepStreamVideoBackendMixin.decode_indices( |
| 36 | data, |
| 37 | frame_idx, |
| 38 | source, |
| 39 | codec=codec, |
| 40 | pool_size=pool_size, |
| 41 | timeout_sec=timeout_sec, |
| 42 | ) |
| 43 | return frames, source, frame_idx, valid |
| 44 | |
| 45 | |
| 46 | class DeepStreamVideoBackendMixin: |
| 47 | """NVIDIA DeepStream (NVDEC) GPU-decode codec utilities. |
| 48 | |
| 49 | Decoding runs on a shared pool of daemon threads inside one CUDA |
| 50 | context (see the ``nvidia-deepstream-videodecode-cu13`` package). The |
| 51 | container bytes are pushed into an ``appsrc`` GStreamer pipeline, so no |
| 52 | local file path is required — HTTP and base64 sources decode identically |
| 53 | to local files. |
| 54 | |
| 55 | Like the OpenCV mixin, this provides only the codec layer. |
| 56 | Frame *selection* lives in the loader's |
| 57 | ``compute_frames_index_to_sample`` and arrives here as an explicit |
| 58 | list of frame indices. |
| 59 | """ |
| 60 | |
| 61 | # Process-wide lazy decode pool, shared across all DeepStream backends. |
| 62 | _pool: ClassVar[Any] = None |
| 63 | _pool_lock: ClassVar[Any] = None |
| 64 | |
| 65 | @classmethod |
| 66 | def _get_pool(cls, pool_size: int | None = None): |
| 67 | """Lazy-initialize the shared decode pool on first use. |
| 68 | |
| 69 | ``pool_size`` (number of decode worker threads) comes from |
| 70 | ``--media-io-kwargs`` (``{"video": {"pool_size": N}}``); when unset it |
| 71 | defaults to the existing ``VLLM_MEDIA_LOADING_THREAD_COUNT`` so no |
| 72 | DeepStream-specific env var is needed. The pool is a process-wide |
| 73 | singleton, so the first decode's value wins. |
| 74 | """ |
| 75 | if cls._pool is not None: |
| 76 | return cls._pool |
| 77 | if cls._pool_lock is None: |
| 78 | cls._pool_lock = threading.Lock() |
| 79 | with cls._pool_lock: |
| 80 | if cls._pool is not None: |
| 81 | return cls._pool |
| 82 | import os |
| 83 | |
| 84 | from nvidia.deepstream_videodecode import DecodePool |
| 85 | |
| 86 | if pool_size is None: |
| 87 | pool_size = int(os.environ.get("VLLM_MEDIA_LOADING_THREAD_COUNT", 8)) |
| 88 | pool_size = max(1, min(int(pool_size), 16)) |
| 89 | logger.info( |
| 90 | "[DeepStream] initializing decode pool with %d workers", |
| 91 | pool_size, |
| 92 | ) |
| 93 | cls._pool = DecodePool(num_workers=pool_size) |
| 94 | return cls._pool |
| 95 | |
| 96 | @classmethod |
| 97 | def decode_indices( |
| 98 | cls, |
| 99 | data: bytes, |
| 100 | frame_indices: list[int], |
| 101 | source: VideoSourceMetadata, |
| 102 | codec: str = "", |
| 103 | pool_size: int | None = None, |
| 104 | timeout_sec: float = 120.0, |
| 105 | ) -> tuple[npt.NDArray, list[int]]: |
| 106 | """Decode the requested frame indices from raw container bytes. |
| 107 | |
| 108 | The whole stream is decoded; the pool keeps exactly the frames whose |
| 109 | decode-order index is in ``frame_indices`` (1:1, frame-exact) and |
| 110 | sends EOS once the last one is matched. |
| 111 | |
| 112 | ``codec`` (e.g. ``"h264"``/``"hevc"``) lets the pool keep its NVDEC |
| 113 | session warm across same-codec streams and rebuild only on a codec |
| 114 | change. Frames are returned as a CPU NHWC uint8 array so the |
| 115 | upstream multimodal parser sees the same shape as the other |
| 116 | backends. |
| 117 | """ |
| 118 | if not frame_indices: |
| 119 | raise ValueError("DeepStream backend received no frame indices") |
| 120 | |
| 121 | result = cls._get_pool(pool_size).decode( |
| 122 | data, |
| 123 | target_indices=frame_indices, |
| 124 | codec=codec, |
| 125 | max_frames=len(frame_indices), |
| 126 | timeout_sec=timeout_sec, |
| 127 | ) |
| 128 | if result.error: |
| 129 | raise ValueError(f"DeepStream decode failed: {result.error}") |
| 130 | if result.frames is None or result.n_kept == 0: |
| 131 | raise ValueError("DeepStream decode produced no frames") |
| 132 | |
| 133 | valid = frame_indices[: result.n_kept] |
| 134 | # GPU -> CPU NHWC uint8 at the codec boundary (one PCIe copy); keeps |
| 135 | # the array shape identical to the OpenCV backend. Copy into |
| 136 | # PINNED host memory (reused across calls by PyTorch's pinned caching |
| 137 | # allocator) so the D2H runs at full PCIe bandwidth (~13 GB/s) rather |
| 138 | # than the ~1 GB/s pageable path that plain ``.cpu()`` takes — ~12x |
| 139 | # faster for a 1080p x8 frame batch (~46ms -> ~4ms). ``numpy()`` keeps |
| 140 | # the pinned tensor alive via the array's base. |
| 141 | import torch |
| 142 | |
| 143 | gpu = result.frames |
| 144 | if gpu.is_cuda: |
| 145 | host = torch.empty(gpu.shape, dtype=gpu.dtype, pin_memory=True) |
| 146 | host.copy_(gpu, non_blocking=True) |
| 147 | torch.cuda.current_stream().synchronize() |
| 148 | arr = host.numpy() |
| 149 | else: |
| 150 | arr = gpu.numpy() |
| 151 | return arr, valid |