From fca1878f6a68e90ce86791c482a4005bdc260ec1 Mon Sep 17 00:00:00 2001 From: Yozer Date: Mon, 13 Jul 2026 21:38:50 +0200 Subject: [PATCH] Support current CUDA contexts --- av/video/frame.pxd | 1 + av/video/frame.py | 37 ++++++++++++++++++++++++++++++++----- av/video/frame.pyi | 10 +++++++++- docs/api/video.rst | 4 ++++ tests/test_dlpack.py | 34 ++++++++++++++++++++++++++++++++-- 5 files changed, 78 insertions(+), 8 deletions(-) diff --git a/av/video/frame.pxd b/av/video/frame.pxd index 90547cb9f..0a96311e0 100644 --- a/av/video/frame.pxd +++ b/av/video/frame.pxd @@ -9,6 +9,7 @@ from av.video.reformatter cimport VideoReformatter cdef class CudaContext: cdef readonly int device_id cdef readonly bint primary_ctx + cdef readonly bint current_ctx cdef lib.AVBufferRef* _device_ref cdef dict _frames_cache cdef lib.AVBufferRef* _get_device_ref(self) diff --git a/av/video/frame.py b/av/video/frame.py index 6d6c3f044..511cb028b 100644 --- a/av/video/frame.py +++ b/av/video/frame.py @@ -22,11 +22,28 @@ @cython.final @cython.cclass class CudaContext: - def __cinit__(self, device_id: cython.int = 0, primary_ctx: cython.bint = True): + """A reusable FFmpeg CUDA context for frames imported with DLPack. + + :param int device_id: CUDA device ordinal. + :param bool primary_ctx: Retain the device's primary CUDA context. + :param bool current_ctx: Wrap the CUDA context current on the calling thread + when this object is first used. This is mutually exclusive with + ``primary_ctx``. + """ + + def __cinit__( + self, + device_id: cython.int = 0, + primary_ctx: cython.bint = True, + current_ctx: cython.bint = False, + ): self.device_id = device_id self.primary_ctx = primary_ctx + self.current_ctx = current_ctx self._device_ref = cython.NULL self._frames_cache = {} + if primary_ctx and current_ctx: + raise ValueError("primary_ctx and current_ctx are mutually exclusive") def __dealloc__(self): ref: cython.pointer[lib.AVBufferRef] @@ -53,9 +70,10 @@ def _get_device_ref(self) -> cython.pointer[lib.AVBufferRef]: device_ref = cython.NULL device_bytes = f"{self.device_id}".encode() c_device: cython.p_char = device_bytes - options: Dictionary = Dictionary( - {"primary_ctx": "1" if self.primary_ctx else "0"} - ) + options_dict = {"primary_ctx": "1" if self.primary_ctx else "0"} + if self.current_ctx: + options_dict["current_ctx"] = "1" + options: Dictionary = Dictionary(options_dict) err_check( lib.av_hwdevice_ctx_create( cython.address(device_ref), @@ -1471,6 +1489,7 @@ def from_dlpack( device_id: int | None = None, primary_ctx: bool = True, cuda_context=None, + current_ctx: bool = False, ): if not isinstance(planes, (tuple, list)): planes = (planes,) @@ -1612,7 +1631,11 @@ def from_dlpack( ctx: CudaContext frames_ref: cython.pointer[lib.AVBufferRef] if cuda_context is None: - ctx = CudaContext(device_id=device_id, primary_ctx=primary_ctx) + ctx = CudaContext( + device_id=device_id, + primary_ctx=primary_ctx, + current_ctx=current_ctx, + ) else: if not isinstance(cuda_context, CudaContext): raise TypeError("cuda_context must be a CudaContext") @@ -1624,6 +1647,10 @@ def from_dlpack( raise ValueError( "cuda_context.primary_ctx does not match primary_ctx" ) + if bool(cuda_context.current_ctx) != bool(current_ctx): + raise ValueError( + "cuda_context.current_ctx does not match current_ctx" + ) ctx = cython.cast(CudaContext, cuda_context) frames_ref = ctx.get_frames_ctx(sw_fmt, width, height) diff --git a/av/video/frame.pyi b/av/video/frame.pyi index 3c8ef6a79..14197f24f 100644 --- a/av/video/frame.pyi +++ b/av/video/frame.pyi @@ -34,7 +34,14 @@ class CudaContext: def device_id(self) -> int: ... @property def primary_ctx(self) -> bool: ... - def __init__(self, device_id: int = 0, primary_ctx: bool = True) -> None: ... + @property + def current_ctx(self) -> bool: ... + def __init__( + self, + device_id: int = 0, + primary_ctx: bool = True, + current_ctx: bool = False, + ) -> None: ... class VideoFrame(Frame): format: VideoFormat @@ -109,4 +116,5 @@ class VideoFrame(Frame): device_id: int | None = None, primary_ctx: bool = True, cuda_context: CudaContext | None = None, + current_ctx: bool = False, ) -> VideoFrame: ... diff --git a/docs/api/video.rst b/docs/api/video.rst index 7e97e38c8..8d349fc1b 100644 --- a/docs/api/video.rst +++ b/docs/api/video.rst @@ -33,6 +33,9 @@ Video Frames .. automodule:: av.video.frame +.. autoclass:: CudaContext + :members: + .. autoclass:: VideoFrame A single video frame. @@ -80,6 +83,7 @@ Conversions .. automethod:: VideoFrame.from_image .. automethod:: VideoFrame.from_ndarray +.. automethod:: VideoFrame.from_dlpack diff --git a/tests/test_dlpack.py b/tests/test_dlpack.py index 2e74f342f..cfff0540c 100644 --- a/tests/test_dlpack.py +++ b/tests/test_dlpack.py @@ -531,6 +531,19 @@ def test_video_frame_from_dlpack_invalid_plane_object_raises_typeerror() -> None VideoFrame.from_dlpack((object(), object()), format="nv12", width=64, height=48) +def test_cuda_context_flags() -> None: + default_ctx = av.video.frame.CudaContext() + assert default_ctx.primary_ctx is True + assert default_ctx.current_ctx is False + + with pytest.raises(ValueError, match="mutually exclusive"): + av.video.frame.CudaContext(primary_ctx=True, current_ctx=True) + + ctx = av.video.frame.CudaContext(primary_ctx=False, current_ctx=True) + assert ctx.primary_ctx is False + assert ctx.current_ctx is True + + def test_video_frame_from_dlpack_cuda_hw_frame_behavior_if_available() -> None: backend = _get_cuda_backend() if backend is None: @@ -621,7 +634,10 @@ def test_video_frame_from_dlpack_cuda_hw_frame_behavior_if_available() -> None: pytest.skip(f"CUDA hwcontext not available in this build/runtime: {e}") -def test_encode_cuda_frame_with_nvenc_if_available() -> None: +@pytest.mark.parametrize( + "use_current_ctx", [False, True], ids=["primary-context", "current-context"] +) +def test_encode_cuda_frame_with_nvenc_if_available(use_current_ctx: bool) -> None: # Issue #2199: a CUDA frame from DLPack should encode on the GPU directly. # Its hw_frames_ctx must propagate to the encoder before avcodec_open2. backend = _get_cuda_backend() @@ -639,7 +655,21 @@ def test_encode_cuda_frame_with_nvenc_if_available() -> None: y = mod.zeros((height, width), dtype=mod.uint8) uv = mod.zeros((height // 2, width // 2, 2), dtype=mod.uint8) - frame = VideoFrame.from_dlpack((y, uv), format="nv12") + if use_current_ctx: + current_ctx = av.video.frame.CudaContext( + device_id=int(y.__dlpack_device__()[1]), + primary_ctx=False, + current_ctx=True, + ) + frame = VideoFrame.from_dlpack( + (y, uv), + format="nv12", + primary_ctx=False, + cuda_context=current_ctx, + current_ctx=True, + ) + else: + frame = VideoFrame.from_dlpack((y, uv), format="nv12") assert frame.format.name == "cuda" assert frame.sw_format is not None and frame.sw_format.name == "nv12"