Remove pin memory from binding (#243)

facebookresearch · Oct 4, 2024 · 5dffc1b · 5dffc1b
1 parent 7efbedc
commit 5dffc1b
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 44 deletions.
diff --git a/src/binding/core/conversion.cpp b/src/binding/core/conversion.cpp
@@ -17,18 +17,14 @@
 
 #include "gil.h"
 
-#include <cstring>
-
 namespace nb = nanobind;
 
 namespace spdl::core {
 namespace {
 template <MediaType media_type>
-CPUBufferPtr convert(
-    const FFmpegFramesPtr<media_type>&& frames,
-    bool pin_memory) {
+CPUBufferPtr convert(const FFmpegFramesPtr<media_type>&& frames) {
   RELEASE_GIL();
-  return convert_frames(frames.get(), pin_memory);
+  return convert_frames(frames.get());
 }
 
 template <MediaType media_type>
@@ -42,48 +38,22 @@ std::vector<const spdl::core::FFmpegFrames<media_type>*> _ref(
 }
 
 template <MediaType media_type>
-CPUBufferPtr batch_convert(
-    std::vector<FFmpegFramesPtr<media_type>>&& frames,
-    bool pin_memory) {
+CPUBufferPtr batch_convert(std::vector<FFmpegFramesPtr<media_type>>&& frames) {
   RELEASE_GIL();
-  return convert_frames(_ref(frames), pin_memory);
+  return convert_frames(_ref(frames));
 }
 } // namespace
 
 void register_conversion(nb::module_& m) {
   ////////////////////////////////////////////////////////////////////////////////
   // Frame conversion
   ////////////////////////////////////////////////////////////////////////////////
-  m.def(
-      "convert_frames",
-      &convert<MediaType::Audio>,
-      nb::arg("frames"),
-      nb::arg("pin_memory") = false);
-  m.def(
-      "convert_frames",
-      &convert<MediaType::Video>,
-      nb::arg("frames"),
-      nb::arg("pin_memory") = false);
-  m.def(
-      "convert_frames",
-      &convert<MediaType::Image>,
-      nb::arg("frames"),
-      nb::arg("pin_memory") = false);
+  m.def("convert_frames", &convert<MediaType::Audio>, nb::arg("frames"));
+  m.def("convert_frames", &convert<MediaType::Video>, nb::arg("frames"));
+  m.def("convert_frames", &convert<MediaType::Image>, nb::arg("frames"));
 
-  m.def(
-      "convert_frames",
-      &batch_convert<MediaType::Audio>,
-      nb::arg("frames"),
-      nb::arg("pin_memory") = false);
-  m.def(
-      "convert_frames",
-      &batch_convert<MediaType::Video>,
-      nb::arg("frames"),
-      nb::arg("pin_memory") = false);
-  m.def(
-      "convert_frames",
-      &batch_convert<MediaType::Image>,
-      nb::arg("frames"),
-      nb::arg("pin_memory") = false);
+  m.def("convert_frames", &batch_convert<MediaType::Audio>, nb::arg("frames"));
+  m.def("convert_frames", &batch_convert<MediaType::Video>, nb::arg("frames"));
+  m.def("convert_frames", &batch_convert<MediaType::Image>, nb::arg("frames"));
 }
 } // namespace spdl::core
diff --git a/src/spdl/io/_composite.py b/src/spdl/io/_composite.py
@@ -420,7 +420,6 @@ async def async_load_image_batch(
     decode_config: DecodeConfig | None = None,
     filter_desc: str | None = _FILTER_DESC_DEFAULT,
     device_config: None = None,
-    pin_memory: bool = False,
     strict: bool = True,
     **kwargs,
 ) -> CPUBuffer: ...
@@ -437,7 +436,6 @@ async def async_load_image_batch(
     decode_config: DecodeConfig | None = None,
     filter_desc: str | None = _FILTER_DESC_DEFAULT,
     device_config: CUDAConfig,
-    pin_memory: bool = False,
     strict: bool = True,
     **kwargs,
 ) -> CUDABuffer: ...
@@ -453,7 +451,6 @@ async def async_load_image_batch(
     decode_config=None,
     filter_desc=_FILTER_DESC_DEFAULT,
     device_config=None,
-    pin_memory=False,
     strict=True,
     **kwargs,
 ):
@@ -544,6 +541,13 @@ async def async_load_image_batch(
     if not srcs:
         raise ValueError("`srcs` must not be empty.")
 
+    if "pin_memory" in kwargs:
+        warnings.warn(
+            "`pin_memory` argument has been removed. Use `storage` instead.",
+            stacklevel=2,
+        )
+        kwargs.pop("pin_memory")
+
     if device_config is None and "cuda_config" in kwargs:
         warnings.warn(
             "The `cuda_config` argument has ben renamed to `device_config`.",
@@ -582,7 +586,7 @@ async def async_load_image_batch(
     if not frames:
         raise RuntimeError("Failed to load all the images.")
 
-    buffer = await _core.async_convert_frames(frames, pin_memory=pin_memory)
+    buffer = await _core.async_convert_frames(frames)
 
     if device_config is not None:
         buffer = await _core.async_transfer_buffer(buffer, device_config=device_config)

diff --git a/src/spdl/io/_core.py b/src/spdl/io/_core.py
@@ -552,6 +552,12 @@ def convert_frames(
             - ``W``: width
             - ``H``: height
     """
+    if "pin_memory" in kwargs:
+        warnings.warn(
+            "`pin_memory` argument has been removed. Use `storage` instead.",
+            stacklevel=2,
+        )
+        kwargs.pop("pin_memory")
     return _libspdl.convert_frames(frames, **kwargs)
 
 
@@ -567,6 +573,12 @@ async def async_convert_frames(
     **kwargs,
 ) -> CPUBuffer:
     """Async version of :py:func:`~spdl.io.convert_frames`."""
+    if "pin_memory" in kwargs:
+        warnings.warn(
+            "`pin_memory` argument has been removed. Use `storage` instead.",
+            stacklevel=2,
+        )
+        kwargs.pop("pin_memory")
     return await run_async(convert_frames, frames, **kwargs)