Address comments

Author: scotts

setup.py

@@ -68,7 +68,7 @@ def run(self):
         super().run()
 
     def build_extension(self, ext):
-        """Call our CMake build system to build libtorchcodec?.so"""
+        """Call our CMake build system to build libtorchcodec*.so"""
         # Setuptools was designed to build one extension (.so file) at a time,
         # calling this method for each Extension object. We're using a
         # CMake-based build where all our extensions are built together at once.

src/torchcodec/decoders/_core/AVIOFileLikeContext.cpp

@@ -29,12 +29,12 @@ AVIOFileLikeContext::AVIOFileLikeContext(py::object fileLike)
 int AVIOFileLikeContext::read(void* opaque, uint8_t* buf, int buf_size) {
   auto fileLike = static_cast<UniquePyObject*>(opaque);
 
+  // Note that we acquire the GIL outside of the loop. This is likely more
+  // efficient than releasing and acquiring it each loop iteration.
+  py::gil_scoped_acquire gil;
   int num_read = 0;
   while (num_read < buf_size) {
     int request = buf_size - num_read;
-    // TODO: It is maybe more efficient to grab the lock once in the
-    // surrounding scope?
-    py::gil_scoped_acquire gil;
     auto chunk = static_cast<std::string>(
         static_cast<py::bytes>((*fileLike)->attr("read")(request)));
     int chunk_len = static_cast<int>(chunk.length());

src/torchcodec/decoders/_core/AVIOFileLikeContext.h

@@ -33,6 +33,11 @@ class AVIOFileLikeContext : public AVIOContextHolder {
   // and that's, at least, hard. For all of the common pitfalls, see:
   //
   //   https://pybind11.readthedocs.io/en/stable/advanced/misc.html#common-sources-of-global-interpreter-lock-errors
+  //
+  // We maintain a reference to the file-like object because the file-like
+  // object that was created on the Python side must live as long as our
+  // potential use. That is, even if there are no more references to the object
+  // on the Python side, we require that the object is still live.
   struct PyObjectDeleter {
     inline void operator()(py::object* obj) const {
       if (obj) {

src/torchcodec/decoders/_core/CMakeLists.txt

@@ -10,8 +10,7 @@ find_package(Python3 ${PYTHON_VERSION} EXACT COMPONENTS Development)
 function(make_torchcodec_sublibrary
     library_name
     sources
-    dependent_libraries
-    ffmpeg_include_dirs)
+    library_dependencies)
 
     add_library(${library_name} SHARED ${sources})
     set_target_properties(${library_name} PROPERTIES CXX_STANDARD 17)
@@ -20,7 +19,6 @@ function(make_torchcodec_sublibrary
         ./../../../../
         "${TORCH_INSTALL_PREFIX}/include"
         ${Python3_INCLUDE_DIRS}
-        ${ffmpeg_include_dirs}
     )
 
     # Avoid adding the "lib" prefix which we already add explicitly.
@@ -29,14 +27,15 @@ function(make_torchcodec_sublibrary
     target_link_libraries(
         ${library_name}
         PUBLIC
-        ${dependent_libraries}
+        ${library_dependencies}
     )
 endfunction()
 
 function(make_torchcodec_libraries
     ffmpeg_major_version
-    ffmpeg_target
-    ffmpeg_include_dirs)
+    ffmpeg_target)
+
+    # TODO: List each library and its purpose.
 
     # Create libtorchcodec_decoderN.so
     set(decoder_library_name "libtorchcodec_decoder${ffmpeg_major_version}")
@@ -52,14 +51,14 @@ function(make_torchcodec_libraries
         list(APPEND decoder_sources CPUOnlyDevice.cpp)
     endif()
 
-    set(decoder_dependent_libraries
+    set(decoder_library_dependencies
         ${ffmpeg_target}
         ${TORCH_LIBRARIES}
         ${Python3_LIBRARIES}
     )
 
     if(ENABLE_CUDA)
-        list(APPEND decoder_dependent_libraries
+        list(APPEND decoder_library_dependencies
             ${CUDA_nppi_LIBRARY}
             ${CUDA_nppicc_LIBRARY}
         )
@@ -68,8 +67,7 @@ function(make_torchcodec_libraries
     make_torchcodec_sublibrary(
         "${decoder_library_name}"
         "${decoder_sources}"
-        "${decoder_dependent_libraries}"
-        "${ffmpeg_include_dirs}"
+        "${decoder_library_dependencies}"
     )
 
     # Create libtorchcodec_custom_opsN.so
@@ -82,7 +80,6 @@ function(make_torchcodec_libraries
         "${custom_ops_library_name}"
         "${custom_ops_sources}"
         "${decoder_library_name}"
-        "${ffmpeg_include_dirs}"
     )
 
     # Create libtorchcodec_pybind_opsN.so
@@ -95,7 +92,6 @@ function(make_torchcodec_libraries
         "${pybind_ops_library_name}"
         "${pybind_ops_sources}"
         "${decoder_library_name}"
-        "${ffmpeg_include_dirs}"
     )
     # pybind11 quirk, see:
     #   https://pybind11.readthedocs.io/en/stable/faq.html#someclass-declared-with-greater-visibility-than-the-type-of-its-field-someclass-member-wattributes
@@ -137,10 +133,10 @@ if(DEFINED ENV{BUILD_AGAINST_ALL_FFMPEG_FROM_S3})
         ${CMAKE_CURRENT_SOURCE_DIR}/fetch_and_expose_non_gpl_ffmpeg_libs.cmake
     )
 
-    make_torchcodec_libraries(7 ffmpeg7 $ffmpeg7_INCLUDE_DIRs)
-    make_torchcodec_libraries(6 ffmpeg6 $ffmpeg6_INCLUDE_DIRS)
-    make_torchcodec_libraries(4 ffmpeg4 $ffmpeg4_INCLUDE_DIRS)
-    make_torchcodec_libraries(5 ffmpeg5 $ffmpeg5_INCLUDE_DIRS)
+    make_torchcodec_libraries(7 ffmpeg7)
+    make_torchcodec_libraries(6 ffmpeg6)
+    make_torchcodec_libraries(4 ffmpeg4)
+    make_torchcodec_libraries(5 ffmpeg5)
 else()
     message(
         STATUS
@@ -180,7 +176,7 @@ else()
         )
     endif()
 
-    make_torchcodec_libraries(${ffmpeg_major_version} PkgConfig::LIBAV ${LIBAV_INCLUDE_DIRS})
+    make_torchcodec_libraries(${ffmpeg_major_version} PkgConfig::LIBAV)
 
     # Expose these values updwards so that the test compilation does not need
     # to re-figure it out. FIXME: it's not great that we just copy-paste the

src/torchcodec/decoders/_core/PyBindOps.cpp

@@ -31,9 +31,9 @@ int64_t create_from_file_like(
     realSeek = seekModeFromString(seek_mode.value());
   }
 
-  auto contextHolder = std::make_unique<AVIOFileLikeContext>(file_like);
+  auto avioContextHolder = std::make_unique<AVIOFileLikeContext>(file_like);
 
-  VideoDecoder* decoder = new VideoDecoder(std::move(contextHolder), realSeek);
+  VideoDecoder* decoder = new VideoDecoder(std::move(avioContextHolder), realSeek);
   return reinterpret_cast<int64_t>(decoder);
 }
 

src/torchcodec/decoders/_core/VideoDecoderOps.cpp

@@ -129,6 +129,7 @@ at::Tensor create_from_tensor(
     at::Tensor video_tensor,
     std::optional<std::string_view> seek_mode) {
   TORCH_CHECK(video_tensor.is_contiguous(), "video_tensor must be contiguous");
+  TORCH_CHECK(video_tensor.scalar_type() == torch::kUInt8, "video_tensor must be kUInt8");
   void* data = video_tensor.mutable_data_ptr();
   size_t length = video_tensor.numel();
 

src/torchcodec/decoders/_core/ops.py

@@ -21,7 +21,7 @@
 _pybind_ops: Optional[ModuleType] = None
 
 
-def load_torchcodec_extension():
+def load_torchcodec_shared_libraries():
     # Successively try to load libtorchcodec7.so, libtorchcodec6.so,
     # libtorchcodec5.so, and libtorchcodec4.so. Each of these correspond to an
     # ffmpeg major version. This should cover all potential ffmpeg versions
@@ -79,7 +79,7 @@ def load_torchcodec_extension():
     )
 
 
-load_torchcodec_extension()
+load_torchcodec_shared_libraries()
 
 
 # Note: We use disallow_in_graph because PyTorch does constant propagation of

test/decoders/test_ops.py

@@ -363,7 +363,7 @@ def test_create_decoder(self, create_from, device):
             decoder = create_from_file_like(open(path, mode="rb", buffering=0), "exact")
         elif create_from == "file_like_bufferedio":
             decoder = create_from_file_like(
-                open(path, mode="rb", buffering=-4096), "exact"
+                open(path, mode="rb", buffering=4096), "exact"
             )
         else:
             raise ValueError("Oops, double check the parametrization of this test!")