Add core support for decoding from Python file-like objects (#564)

Co-authored-by: Nicolas Hug <[email protected]>

Author: scotts

.github/workflows/cpp_tests.yaml

@@ -34,12 +34,12 @@ jobs:
           python-version: '3.12'
       - name: Update pip
         run: python -m pip install --upgrade pip
-      - name: Install dependencies
+      - name: Install torch dependencies
         run: |
           python -m pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu
-      - name: Install ffmpeg and pkg-config
+      - name: Install ffmpeg, pkg-config and pybind11
         run: |
-          conda install "ffmpeg=${{ matrix.ffmpeg-version-for-tests }}" pkg-config -c conda-forge
+          conda install "ffmpeg=${{ matrix.ffmpeg-version-for-tests }}" pkg-config pybind11 -c conda-forge
           ffmpeg -version
       - name: Build and run C++ tests
         run: |

.github/workflows/docs.yaml

@@ -38,6 +38,7 @@ jobs:
       test-infra-repository: pytorch/test-infra
       test-infra-ref: main
       build-matrix: ${{ needs.generate-matrix.outputs.matrix }}
+      pre-script: packaging/pre_build_script.sh
       post-script: packaging/post_build_script.sh
       smoke-test-script: packaging/fake_smoke_test.py
       package-name: torchcodec

.github/workflows/lint.yaml

@@ -63,7 +63,7 @@ jobs:
       - name: Install dependencies and FFmpeg
         run: |
           python -m pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu
-          conda install "ffmpeg=7.0.1" pkg-config -c conda-forge
+          conda install "ffmpeg=7.0.1" pkg-config pybind11 -c conda-forge
           ffmpeg -version
       - name: Build and install torchcodec
         run: |

.github/workflows/linux_cuda_wheel.yaml

@@ -48,6 +48,7 @@ jobs:
       test-infra-repository: pytorch/test-infra
       test-infra-ref: main
       build-matrix: ${{ needs.generate-matrix.outputs.matrix }}
+      pre-script: packaging/pre_build_script.sh
       post-script: packaging/post_build_script.sh
       smoke-test-script: packaging/fake_smoke_test.py
       package-name: torchcodec

.github/workflows/linux_wheel.yaml

@@ -49,6 +49,7 @@ jobs:
       test-infra-repository: pytorch/test-infra
       test-infra-ref: main
       build-matrix: ${{ needs.generate-matrix.outputs.matrix }}
+      pre-script: packaging/pre_build_script.sh
       post-script: packaging/post_build_script.sh
       smoke-test-script: packaging/fake_smoke_test.py
       package-name: torchcodec

.github/workflows/macos_wheel.yaml

@@ -49,6 +49,7 @@ jobs:
       test-infra-repository: pytorch/test-infra
       test-infra-ref: main
       build-matrix: ${{ needs.generate-matrix.outputs.matrix }}
+      pre-script: packaging/pre_build_script.sh
       post-script: packaging/post_build_script.sh
       smoke-test-script: packaging/fake_smoke_test.py
       runner-type: macos-m1-stable

CONTRIBUTING.md

@@ -20,6 +20,7 @@ test locally you will need the following dependencies:
   installation already.
 - cmake
 - pkg-config
+- pybind11
 - FFmpeg
 - PyTorch nightly
 
@@ -29,7 +30,7 @@ Start by installing the **nightly** build of PyTorch following the
 Then, the easiest way to install the rest of the dependencies is to run:
 
 ```bash
-conda install cmake pkg-config ffmpeg -c conda-forge
+conda install cmake pkg-config pbyind11 ffmpeg -c conda-forge
 ```
 
 ### Clone and build

packaging/pre_build_script.sh

@@ -0,0 +1,9 @@
+#!/bin/bash
+
+set -ex
+
+# We need to install pybind11 because we need its CMake helpers in order to
+# compile correctly on Mac. Pybind11 is actually a C++ header-only library,
+# and PyTorch actually has it included. PyTorch, however, does not have the
+# CMake helpers.
+conda install -y pybind11 -c conda-forge

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.
@@ -136,21 +136,27 @@ def copy_extensions_to_source(self):
         This is called by setuptools at the end of .run() during editable installs.
         """
         self.get_finalized_command("build_py")
-        extension = ""
+        extensions = []
         if sys.platform == "linux":
-            extension = "so"
+            extensions = ["so"]
         elif sys.platform == "darwin":
-            extension = "dylib"
+            # Mac has BOTH .dylib and .so as library extensions. Short version
+            # is that a .dylib is a shared library that can be both dynamically
+            # loaded and depended on by other libraries; a .so can only be a
+            # dynamically loaded module. For more, see:
+            #   https://stackoverflow.com/a/2339910
+            extensions = ["dylib", "so"]
         else:
             raise NotImplementedError(
                 "Platforms other than linux/darwin are not supported yet"
             )
 
-        for so_file in self._install_prefix.glob(f"*.{extension}"):
-            assert "libtorchcodec" in so_file.name
-            destination = Path("src/torchcodec/") / so_file.name
-            print(f"Copying {so_file} to {destination}")
-            self.copy_file(so_file, destination, level=self.verbose)
+        for ext in extensions:
+            for lib_file in self._install_prefix.glob(f"*.{ext}"):
+                assert "libtorchcodec" in lib_file.name
+                destination = Path("src/torchcodec/") / lib_file.name
+                print(f"Copying {lib_file} to {destination}")
+                self.copy_file(lib_file, destination, level=self.verbose)
 
 
 NOT_A_LICENSE_VIOLATION_VAR = "I_CONFIRM_THIS_IS_NOT_A_LICENSE_VIOLATION"

src/torchcodec/_internally_replaced_utils.py

@@ -7,11 +7,12 @@
 import importlib
 import sys
 from pathlib import Path
+from types import ModuleType
 
 
 # Copy pasted from torchvision
 # https://github.com/pytorch/vision/blob/947ae1dc71867f28021d5bc0ff3a19c249236e2a/torchvision/_internally_replaced_utils.py#L25
-def _get_extension_path(lib_name):
+def _get_extension_path(lib_name: str) -> str:
     extension_suffixes = []
     if sys.platform == "linux":
         extension_suffixes = importlib.machinery.EXTENSION_SUFFIXES
@@ -31,6 +32,22 @@ def _get_extension_path(lib_name):
     )
     ext_specs = extfinder.find_spec(lib_name)
     if ext_specs is None:
-        raise ImportError
+        raise ImportError(f"No spec found for {lib_name}")
+
+    if ext_specs.origin is None:
+        raise ImportError(f"Existing spec found for {lib_name} does not have an origin")
 
     return ext_specs.origin
+
+
+def _load_pybind11_module(module_name: str, library_path: str) -> ModuleType:
+    spec = importlib.util.spec_from_file_location(
+        module_name,
+        library_path,
+    )
+    if spec is None:
+        raise ImportError(
+            f"Unable to load spec for module {module_name} from path {library_path}"
+        )
+
+    return importlib.util.module_from_spec(spec)