[Multimodal] Reorganize video decoder backends (#49155)
Signed-off-by: Isotr0py <Isotr0py@outlook.com>
This commit is contained in:
@@ -21,8 +21,10 @@ from vllm.multimodal.media import ImageMediaIO, VideoMediaIO
|
||||
from vllm.multimodal.video import (
|
||||
PYNVVIDEOCODEC_VIDEO_BACKEND,
|
||||
VIDEO_LOADER_REGISTRY,
|
||||
PyNvVideoCodecVideoBackend,
|
||||
VideoLoader,
|
||||
)
|
||||
from vllm.multimodal.video_decoders.pynvvideocodec import (
|
||||
PyNvVideoCodecVideoBackendMixin,
|
||||
_pynvvc_frames_to_nhwc,
|
||||
)
|
||||
|
||||
@@ -381,13 +383,15 @@ def test_pynvvideocodec_unrelated_error_propagates(
|
||||
raise original_error
|
||||
|
||||
monkeypatch.setattr(
|
||||
PyNvVideoCodecVideoBackend,
|
||||
PyNvVideoCodecVideoBackendMixin,
|
||||
"_read_source_metadata",
|
||||
classmethod(raise_unrelated_error),
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError) as exc_info:
|
||||
PyNvVideoCodecVideoBackend.decode_frames_pynvvideocodec(b"video", None)
|
||||
PyNvVideoCodecVideoBackendMixin.decode_frames_pynvvideocodec(
|
||||
None, b"video", None
|
||||
)
|
||||
|
||||
assert exc_info.value is original_error
|
||||
|
||||
|
||||
@@ -6,7 +6,6 @@ import time
|
||||
|
||||
import pytest
|
||||
|
||||
import vllm.config.multimodal as multimodal_config_module
|
||||
from vllm.config.multimodal import MultiModalConfig
|
||||
from vllm.multimodal.gpu_ipc_memory import (
|
||||
MultiModalGPUMemoryPool,
|
||||
@@ -15,10 +14,10 @@ from vllm.multimodal.gpu_ipc_memory import (
|
||||
reserve_mm_ipc_gpu_memory,
|
||||
set_mm_gpu_ipc_pool,
|
||||
)
|
||||
from vllm.multimodal.video import (
|
||||
from vllm.multimodal.video_decoders import PYNVVIDEOCODEC_VIDEO_BACKEND
|
||||
from vllm.multimodal.video_decoders.pynvvideocodec import (
|
||||
PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES,
|
||||
PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES,
|
||||
PYNVVIDEOCODEC_VIDEO_BACKEND,
|
||||
)
|
||||
from vllm.utils.mem_constants import GiB_bytes
|
||||
|
||||
@@ -186,11 +185,7 @@ def test_reserve_mm_ipc_gpu_memory_raw_frame_budget_only(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
video_backend: str | None,
|
||||
):
|
||||
monkeypatch.setattr(
|
||||
multimodal_config_module.envs,
|
||||
"VLLM_VIDEO_LOADER_BACKEND",
|
||||
"opencv",
|
||||
)
|
||||
monkeypatch.setenv("VLLM_VIDEO_LOADER_BACKEND", "opencv")
|
||||
mm_config = _mm_config(
|
||||
mm_ipc_gpu_memory_gb=0.25,
|
||||
video_backend=video_backend,
|
||||
@@ -202,11 +197,7 @@ def test_reserve_mm_ipc_gpu_memory_raw_frame_budget_only(
|
||||
def test_reserve_mm_ipc_gpu_memory_includes_pynvvideocodec_decoder_budget(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
monkeypatch.setattr(
|
||||
multimodal_config_module.envs,
|
||||
"VLLM_VIDEO_LOADER_BACKEND",
|
||||
"opencv",
|
||||
)
|
||||
monkeypatch.setenv("VLLM_VIDEO_LOADER_BACKEND", "opencv")
|
||||
mm_config = _mm_config(
|
||||
mm_ipc_gpu_memory_gb=0.25,
|
||||
video_backend=PYNVVIDEOCODEC_VIDEO_BACKEND,
|
||||
@@ -221,11 +212,7 @@ def test_reserve_mm_ipc_gpu_memory_includes_pynvvideocodec_decoder_budget(
|
||||
def test_reserve_mm_ipc_gpu_memory_uses_env_video_backend(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
monkeypatch.setattr(
|
||||
multimodal_config_module.envs,
|
||||
"VLLM_VIDEO_LOADER_BACKEND",
|
||||
PYNVVIDEOCODEC_VIDEO_BACKEND,
|
||||
)
|
||||
monkeypatch.setenv("VLLM_VIDEO_LOADER_BACKEND", PYNVVIDEOCODEC_VIDEO_BACKEND)
|
||||
available_bytes = 4 * GiB_bytes
|
||||
|
||||
assert reserve_mm_ipc_gpu_memory(available_bytes, _mm_config()) == (
|
||||
@@ -236,11 +223,7 @@ def test_reserve_mm_ipc_gpu_memory_uses_env_video_backend(
|
||||
def test_reserve_mm_ipc_gpu_memory_scales_decoder_budget_by_api_servers(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
monkeypatch.setattr(
|
||||
multimodal_config_module.envs,
|
||||
"VLLM_VIDEO_LOADER_BACKEND",
|
||||
PYNVVIDEOCODEC_VIDEO_BACKEND,
|
||||
)
|
||||
monkeypatch.setenv("VLLM_VIDEO_LOADER_BACKEND", PYNVVIDEOCODEC_VIDEO_BACKEND)
|
||||
available_bytes = 8 * GiB_bytes
|
||||
|
||||
assert reserve_mm_ipc_gpu_memory(
|
||||
@@ -253,11 +236,7 @@ def test_reserve_mm_ipc_gpu_memory_scales_decoder_budget_by_api_servers(
|
||||
def test_reserve_mm_ipc_gpu_memory_uses_configured_hw_decoders(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
monkeypatch.setattr(
|
||||
multimodal_config_module.envs,
|
||||
"VLLM_VIDEO_LOADER_BACKEND",
|
||||
"opencv",
|
||||
)
|
||||
monkeypatch.setenv("VLLM_VIDEO_LOADER_BACKEND", "opencv")
|
||||
available_bytes = 4 * GiB_bytes
|
||||
mm_config = _mm_config(
|
||||
video_backend=PYNVVIDEOCODEC_VIDEO_BACKEND,
|
||||
|
||||
+155
-18
@@ -2,6 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import itertools
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
from contextlib import ExitStack, contextmanager
|
||||
@@ -16,24 +17,26 @@ from transformers.video_utils import VideoMetadata
|
||||
from vllm.assets.base import get_vllm_public_assets
|
||||
from vllm.models.minimax_m3.common.mm_preprocess import MiniMaxM3VideoBackend
|
||||
from vllm.multimodal.video import (
|
||||
PYNVVIDEOCODEC_DECODER_CACHE_SIZE,
|
||||
PYNVVIDEOCODEC_VIDEO_BACKEND,
|
||||
VIDEO_LOADER_REGISTRY,
|
||||
DynamicVideoBackend,
|
||||
GLM46VVideoBackend,
|
||||
Molmo2VideoBackend,
|
||||
PyNvVideoCodecDecoderSlot,
|
||||
PyNvVideoCodecVideoBackend,
|
||||
PyNvVideoCodecVideoBackendMixin,
|
||||
Qwen2VLVideoBackend,
|
||||
Qwen3VLVideoBackend,
|
||||
VideoBackend,
|
||||
VideoLoader,
|
||||
VideoSourceMetadata,
|
||||
VideoTargetMetadata,
|
||||
_pynv_decoder_pool,
|
||||
get_video_loader_backend_for_processor,
|
||||
)
|
||||
from vllm.multimodal.video_decoders import decode_video, resolve_video_backend_kwargs
|
||||
from vllm.multimodal.video_decoders.pynvvideocodec import (
|
||||
PYNVVIDEOCODEC_DECODER_CACHE_SIZE,
|
||||
PyNvVideoCodecDecoderSlot,
|
||||
PyNvVideoCodecVideoBackendMixin,
|
||||
_pynv_decoder_pool,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.transformers_utils.processor import get_video_processor_cls_name_from_config
|
||||
|
||||
@@ -99,6 +102,126 @@ def test_video_loader_type_doesnt_exist():
|
||||
VIDEO_LOADER_REGISTRY.load("non_existing_video_loader")
|
||||
|
||||
|
||||
def test_video_decoder_backends_are_lazy_imported():
|
||||
code = """
|
||||
import sys
|
||||
import vllm.multimodal.video # noqa: F401
|
||||
|
||||
backend_modules = {
|
||||
f"vllm.multimodal.video_decoders.{backend}"
|
||||
for backend in ("opencv", "pyav", "torchcodec", "pynvvideocodec", "deepstream")
|
||||
}
|
||||
loaded = sorted(backend_modules & sys.modules.keys())
|
||||
assert not loaded, loaded
|
||||
"""
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", code],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
assert result.returncode == 0, result.stderr
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"backend",
|
||||
["opencv", "pyav", "torchcodec", "pynvvideocodec", "deepstream"],
|
||||
)
|
||||
def test_decode_video_imports_only_selected_backend(
|
||||
backend: str,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
imports = []
|
||||
decoded = object()
|
||||
|
||||
def fake_decoder(*args, **kwargs):
|
||||
return decoded
|
||||
|
||||
class FakeBackendModule:
|
||||
pass
|
||||
|
||||
setattr(FakeBackendModule, f"decode_{backend}", fake_decoder)
|
||||
|
||||
def fake_import_module(name: str, package: str):
|
||||
imports.append((name, package))
|
||||
return FakeBackendModule
|
||||
|
||||
monkeypatch.setattr(
|
||||
"vllm.multimodal.video_decoders.import_module", fake_import_module
|
||||
)
|
||||
result = decode_video(
|
||||
backend,
|
||||
loader_cls=None,
|
||||
data=b"",
|
||||
target=VideoTargetMetadata(-1, -1, 300),
|
||||
sampling_kwargs={},
|
||||
backend_kwargs={},
|
||||
frame_recovery=False,
|
||||
)
|
||||
|
||||
assert result is decoded
|
||||
assert imports == [(f".{backend}", "vllm.multimodal.video_decoders")]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("backend", "kwargs", "expected_sampling", "expected_backend"),
|
||||
[
|
||||
(
|
||||
"torchcodec",
|
||||
{"min_frames": 4, "num_ffmpeg_threads": 2, "seek_mode": "approximate"},
|
||||
{"min_frames": 4},
|
||||
{"num_ffmpeg_threads": 2, "seek_mode": "approximate"},
|
||||
),
|
||||
(
|
||||
"deepstream",
|
||||
{"max_frames": 16, "pool_size": 3, "timeout_sec": 10.0},
|
||||
{"max_frames": 16},
|
||||
{"pool_size": 3, "timeout_sec": 10.0},
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_video_backend_kwargs_are_separated_from_sampling_kwargs(
|
||||
backend: str,
|
||||
kwargs: dict,
|
||||
expected_sampling: dict,
|
||||
expected_backend: dict,
|
||||
):
|
||||
original_kwargs = dict(kwargs)
|
||||
sampling_kwargs, backend_kwargs = resolve_video_backend_kwargs(backend, kwargs)
|
||||
|
||||
assert sampling_kwargs == expected_sampling
|
||||
assert backend_kwargs == expected_backend
|
||||
assert kwargs == original_kwargs
|
||||
|
||||
|
||||
def test_video_backend_rejects_options_for_another_decoder():
|
||||
with pytest.raises(
|
||||
ValueError, match="num_ffmpeg_threads is not supported by the 'pyav' backend"
|
||||
):
|
||||
resolve_video_backend_kwargs("pyav", {"num_ffmpeg_threads": 2})
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("backend", "error"),
|
||||
[
|
||||
("pyav", AssertionError),
|
||||
(PYNVVIDEOCODEC_VIDEO_BACKEND, ValueError),
|
||||
],
|
||||
)
|
||||
def test_video_decoder_spec_validates_frame_recovery(
|
||||
backend: str, error: type[Exception]
|
||||
):
|
||||
with pytest.raises(error, match="frame_recovery is not supported"):
|
||||
decode_video(
|
||||
backend,
|
||||
loader_cls=None,
|
||||
data=b"",
|
||||
target=VideoTargetMetadata(-1, -1, 300),
|
||||
sampling_kwargs={},
|
||||
backend_kwargs={},
|
||||
frame_recovery=True,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA")
|
||||
def test_pynvvideocodec_backend_accounts_raw_decoded_frames(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
@@ -145,7 +268,9 @@ def test_pynvvideocodec_backend_accounts_raw_decoded_frames(
|
||||
"vllm.multimodal.gpu_ipc_memory.get_mm_gpu_ipc_pool", lambda: pool
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
PyNvVideoCodecVideoBackend, "_decode_to_pinned_host", classmethod(fake_decode)
|
||||
PyNvVideoCodecVideoBackendMixin,
|
||||
"_decode_to_pinned_host",
|
||||
classmethod(fake_decode),
|
||||
)
|
||||
|
||||
loader = VIDEO_LOADER_REGISTRY.load(PYNVVIDEOCODEC_VIDEO_BACKEND)
|
||||
@@ -205,7 +330,9 @@ def test_pynvvideocodec_codec_uses_dynamic_sampling_strategy(
|
||||
"vllm.multimodal.gpu_ipc_memory.get_mm_gpu_ipc_pool", lambda: pool
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
DynamicVideoBackend, "_decode_to_pinned_host", classmethod(fake_decode)
|
||||
PyNvVideoCodecVideoBackendMixin,
|
||||
"_decode_to_pinned_host",
|
||||
classmethod(fake_decode),
|
||||
)
|
||||
|
||||
loader = VIDEO_LOADER_REGISTRY.load("opencv_dynamic")
|
||||
@@ -273,7 +400,7 @@ def test_pynvvideocodec_decoder_slots_are_bounded(
|
||||
|
||||
create_count = 0
|
||||
with _fresh_decoder_pool():
|
||||
PyNvVideoCodecVideoBackend._configure_decoder_slots(hw_decoders)
|
||||
PyNvVideoCodecVideoBackendMixin._configure_decoder_slots(hw_decoders)
|
||||
|
||||
def fake_create_slot(cls):
|
||||
nonlocal create_count
|
||||
@@ -281,7 +408,7 @@ def test_pynvvideocodec_decoder_slots_are_bounded(
|
||||
return FakeSlot()
|
||||
|
||||
monkeypatch.setattr(
|
||||
PyNvVideoCodecVideoBackend,
|
||||
PyNvVideoCodecVideoBackendMixin,
|
||||
"_create_decoder_slot",
|
||||
classmethod(fake_create_slot),
|
||||
)
|
||||
@@ -291,12 +418,16 @@ def test_pynvvideocodec_decoder_slots_are_bounded(
|
||||
|
||||
with ExitStack() as stack:
|
||||
retained_slots = [
|
||||
stack.enter_context(PyNvVideoCodecVideoBackend._borrow_decoder_slot())
|
||||
stack.enter_context(
|
||||
PyNvVideoCodecVideoBackendMixin._borrow_decoder_slot()
|
||||
)
|
||||
for _ in range(hw_decoders)
|
||||
]
|
||||
|
||||
def borrow_extra_slot():
|
||||
with PyNvVideoCodecVideoBackend._borrow_decoder_slot() as extra_slot:
|
||||
with (
|
||||
PyNvVideoCodecVideoBackendMixin._borrow_decoder_slot()
|
||||
) as extra_slot:
|
||||
seen_slots.append(extra_slot)
|
||||
borrowed.set()
|
||||
|
||||
@@ -317,11 +448,11 @@ def test_pynvvideocodec_decoder_slots_are_configured_once(
|
||||
):
|
||||
monkeypatch.setattr(_pynv_decoder_pool, "max_slots", None)
|
||||
|
||||
PyNvVideoCodecVideoBackend._configure_decoder_slots(2)
|
||||
PyNvVideoCodecVideoBackend._configure_decoder_slots(2)
|
||||
PyNvVideoCodecVideoBackendMixin._configure_decoder_slots(2)
|
||||
PyNvVideoCodecVideoBackendMixin._configure_decoder_slots(2)
|
||||
|
||||
with pytest.raises(RuntimeError, match="already configured as 2, got 3"):
|
||||
PyNvVideoCodecVideoBackend._configure_decoder_slots(3)
|
||||
PyNvVideoCodecVideoBackendMixin._configure_decoder_slots(3)
|
||||
|
||||
|
||||
def test_pynvvideocodec_failed_rebuild_invalidates_decoder_slot():
|
||||
@@ -367,7 +498,7 @@ def test_pynvvideocodec_failed_rebuild_invalidates_decoder_slot():
|
||||
|
||||
with (
|
||||
pytest.raises(RuntimeError, match="construct failed"),
|
||||
PyNvVideoCodecVideoBackend._borrow_decoder_slot() as borrowed,
|
||||
PyNvVideoCodecVideoBackendMixin._borrow_decoder_slot() as borrowed,
|
||||
):
|
||||
assert borrowed is slot
|
||||
borrowed.get_decoder(
|
||||
@@ -456,6 +587,12 @@ def test_pynvvideocodec_cross_subclass_shares_single_pool():
|
||||
counters via ClassVar shadowing.
|
||||
"""
|
||||
|
||||
class MixinSubclassA(PyNvVideoCodecVideoBackendMixin):
|
||||
pass
|
||||
|
||||
class MixinSubclassB(PyNvVideoCodecVideoBackendMixin):
|
||||
pass
|
||||
|
||||
class FakeSlot:
|
||||
pass
|
||||
|
||||
@@ -475,8 +612,8 @@ def test_pynvvideocodec_cross_subclass_shares_single_pool():
|
||||
)
|
||||
try:
|
||||
with ExitStack() as stack:
|
||||
stack.enter_context(VideoBackend._borrow_decoder_slot())
|
||||
stack.enter_context(Qwen3VLVideoBackend._borrow_decoder_slot())
|
||||
stack.enter_context(MixinSubclassA._borrow_decoder_slot())
|
||||
stack.enter_context(MixinSubclassB._borrow_decoder_slot())
|
||||
assert pool.active == 2
|
||||
|
||||
blocked = threading.Event()
|
||||
@@ -484,7 +621,7 @@ def test_pynvvideocodec_cross_subclass_shares_single_pool():
|
||||
|
||||
def try_borrow():
|
||||
blocked.set()
|
||||
with Qwen2VLVideoBackend._borrow_decoder_slot():
|
||||
with MixinSubclassB._borrow_decoder_slot():
|
||||
acquired.set()
|
||||
|
||||
t = threading.Thread(target=try_borrow)
|
||||
|
||||
@@ -188,11 +188,13 @@ def reserve_mm_ipc_gpu_memory(
|
||||
return available_kv_cache_memory_bytes
|
||||
|
||||
from vllm import envs
|
||||
from vllm.multimodal.video import (
|
||||
PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES,
|
||||
PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES,
|
||||
from vllm.multimodal.video_decoders import (
|
||||
PYNVVIDEOCODEC_DEFAULT_HW_DECODERS,
|
||||
PYNVVIDEOCODEC_VIDEO_BACKEND,
|
||||
)
|
||||
from vllm.multimodal.video_decoders.pynvvideocodec import (
|
||||
PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES,
|
||||
PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES,
|
||||
validate_pynvvideocodec_hw_decoders,
|
||||
)
|
||||
|
||||
|
||||
+66
-1045
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,110 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from importlib import import_module
|
||||
from typing import Any, Literal
|
||||
|
||||
from .base import (
|
||||
PYNVVIDEOCODEC_DEFAULT_HW_DECODERS,
|
||||
PYNVVIDEOCODEC_VIDEO_BACKEND,
|
||||
VideoSourceMetadata,
|
||||
VideoTargetMetadata,
|
||||
check_frame_pixel_limit,
|
||||
)
|
||||
|
||||
VideoDecoderBackend = Literal[
|
||||
"opencv", "pyav", "torchcodec", "pynvvideocodec", "deepstream"
|
||||
]
|
||||
|
||||
_BACKEND_OPTION_DEFAULTS: dict[str, dict[str, Any]] = {
|
||||
"opencv": {},
|
||||
"pyav": {},
|
||||
"torchcodec": {
|
||||
"num_ffmpeg_threads": 0,
|
||||
"seek_mode": "exact",
|
||||
},
|
||||
PYNVVIDEOCODEC_VIDEO_BACKEND: {
|
||||
"hw_decoders": PYNVVIDEOCODEC_DEFAULT_HW_DECODERS,
|
||||
},
|
||||
"deepstream": {
|
||||
"pool_size": None,
|
||||
"timeout_sec": 120.0,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _get_backend_option_defaults(backend: str) -> dict[str, Any]:
|
||||
try:
|
||||
return _BACKEND_OPTION_DEFAULTS[backend]
|
||||
except KeyError:
|
||||
valid = ", ".join(repr(name) for name in _BACKEND_OPTION_DEFAULTS)
|
||||
raise ValueError(
|
||||
f"Unknown video codec backend {backend!r}; valid options: {valid}."
|
||||
) from None
|
||||
|
||||
|
||||
def resolve_video_backend_kwargs(
|
||||
backend: str,
|
||||
kwargs: dict[str, Any],
|
||||
) -> tuple[dict[str, Any], dict[str, Any]]:
|
||||
"""Split frame-sampling kwargs from options owned by a decoder backend."""
|
||||
defaults = _get_backend_option_defaults(backend)
|
||||
sampling_kwargs = dict(kwargs)
|
||||
backend_kwargs = dict(defaults)
|
||||
backend_option_names = {
|
||||
name for options in _BACKEND_OPTION_DEFAULTS.values() for name in options
|
||||
}
|
||||
misplaced = (sampling_kwargs.keys() & backend_option_names) - defaults.keys()
|
||||
if misplaced:
|
||||
names = ", ".join(sorted(misplaced))
|
||||
raise ValueError(f"{names} is not supported by the {backend!r} backend")
|
||||
|
||||
for name in defaults:
|
||||
if name in sampling_kwargs:
|
||||
backend_kwargs[name] = sampling_kwargs.pop(name)
|
||||
|
||||
return sampling_kwargs, backend_kwargs
|
||||
|
||||
|
||||
def decode_video(
|
||||
backend: str,
|
||||
loader_cls,
|
||||
data: bytes,
|
||||
target: VideoTargetMetadata,
|
||||
sampling_kwargs: dict[str, Any],
|
||||
backend_kwargs: dict[str, Any],
|
||||
*,
|
||||
frame_recovery: bool,
|
||||
):
|
||||
"""Decode a sampled video, importing only the selected backend."""
|
||||
_get_backend_option_defaults(backend)
|
||||
if frame_recovery and backend != "opencv":
|
||||
error = (
|
||||
ValueError if backend == PYNVVIDEOCODEC_VIDEO_BACKEND else AssertionError
|
||||
)
|
||||
raise error(f"frame_recovery is not supported by the {backend!r} backend")
|
||||
|
||||
decoder_kwargs = dict(backend_kwargs)
|
||||
if backend == "opencv":
|
||||
decoder_kwargs["frame_recovery"] = frame_recovery
|
||||
module = import_module(f".{backend}", __name__)
|
||||
decoder = getattr(module, f"decode_{backend}")
|
||||
return decoder(
|
||||
loader_cls,
|
||||
data,
|
||||
target,
|
||||
sampling_kwargs,
|
||||
**decoder_kwargs,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"PYNVVIDEOCODEC_DEFAULT_HW_DECODERS",
|
||||
"PYNVVIDEOCODEC_VIDEO_BACKEND",
|
||||
"VideoDecoderBackend",
|
||||
"VideoSourceMetadata",
|
||||
"VideoTargetMetadata",
|
||||
"check_frame_pixel_limit",
|
||||
"decode_video",
|
||||
"resolve_video_backend_kwargs",
|
||||
]
|
||||
@@ -0,0 +1,37 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from typing import Literal, NamedTuple
|
||||
|
||||
from vllm import envs
|
||||
|
||||
PYNVVIDEOCODEC_VIDEO_BACKEND: Literal["pynvvideocodec"] = "pynvvideocodec"
|
||||
PYNVVIDEOCODEC_DEFAULT_HW_DECODERS = 2
|
||||
|
||||
|
||||
class VideoTargetMetadata(NamedTuple):
|
||||
"""Metadata describing the requested video sample."""
|
||||
|
||||
num_frames: int
|
||||
fps: float
|
||||
max_duration: float
|
||||
|
||||
|
||||
class VideoSourceMetadata(NamedTuple):
|
||||
"""Metadata describing the encoded video source."""
|
||||
|
||||
total_frames_num: int
|
||||
original_fps: float
|
||||
duration: float
|
||||
|
||||
|
||||
def check_frame_pixel_limit(width: int, height: int) -> None:
|
||||
"""Reject video frames exceeding ``VLLM_MAX_IMAGE_PIXELS``."""
|
||||
max_pixels = envs.VLLM_MAX_IMAGE_PIXELS
|
||||
if max_pixels > 0 and width * height > max_pixels:
|
||||
raise ValueError(
|
||||
f"Video frame dimensions {width}x{height} "
|
||||
f"({width * height} pixels) exceed the maximum of "
|
||||
f"{max_pixels} pixels. Set VLLM_MAX_IMAGE_PIXELS to "
|
||||
f"increase this limit."
|
||||
)
|
||||
@@ -0,0 +1,151 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import threading
|
||||
from typing import Any, ClassVar
|
||||
|
||||
import numpy.typing as npt
|
||||
|
||||
from vllm.logger import init_logger
|
||||
|
||||
from .base import VideoSourceMetadata, VideoTargetMetadata, check_frame_pixel_limit
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def decode_deepstream(
|
||||
loader_cls,
|
||||
data: bytes,
|
||||
target: VideoTargetMetadata,
|
||||
sampling_kwargs: dict,
|
||||
*,
|
||||
pool_size: int | None = None,
|
||||
timeout_sec: float = 120.0,
|
||||
) -> tuple[npt.NDArray, VideoSourceMetadata, list[int], list[int]]:
|
||||
from nvidia.deepstream_videodecode import probe_metadata
|
||||
|
||||
total_frames, original_fps, duration, _width, _height, codec = probe_metadata(data)
|
||||
check_frame_pixel_limit(_width, _height)
|
||||
source = loader_cls._prepare_source(
|
||||
VideoSourceMetadata(total_frames, original_fps, duration)
|
||||
)
|
||||
frame_idx = loader_cls.compute_frames_index_to_sample(
|
||||
source=source, target=target, **sampling_kwargs
|
||||
)
|
||||
frames, valid = DeepStreamVideoBackendMixin.decode_indices(
|
||||
data,
|
||||
frame_idx,
|
||||
source,
|
||||
codec=codec,
|
||||
pool_size=pool_size,
|
||||
timeout_sec=timeout_sec,
|
||||
)
|
||||
return frames, source, frame_idx, valid
|
||||
|
||||
|
||||
class DeepStreamVideoBackendMixin:
|
||||
"""NVIDIA DeepStream (NVDEC) GPU-decode codec utilities.
|
||||
|
||||
Decoding runs on a shared pool of daemon threads inside one CUDA
|
||||
context (see the ``nvidia-deepstream-videodecode-cu13`` package). The
|
||||
container bytes are pushed into an ``appsrc`` GStreamer pipeline, so no
|
||||
local file path is required — HTTP and base64 sources decode identically
|
||||
to local files.
|
||||
|
||||
Like the OpenCV/PyAV mixins, this provides only the codec layer.
|
||||
Frame *selection* lives in the loader's
|
||||
``compute_frames_index_to_sample`` and arrives here as an explicit
|
||||
list of frame indices.
|
||||
"""
|
||||
|
||||
# Process-wide lazy decode pool, shared across all DeepStream backends.
|
||||
_pool: ClassVar[Any] = None
|
||||
_pool_lock: ClassVar[Any] = None
|
||||
|
||||
@classmethod
|
||||
def _get_pool(cls, pool_size: int | None = None):
|
||||
"""Lazy-initialize the shared decode pool on first use.
|
||||
|
||||
``pool_size`` (number of decode worker threads) comes from
|
||||
``--media-io-kwargs`` (``{"video": {"pool_size": N}}``); when unset it
|
||||
defaults to the existing ``VLLM_MEDIA_LOADING_THREAD_COUNT`` so no
|
||||
DeepStream-specific env var is needed. The pool is a process-wide
|
||||
singleton, so the first decode's value wins.
|
||||
"""
|
||||
if cls._pool is not None:
|
||||
return cls._pool
|
||||
if cls._pool_lock is None:
|
||||
cls._pool_lock = threading.Lock()
|
||||
with cls._pool_lock:
|
||||
if cls._pool is not None:
|
||||
return cls._pool
|
||||
import os
|
||||
|
||||
from nvidia.deepstream_videodecode import DecodePool
|
||||
|
||||
if pool_size is None:
|
||||
pool_size = int(os.environ.get("VLLM_MEDIA_LOADING_THREAD_COUNT", 8))
|
||||
pool_size = max(1, min(int(pool_size), 16))
|
||||
logger.info(
|
||||
"[DeepStream] initializing decode pool with %d workers",
|
||||
pool_size,
|
||||
)
|
||||
cls._pool = DecodePool(num_workers=pool_size)
|
||||
return cls._pool
|
||||
|
||||
@classmethod
|
||||
def decode_indices(
|
||||
cls,
|
||||
data: bytes,
|
||||
frame_indices: list[int],
|
||||
source: VideoSourceMetadata,
|
||||
codec: str = "",
|
||||
pool_size: int | None = None,
|
||||
timeout_sec: float = 120.0,
|
||||
) -> tuple[npt.NDArray, list[int]]:
|
||||
"""Decode the requested frame indices from raw container bytes.
|
||||
|
||||
The whole stream is decoded; the pool keeps exactly the frames whose
|
||||
decode-order index is in ``frame_indices`` (1:1, frame-exact) and
|
||||
sends EOS once the last one is matched.
|
||||
|
||||
``codec`` (e.g. ``"h264"``/``"hevc"``) lets the pool keep its NVDEC
|
||||
session warm across same-codec streams and rebuild only on a codec
|
||||
change. Frames are returned as a CPU NHWC uint8 array so the
|
||||
upstream multimodal parser sees the same shape as the other
|
||||
backends.
|
||||
"""
|
||||
if not frame_indices:
|
||||
raise ValueError("DeepStream backend received no frame indices")
|
||||
|
||||
result = cls._get_pool(pool_size).decode(
|
||||
data,
|
||||
target_indices=frame_indices,
|
||||
codec=codec,
|
||||
max_frames=len(frame_indices),
|
||||
timeout_sec=timeout_sec,
|
||||
)
|
||||
if result.error:
|
||||
raise ValueError(f"DeepStream decode failed: {result.error}")
|
||||
if result.frames is None or result.n_kept == 0:
|
||||
raise ValueError("DeepStream decode produced no frames")
|
||||
|
||||
valid = frame_indices[: result.n_kept]
|
||||
# GPU -> CPU NHWC uint8 at the codec boundary (one PCIe copy); keeps
|
||||
# the array shape identical to the OpenCV/PyAV backends. Copy into
|
||||
# PINNED host memory (reused across calls by PyTorch's pinned caching
|
||||
# allocator) so the D2H runs at full PCIe bandwidth (~13 GB/s) rather
|
||||
# than the ~1 GB/s pageable path that plain ``.cpu()`` takes — ~12x
|
||||
# faster for a 1080p x8 frame batch (~46ms -> ~4ms). ``numpy()`` keeps
|
||||
# the pinned tensor alive via the array's base.
|
||||
import torch
|
||||
|
||||
gpu = result.frames
|
||||
if gpu.is_cuda:
|
||||
host = torch.empty(gpu.shape, dtype=gpu.dtype, pin_memory=True)
|
||||
host.copy_(gpu, non_blocking=True)
|
||||
torch.cuda.current_stream().synchronize()
|
||||
arr = host.numpy()
|
||||
else:
|
||||
arr = gpu.numpy()
|
||||
return arr, valid
|
||||
@@ -0,0 +1,295 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from io import BytesIO
|
||||
|
||||
import numpy as np
|
||||
import numpy.typing as npt
|
||||
|
||||
from vllm.logger import init_logger
|
||||
from vllm.utils.import_utils import PlaceholderModule
|
||||
|
||||
from .base import (
|
||||
VideoSourceMetadata,
|
||||
VideoTargetMetadata,
|
||||
check_frame_pixel_limit,
|
||||
)
|
||||
|
||||
try:
|
||||
import cv2
|
||||
import cv2.videoio_registry as vr
|
||||
except ImportError:
|
||||
cv2 = PlaceholderModule("cv2")
|
||||
vr = PlaceholderModule("cv2").placeholder_attr("videoio_registry")
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def decode_opencv(
|
||||
loader_cls,
|
||||
data: bytes,
|
||||
target: VideoTargetMetadata,
|
||||
sampling_kwargs: dict,
|
||||
*,
|
||||
frame_recovery: bool = False,
|
||||
) -> tuple[npt.NDArray, VideoSourceMetadata, list[int], list[int]]:
|
||||
cap = OpenCVVideoBackendMixin.open_video_capture(data)
|
||||
check_frame_pixel_limit(
|
||||
int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
|
||||
int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)),
|
||||
)
|
||||
source = loader_cls._prepare_source(OpenCVVideoBackendMixin.get_video_metadata(cap))
|
||||
frame_idx = loader_cls.compute_frames_index_to_sample(
|
||||
source=source, target=target, **sampling_kwargs
|
||||
)
|
||||
frames, valid = OpenCVVideoBackendMixin.read_frames(
|
||||
cap,
|
||||
frame_idx,
|
||||
total_frames_num=source.total_frames_num,
|
||||
frame_recovery=frame_recovery,
|
||||
)
|
||||
return frames, source, frame_idx, valid
|
||||
|
||||
|
||||
class OpenCVVideoBackendMixin:
|
||||
@staticmethod
|
||||
def get_cv2_video_api():
|
||||
api_pref = None
|
||||
for backend in vr.getStreamBufferedBackends():
|
||||
if not vr.hasBackend(backend):
|
||||
continue
|
||||
if not vr.isBackendBuiltIn(backend):
|
||||
_, abi, api = vr.getStreamBufferedBackendPluginVersion(backend)
|
||||
if abi < 1 or (abi == 1 and api < 2):
|
||||
continue
|
||||
api_pref = backend
|
||||
break
|
||||
return api_pref
|
||||
|
||||
@classmethod
|
||||
def open_video_capture(cls, data: bytes) -> "cv2.VideoCapture":
|
||||
backend = cls.get_cv2_video_api()
|
||||
cap = cv2.VideoCapture(BytesIO(data), backend, [])
|
||||
if not cap.isOpened():
|
||||
raise ValueError("Could not open video stream")
|
||||
return cap
|
||||
|
||||
@staticmethod
|
||||
def get_video_metadata(cap: "cv2.VideoCapture") -> VideoSourceMetadata:
|
||||
total_frames_num = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
original_fps = cap.get(cv2.CAP_PROP_FPS)
|
||||
duration = total_frames_num / original_fps if original_fps > 0 else 0
|
||||
return VideoSourceMetadata(
|
||||
total_frames_num=total_frames_num,
|
||||
original_fps=original_fps,
|
||||
duration=duration,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _can_use_for_recovery(
|
||||
cls,
|
||||
idx: int,
|
||||
failed_frames: list[int],
|
||||
next_target_map: dict[int, int],
|
||||
total_frames: int,
|
||||
) -> bool:
|
||||
"""Check if current frame can recover the oldest failed frame."""
|
||||
if not failed_frames:
|
||||
return False
|
||||
oldest_failed = failed_frames[0]
|
||||
limit = next_target_map.get(oldest_failed, total_frames)
|
||||
return idx < limit
|
||||
|
||||
@classmethod
|
||||
def _read_frames_with_recovery(
|
||||
cls,
|
||||
cap: "cv2.VideoCapture",
|
||||
frame_indices: list[int],
|
||||
total_frames: int,
|
||||
) -> tuple[npt.NDArray, list[int], dict[int, int]]:
|
||||
"""
|
||||
Read frames with dynamic window forward-scan recovery.
|
||||
|
||||
When a target frame fails to load, the next successfully grabbed
|
||||
frame (before the next target frame) will be used to recover it.
|
||||
|
||||
Args:
|
||||
cap: OpenCV VideoCapture object
|
||||
frame_indices: Sorted list of target frame indices to load
|
||||
total_frames: Total number of frames in the video
|
||||
|
||||
Returns:
|
||||
Tuple of (frames_array, valid_frame_indices, recovered_map)
|
||||
- frames_array: Array of loaded frames
|
||||
- valid_frame_indices: List of frame indices that were loaded
|
||||
- recovered_map: Dict mapping recovered_idx -> source_idx
|
||||
"""
|
||||
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||||
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||||
|
||||
assert width > 0 and height > 0, (
|
||||
f"Invalid video frame size: width={width}, height={height}"
|
||||
)
|
||||
|
||||
frame_idx_set = set(frame_indices)
|
||||
max_frame_idx = frame_indices[-1] if frame_indices else 0
|
||||
|
||||
# Build map: target_idx -> next_target_idx (for recovery window)
|
||||
next_target_map: dict[int, int] = {}
|
||||
for k in range(len(frame_indices) - 1):
|
||||
next_target_map[frame_indices[k]] = frame_indices[k + 1]
|
||||
next_target_map[frame_indices[-1]] = total_frames
|
||||
|
||||
frames_list: list[npt.NDArray] = []
|
||||
valid_frame_indices: list[int] = []
|
||||
failed_frames_idx: list[int] = []
|
||||
recovered_map: dict[int, int] = {}
|
||||
|
||||
i = 0
|
||||
for idx in range(max_frame_idx + 1):
|
||||
is_target_frame = idx in frame_idx_set
|
||||
|
||||
# Attempt to grab the current frame
|
||||
ok = cap.grab()
|
||||
|
||||
if not ok:
|
||||
if is_target_frame:
|
||||
logger.debug(
|
||||
"Failed to grab frame %d during video loading.",
|
||||
idx,
|
||||
)
|
||||
failed_frames_idx.append(idx)
|
||||
continue
|
||||
|
||||
# Check if we should retrieve: target frame OR can recover a failed one
|
||||
can_recover = cls._can_use_for_recovery(
|
||||
idx, failed_frames_idx, next_target_map, total_frames
|
||||
)
|
||||
|
||||
if is_target_frame or can_recover:
|
||||
ret, frame = cap.retrieve()
|
||||
|
||||
if ret and frame is not None and frame.size > 0:
|
||||
rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||||
frames_list.append(rgb_frame)
|
||||
valid_frame_indices.append(idx)
|
||||
i += 1
|
||||
|
||||
if can_recover:
|
||||
recovered_idx = failed_frames_idx.pop(0)
|
||||
recovered_map[recovered_idx] = idx
|
||||
logger.info(
|
||||
"Recovered frame %d using frame %d (delay: %d)",
|
||||
recovered_idx,
|
||||
idx,
|
||||
idx - recovered_idx,
|
||||
)
|
||||
elif is_target_frame:
|
||||
logger.debug(
|
||||
"Failed to retrieve frame %d during video loading.",
|
||||
idx,
|
||||
)
|
||||
failed_frames_idx.append(idx)
|
||||
|
||||
# Log any remaining failed frames
|
||||
for failed_idx in failed_frames_idx:
|
||||
logger.debug(
|
||||
"Frame %d could not be recovered (end of video).",
|
||||
failed_idx,
|
||||
)
|
||||
|
||||
# Stack frames
|
||||
if frames_list:
|
||||
frames = np.stack(frames_list)
|
||||
else:
|
||||
frames = np.empty((0, height, width, 3), dtype=np.uint8)
|
||||
|
||||
return frames, valid_frame_indices, recovered_map
|
||||
|
||||
@classmethod
|
||||
def _read_frames_no_recovery(
|
||||
cls,
|
||||
cap,
|
||||
frame_indices: set[int],
|
||||
max_frame_idx: int,
|
||||
) -> tuple[npt.NDArray, list[int]]:
|
||||
num_expected_frames = len(frame_indices)
|
||||
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||||
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||||
frames = np.empty((num_expected_frames, height, width, 3), dtype=np.uint8)
|
||||
|
||||
i = 0
|
||||
valid_frame_indices = []
|
||||
for idx in range(max_frame_idx + 1):
|
||||
ok = cap.grab()
|
||||
if not ok:
|
||||
# Frame is broken/unreadable, skip it
|
||||
if idx in frame_indices:
|
||||
logger.debug(
|
||||
"Failed to grab frame %d during video loading. "
|
||||
"This frame will be skipped.",
|
||||
idx,
|
||||
)
|
||||
continue
|
||||
if idx in frame_indices:
|
||||
ret, frame = cap.retrieve()
|
||||
if ret:
|
||||
frames[i] = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||||
valid_frame_indices.append(idx)
|
||||
i += 1
|
||||
else:
|
||||
# retrieve() failed even though grab() succeeded
|
||||
logger.debug(
|
||||
"Failed to retrieve frame %d during video loading. "
|
||||
"This frame will be skipped.",
|
||||
idx,
|
||||
)
|
||||
|
||||
valid_num_frames = len(valid_frame_indices)
|
||||
if valid_num_frames < num_expected_frames:
|
||||
logger.warning(
|
||||
"Video loading completed with %d broken/unreadable frames. "
|
||||
"Expected %d frames but only loaded %d frames.",
|
||||
num_expected_frames - valid_num_frames,
|
||||
num_expected_frames,
|
||||
valid_num_frames,
|
||||
)
|
||||
|
||||
return frames[:valid_num_frames], valid_frame_indices
|
||||
|
||||
@classmethod
|
||||
def read_frames(
|
||||
cls,
|
||||
cap: "cv2.VideoCapture",
|
||||
frame_idx: list[int],
|
||||
total_frames_num: int,
|
||||
*,
|
||||
frame_recovery: bool = False,
|
||||
) -> tuple[npt.NDArray, list[int]]:
|
||||
if frame_recovery:
|
||||
num_frames_to_sample = len(frame_idx)
|
||||
frames, valid_frame_indices, recovered_map = cls._read_frames_with_recovery(
|
||||
cap, frame_idx, total_frames_num
|
||||
)
|
||||
|
||||
if recovered_map:
|
||||
logger.info(
|
||||
"Frame recovery: %d frames recovered using forward scan.",
|
||||
len(recovered_map),
|
||||
)
|
||||
else:
|
||||
frame_idx_set = set(frame_idx)
|
||||
num_frames_to_sample = len(frame_idx_set)
|
||||
frames, valid_frame_indices = cls._read_frames_no_recovery(
|
||||
cap, frame_idx_set, max(frame_idx)
|
||||
)
|
||||
valid_num_frames = len(valid_frame_indices)
|
||||
if valid_num_frames < num_frames_to_sample:
|
||||
logger.warning(
|
||||
"Video loading completed with %d broken/unreadable frames. "
|
||||
"Expected to sample %d frames but only loaded %d frames.",
|
||||
num_frames_to_sample - valid_num_frames,
|
||||
num_frames_to_sample,
|
||||
valid_num_frames,
|
||||
)
|
||||
return frames, valid_frame_indices
|
||||
@@ -0,0 +1,111 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from io import BytesIO
|
||||
|
||||
import numpy as np
|
||||
import numpy.typing as npt
|
||||
|
||||
from vllm.utils.import_utils import PlaceholderModule
|
||||
|
||||
from .base import (
|
||||
VideoSourceMetadata,
|
||||
VideoTargetMetadata,
|
||||
check_frame_pixel_limit,
|
||||
)
|
||||
|
||||
try:
|
||||
import av
|
||||
except ImportError:
|
||||
av = PlaceholderModule("av") # type: ignore[assignment]
|
||||
|
||||
|
||||
def decode_pyav(
|
||||
loader_cls,
|
||||
data: bytes,
|
||||
target: VideoTargetMetadata,
|
||||
sampling_kwargs: dict,
|
||||
) -> tuple[npt.NDArray, VideoSourceMetadata, list[int], list[int]]:
|
||||
with av.open(BytesIO(data)) as container:
|
||||
stream = container.streams.video[0]
|
||||
check_frame_pixel_limit(stream.width, stream.height)
|
||||
source = loader_cls._prepare_source(
|
||||
PyAVVideoBackendMixin.get_metadata(container)
|
||||
)
|
||||
frame_idx = loader_cls.compute_frames_index_to_sample(
|
||||
source=source, target=target, **sampling_kwargs
|
||||
)
|
||||
frames, valid = PyAVVideoBackendMixin.decode_frames(
|
||||
container, frame_idx, source.original_fps, source.duration
|
||||
)
|
||||
return frames, source, frame_idx, valid
|
||||
|
||||
|
||||
class PyAVVideoBackendMixin:
|
||||
"""PyAV (in-process FFmpeg bindings) codec utilities.
|
||||
|
||||
Reads stream metadata and decodes target frames via per-frame
|
||||
``container.seek()``. The seek releases the GIL between frames and
|
||||
scales with the number of sampled frames rather than the video
|
||||
length, enabling concurrent decoding under serving load.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def get_metadata(
|
||||
container: "av.container.InputContainer",
|
||||
) -> VideoSourceMetadata:
|
||||
if not container.streams.video:
|
||||
raise ValueError("No video streams found in container")
|
||||
stream = container.streams.video[0]
|
||||
total_frames = stream.frames or 0
|
||||
fps = float(stream.average_rate) if stream.average_rate else 0.0
|
||||
duration = float(stream.duration * stream.time_base) if stream.duration else 0.0
|
||||
if total_frames == 0 and duration > 0 and fps > 0:
|
||||
total_frames = int(duration * fps)
|
||||
return VideoSourceMetadata(total_frames, fps, duration)
|
||||
|
||||
@staticmethod
|
||||
def decode_frames(
|
||||
container: "av.container.InputContainer",
|
||||
frame_indices: list[int],
|
||||
fps: float,
|
||||
duration: float,
|
||||
) -> tuple[npt.NDArray, list[int]]:
|
||||
"""Decode target frames via per-frame seek + forward decode to PTS."""
|
||||
stream = container.streams.video[0]
|
||||
# SLICE parallelizes within a single frame without the
|
||||
# one-frame-per-thread latency penalty of FRAME threading.
|
||||
stream.thread_type = "SLICE"
|
||||
time_base = stream.time_base
|
||||
|
||||
frames_list: list[npt.NDArray] = []
|
||||
valid_indices: list[int] = []
|
||||
frame_interval = 1.0 / fps if fps > 0 else 0.1
|
||||
max_ts = max(0.0, duration - frame_interval) if duration > 0 else float("inf")
|
||||
|
||||
decoder = None
|
||||
last_pts = None
|
||||
for idx in frame_indices:
|
||||
ts = min(idx / fps, max_ts) if fps > 0 else 0.0
|
||||
pts = int(ts / time_base)
|
||||
# seek() snaps backward to a keyframe; reuse the running decoder
|
||||
# while targets advance monotonically to avoid re-decoding the
|
||||
# GOP prefix once per requested frame.
|
||||
if decoder is None or last_pts is None or pts <= last_pts:
|
||||
container.seek(pts, stream=stream)
|
||||
decoder = container.decode(video=0)
|
||||
chosen = None
|
||||
for frame in decoder:
|
||||
if frame.pts is not None and frame.pts >= pts:
|
||||
chosen = frame
|
||||
last_pts = frame.pts
|
||||
break
|
||||
if chosen is not None:
|
||||
frames_list.append(chosen.to_ndarray(format="rgb24"))
|
||||
valid_indices.append(idx)
|
||||
else:
|
||||
decoder = None
|
||||
|
||||
if not frames_list:
|
||||
return np.empty((0,), dtype=np.uint8), valid_indices
|
||||
return np.stack(frames_list), valid_indices
|
||||
@@ -0,0 +1,372 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
import threading
|
||||
from contextlib import contextmanager, suppress
|
||||
from typing import ClassVar, NamedTuple
|
||||
|
||||
import numpy as np
|
||||
import numpy.typing as npt
|
||||
|
||||
from vllm.logger import init_logger
|
||||
from vllm.utils.mem_constants import MiB_bytes
|
||||
|
||||
from .base import (
|
||||
PYNVVIDEOCODEC_DEFAULT_HW_DECODERS,
|
||||
VideoSourceMetadata,
|
||||
VideoTargetMetadata,
|
||||
check_frame_pixel_limit,
|
||||
)
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def decode_pynvvideocodec(
|
||||
loader_cls,
|
||||
data: bytes,
|
||||
target: VideoTargetMetadata,
|
||||
sampling_kwargs: dict,
|
||||
*,
|
||||
hw_decoders: int = PYNVVIDEOCODEC_DEFAULT_HW_DECODERS,
|
||||
) -> tuple[npt.NDArray, VideoSourceMetadata, list[int], list[int]]:
|
||||
PyNvVideoCodecVideoBackendMixin._configure_decoder_slots(hw_decoders)
|
||||
return PyNvVideoCodecVideoBackendMixin.decode_frames_pynvvideocodec(
|
||||
loader_cls,
|
||||
data,
|
||||
target,
|
||||
**sampling_kwargs,
|
||||
)
|
||||
|
||||
|
||||
class PyNvVideoCodecSourceMetadata(NamedTuple):
|
||||
"""Metadata needed before GPU video decode."""
|
||||
|
||||
source: VideoSourceMetadata
|
||||
width: int
|
||||
height: int
|
||||
|
||||
|
||||
# Per-decoder upper bound reserved for persistent PyNvVideoCodec surfaces.
|
||||
PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES = 128 * MiB_bytes
|
||||
PYNVVIDEOCODEC_DECODER_CACHE_SIZE = 2
|
||||
# Per-API-server CUDA context and driver allocation, measured with
|
||||
# PyNvVideoCodec 2.0.4 on H100.
|
||||
PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES = int(1.8 * 1024 * MiB_bytes)
|
||||
|
||||
|
||||
def validate_pynvvideocodec_hw_decoders(hw_decoders: object) -> int:
|
||||
if (
|
||||
isinstance(hw_decoders, bool)
|
||||
or not isinstance(hw_decoders, int)
|
||||
or hw_decoders < 1
|
||||
):
|
||||
raise ValueError("hw_decoders must be a positive integer")
|
||||
return hw_decoders
|
||||
|
||||
|
||||
def _pynvvideocodec_exception_types(nvc) -> tuple[type[Exception], ...]:
|
||||
return tuple(
|
||||
exception_type
|
||||
for name in dir(nvc)
|
||||
if name.startswith("PyNvVCException")
|
||||
and isinstance((exception_type := getattr(nvc, name)), type)
|
||||
and issubclass(exception_type, Exception)
|
||||
)
|
||||
|
||||
|
||||
def _pynvvc_frames_to_nhwc(frames):
|
||||
"""Return a stacked PyNvVideoCodec frame batch as contiguous NHWC."""
|
||||
if frames.shape[-1] != 3 and frames.shape[-3] == 3:
|
||||
frames = frames.permute(0, 2, 3, 1)
|
||||
return frames.contiguous()
|
||||
|
||||
|
||||
class PyNvVideoCodecDecoderSlot:
|
||||
"""A retained PyNv decoder slot and its CUDA stream.
|
||||
|
||||
The decoder is reused across requests: ``reconfigure_decoder`` repoints the
|
||||
existing decoder at each new source instead of paying a fresh
|
||||
``SimpleDecoder`` construction per request. Construction (CUVID parser +
|
||||
decoder + surface-pool allocation) is the dominant per-request cost, so
|
||||
reconfiguring is far cheaper. A single decoder serves both metadata
|
||||
(``len``/``get_stream_metadata``) and frame decode -- no separate
|
||||
metadata decoder.
|
||||
"""
|
||||
|
||||
def __init__(self, stream) -> None:
|
||||
self.stream = stream
|
||||
self.decoder = None
|
||||
self.source_path: str | None = None
|
||||
|
||||
def invalidate(self) -> None:
|
||||
self.decoder = None
|
||||
self.source_path = None
|
||||
|
||||
def _construct(self, file_path: str, nvc, device_index: int) -> None:
|
||||
self.invalidate()
|
||||
decoder = nvc.SimpleDecoder(
|
||||
file_path,
|
||||
output_color_type=nvc.OutputColorType.RGB,
|
||||
use_device_memory=True,
|
||||
need_scanned_stream_metadata=True,
|
||||
gpu_id=device_index,
|
||||
cuda_stream=self.stream.cuda_stream,
|
||||
decoder_cache_size=PYNVVIDEOCODEC_DECODER_CACHE_SIZE,
|
||||
)
|
||||
self.decoder = decoder
|
||||
self.source_path = file_path
|
||||
|
||||
def get_decoder(self, file_path: str, nvc, device_index: int):
|
||||
if self.decoder is None:
|
||||
self._construct(file_path, nvc, device_index)
|
||||
elif self.source_path != file_path:
|
||||
try:
|
||||
self.decoder.reconfigure_decoder(file_path)
|
||||
self.source_path = file_path
|
||||
except Exception:
|
||||
# reconfigure unsupported/unsafe for this source -> rebuild.
|
||||
self._construct(file_path, nvc, device_index)
|
||||
return self.decoder
|
||||
|
||||
|
||||
class _PyNvDecoderPool:
|
||||
"""Process-wide singleton managing PyNvVideoCodec decoder slot state.
|
||||
|
||||
Prevents subclass counter shadowing (GHSA-j682-9xp5-rrf3) by storing
|
||||
all mutable pool state in a single module-level instance rather than
|
||||
in ClassVar attributes that get shadowed by Python's augmented
|
||||
assignment semantics on subclasses.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.slots: list[PyNvVideoCodecDecoderSlot] = []
|
||||
self.active: int = 0
|
||||
self.cond: threading.Condition = threading.Condition()
|
||||
self.max_slots: int | None = None
|
||||
|
||||
def configure(self, hw_decoders: int) -> None:
|
||||
with self.cond:
|
||||
if self.max_slots is None:
|
||||
self.max_slots = hw_decoders
|
||||
elif self.max_slots != hw_decoders:
|
||||
raise RuntimeError(
|
||||
"PyNvVideoCodec decoder count is already configured as "
|
||||
f"{self.max_slots}, got {hw_decoders}"
|
||||
)
|
||||
|
||||
|
||||
_pynv_decoder_pool = _PyNvDecoderPool()
|
||||
|
||||
|
||||
class PyNvVideoCodecVideoBackendMixin:
|
||||
"""PyNvVideoCodec utilities for GPU-backed frame decode."""
|
||||
|
||||
_DEVICE_INDEX: ClassVar[int] = 0
|
||||
|
||||
@classmethod
|
||||
def _create_decoder_slot(cls) -> PyNvVideoCodecDecoderSlot:
|
||||
import torch
|
||||
|
||||
return PyNvVideoCodecDecoderSlot(torch.cuda.Stream(device=cls._DEVICE_INDEX))
|
||||
|
||||
@classmethod
|
||||
def _configure_decoder_slots(cls, hw_decoders: object) -> None:
|
||||
hw_decoders = validate_pynvvideocodec_hw_decoders(hw_decoders)
|
||||
_pynv_decoder_pool.configure(hw_decoders)
|
||||
|
||||
@staticmethod
|
||||
@contextmanager
|
||||
def _torch_stream_context(stream):
|
||||
import torch
|
||||
|
||||
torch.accelerator.set_device_index(stream.device.index)
|
||||
previous_stream = torch.accelerator.current_stream()
|
||||
torch.accelerator.set_stream(stream)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
torch.accelerator.set_stream(previous_stream)
|
||||
|
||||
@classmethod
|
||||
@contextmanager
|
||||
def _borrow_decoder_slot(cls):
|
||||
pool = _pynv_decoder_pool
|
||||
create_slot = False
|
||||
with pool.cond:
|
||||
if pool.max_slots is None:
|
||||
raise RuntimeError("PyNvVideoCodec decoder slots are not configured")
|
||||
while True:
|
||||
if pool.slots:
|
||||
slot = pool.slots.pop()
|
||||
break
|
||||
if pool.active < pool.max_slots:
|
||||
pool.active += 1
|
||||
create_slot = True
|
||||
break
|
||||
pool.cond.wait()
|
||||
|
||||
if create_slot:
|
||||
try:
|
||||
slot = cls._create_decoder_slot()
|
||||
except Exception:
|
||||
with pool.cond:
|
||||
pool.active -= 1
|
||||
pool.cond.notify()
|
||||
raise
|
||||
|
||||
borrow_succeeded = False
|
||||
try:
|
||||
yield slot
|
||||
borrow_succeeded = True
|
||||
finally:
|
||||
if not borrow_succeeded:
|
||||
slot.invalidate()
|
||||
with pool.cond:
|
||||
pool.slots.append(slot)
|
||||
pool.cond.notify()
|
||||
|
||||
@staticmethod
|
||||
def _metadata_value(metadata, *names: str, default=None):
|
||||
for name in names:
|
||||
value = getattr(metadata, name, None)
|
||||
if value is not None:
|
||||
return value
|
||||
return default
|
||||
|
||||
@classmethod
|
||||
def _read_source_metadata(
|
||||
cls,
|
||||
file_path: str,
|
||||
nvc,
|
||||
) -> PyNvVideoCodecSourceMetadata:
|
||||
with cls._borrow_decoder_slot() as decoder_slot:
|
||||
with cls._torch_stream_context(decoder_slot.stream):
|
||||
decoder = decoder_slot.get_decoder(
|
||||
file_path, nvc, device_index=cls._DEVICE_INDEX
|
||||
)
|
||||
metadata = decoder.get_stream_metadata()
|
||||
total_frames_num = len(decoder)
|
||||
width = int(cls._metadata_value(metadata, "width", default=0))
|
||||
height = int(cls._metadata_value(metadata, "height", default=0))
|
||||
original_fps = float(
|
||||
cls._metadata_value(
|
||||
metadata,
|
||||
"average_fps",
|
||||
"avg_frame_rate",
|
||||
"frame_rate",
|
||||
"frameRate",
|
||||
default=0.0,
|
||||
)
|
||||
)
|
||||
duration = float(
|
||||
cls._metadata_value(metadata, "duration", default=0.0)
|
||||
or (total_frames_num / original_fps if original_fps > 0 else 0.0)
|
||||
)
|
||||
if total_frames_num <= 0:
|
||||
raise ValueError("Could not determine video frame count")
|
||||
if width <= 0 or height <= 0:
|
||||
raise ValueError("Could not determine video dimensions")
|
||||
return PyNvVideoCodecSourceMetadata(
|
||||
source=VideoSourceMetadata(total_frames_num, original_fps, duration),
|
||||
width=width,
|
||||
height=height,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _decode_to_pinned_host(
|
||||
cls,
|
||||
file_path: str,
|
||||
frame_idx: list[int],
|
||||
nvc,
|
||||
) -> npt.NDArray:
|
||||
import torch
|
||||
|
||||
if not frame_idx:
|
||||
return np.empty((0,), dtype=np.uint8)
|
||||
|
||||
with cls._borrow_decoder_slot() as decoder_slot:
|
||||
stream = decoder_slot.stream
|
||||
with cls._torch_stream_context(stream):
|
||||
try:
|
||||
decoder = decoder_slot.get_decoder(
|
||||
file_path, nvc, device_index=cls._DEVICE_INDEX
|
||||
)
|
||||
decoded_frames = decoder.get_batch_frames_by_index(frame_idx)
|
||||
except Exception as exc:
|
||||
if not isinstance(
|
||||
exc,
|
||||
_pynvvideocodec_exception_types(nvc) + (IndexError,),
|
||||
):
|
||||
raise
|
||||
raise ValueError("Invalid or unsupported video file.") from exc
|
||||
if len(decoded_frames) < len(frame_idx):
|
||||
logger.warning(
|
||||
"pynvvideocodec video loading: expected %d frames but got %d.",
|
||||
len(frame_idx),
|
||||
len(decoded_frames),
|
||||
)
|
||||
torch_frames = [torch.from_dlpack(frame) for frame in decoded_frames]
|
||||
if not torch_frames:
|
||||
return np.empty((0,), dtype=np.uint8)
|
||||
device_frames = torch.stack(torch_frames)
|
||||
if device_frames.ndim != 4:
|
||||
raise ValueError(
|
||||
"PyNvVideoCodec returned frames with unexpected shape "
|
||||
f"{tuple(device_frames.shape)}"
|
||||
)
|
||||
device_frames = _pynvvc_frames_to_nhwc(device_frames)
|
||||
host_frames = torch.empty(
|
||||
device_frames.shape,
|
||||
dtype=device_frames.dtype,
|
||||
device="cpu",
|
||||
pin_memory=True,
|
||||
)
|
||||
host_frames.copy_(device_frames, non_blocking=True)
|
||||
stream.synchronize()
|
||||
host_array = host_frames.numpy()
|
||||
del decoded_frames, torch_frames, device_frames
|
||||
return host_array
|
||||
|
||||
@classmethod
|
||||
def decode_frames_pynvvideocodec(
|
||||
cls,
|
||||
loader_cls,
|
||||
data: bytes,
|
||||
target: VideoTargetMetadata,
|
||||
**kwargs,
|
||||
) -> tuple[npt.NDArray, VideoSourceMetadata, list[int], list[int]]:
|
||||
import PyNvVideoCodec as nvc
|
||||
|
||||
from vllm.multimodal.gpu_ipc_memory import get_mm_gpu_ipc_pool
|
||||
|
||||
temp_fd, temp_path = tempfile.mkstemp(suffix=".mp4")
|
||||
try:
|
||||
with os.fdopen(temp_fd, "wb") as temp_file:
|
||||
temp_file.write(data)
|
||||
|
||||
try:
|
||||
gpu_source = cls._read_source_metadata(temp_path, nvc)
|
||||
except Exception as exc:
|
||||
if not isinstance(exc, _pynvvideocodec_exception_types(nvc)):
|
||||
raise
|
||||
raise ValueError("Invalid or unsupported video file.") from exc
|
||||
check_frame_pixel_limit(gpu_source.width, gpu_source.height)
|
||||
source = loader_cls._prepare_source(gpu_source.source)
|
||||
frame_idx = loader_cls.compute_frames_index_to_sample(
|
||||
source=source, target=target, **kwargs
|
||||
)
|
||||
raw_frame_bytes = len(frame_idx) * gpu_source.height * gpu_source.width * 3
|
||||
pool = get_mm_gpu_ipc_pool()
|
||||
if pool is None or raw_frame_bytes == 0:
|
||||
frames = cls._decode_to_pinned_host(temp_path, frame_idx, nvc)
|
||||
else:
|
||||
with pool.acquire(raw_frame_bytes):
|
||||
frames = cls._decode_to_pinned_host(temp_path, frame_idx, nvc)
|
||||
finally:
|
||||
with suppress(FileNotFoundError):
|
||||
os.unlink(temp_path)
|
||||
|
||||
valid_frame_indices = frame_idx[: int(frames.shape[0])]
|
||||
return frames, source, frame_idx, valid_frame_indices
|
||||
@@ -0,0 +1,100 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from typing import Literal
|
||||
|
||||
import numpy as np
|
||||
import numpy.typing as npt
|
||||
|
||||
from vllm.utils.import_utils import PlaceholderModule, check_torchcodec_available
|
||||
|
||||
from .base import (
|
||||
VideoSourceMetadata,
|
||||
VideoTargetMetadata,
|
||||
check_frame_pixel_limit,
|
||||
)
|
||||
|
||||
try:
|
||||
from torchcodec.decoders import VideoDecoder
|
||||
except (ImportError, RuntimeError):
|
||||
VideoDecoder = PlaceholderModule("torchcodec").placeholder_attr( # type: ignore[assignment]
|
||||
"decoders.VideoDecoder"
|
||||
)
|
||||
|
||||
|
||||
def decode_torchcodec(
|
||||
loader_cls,
|
||||
data: bytes,
|
||||
target: VideoTargetMetadata,
|
||||
sampling_kwargs: dict,
|
||||
*,
|
||||
num_ffmpeg_threads: int = 0,
|
||||
seek_mode: Literal["exact", "approximate"] = "exact",
|
||||
) -> tuple[npt.NDArray, VideoSourceMetadata, list[int], list[int]]:
|
||||
check_torchcodec_available()
|
||||
decoder = TorchCodecVideoBackendMixin.make_torchcodec_decoder(
|
||||
data,
|
||||
num_ffmpeg_threads=num_ffmpeg_threads,
|
||||
seek_mode=seek_mode,
|
||||
)
|
||||
check_frame_pixel_limit(
|
||||
decoder.metadata.width or 0,
|
||||
decoder.metadata.height or 0,
|
||||
)
|
||||
source = loader_cls._prepare_source(
|
||||
TorchCodecVideoBackendMixin.get_torchcodec_metadata(decoder)
|
||||
)
|
||||
frame_idx = loader_cls.compute_frames_index_to_sample(
|
||||
source=source, target=target, **sampling_kwargs
|
||||
)
|
||||
frames, valid = TorchCodecVideoBackendMixin.decode_torchcodec_frames(
|
||||
decoder, frame_idx
|
||||
)
|
||||
return frames, source, frame_idx, valid
|
||||
|
||||
|
||||
class TorchCodecVideoBackendMixin:
|
||||
"""TorchCodec (FFmpeg-backed, PyTorch-native) codec utilities.
|
||||
|
||||
Builds a :class:`~torchcodec.decoders.VideoDecoder` over the in-memory
|
||||
bytes and extracts the sampled indices with a single batched
|
||||
``get_frames_at`` call, while releasing the GIL during decode.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def make_torchcodec_decoder(
|
||||
data: bytes,
|
||||
*,
|
||||
num_ffmpeg_threads: int = 0,
|
||||
seek_mode: Literal["exact", "approximate"] = "exact",
|
||||
) -> "VideoDecoder":
|
||||
# NHWC matches the (num_frames, H, W, 3) uint8 RGB layout the rest
|
||||
# of the pipeline expects, avoiding a transpose.
|
||||
return VideoDecoder(
|
||||
data,
|
||||
dimension_order="NHWC",
|
||||
num_ffmpeg_threads=num_ffmpeg_threads,
|
||||
seek_mode=seek_mode,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def get_torchcodec_metadata(decoder: "VideoDecoder") -> VideoSourceMetadata:
|
||||
md = decoder.metadata
|
||||
total_frames = md.num_frames or 0
|
||||
fps = float(md.average_fps) if md.average_fps else 0.0
|
||||
duration = float(md.duration_seconds) if md.duration_seconds else 0.0
|
||||
if total_frames == 0 and duration > 0 and fps > 0:
|
||||
total_frames = int(duration * fps)
|
||||
return VideoSourceMetadata(total_frames, fps, duration)
|
||||
|
||||
@staticmethod
|
||||
def decode_torchcodec_frames(
|
||||
decoder: "VideoDecoder",
|
||||
frame_indices: list[int],
|
||||
) -> tuple[npt.NDArray, list[int]]:
|
||||
"""Decode the requested indices in one batched, index-exact call."""
|
||||
if not frame_indices:
|
||||
return np.empty((0,), dtype=np.uint8), []
|
||||
# Note: torchcodec releases the GIL for the entire call
|
||||
batch = decoder.get_frames_at(frame_indices)
|
||||
return batch.data.numpy(), list(frame_indices)
|
||||
Reference in New Issue
Block a user