Generalize pybind_ops module name (#759)

Author: scotts

src/torchcodec/_core/CMakeLists.txt

@@ -151,6 +151,15 @@ function(make_torchcodec_libraries
         PUBLIC
       "-fvisibility=hidden"
     )
+    # The value we use here must match the value we return from
+    # _get_pybind_ops_module_name() on the Python side. If the values do not
+    # match, then we will be unable to import the C++ shared library as a
+    # Python module at runtime.
+    target_compile_definitions(
+        ${pybind_ops_library_name}
+        PRIVATE
+        PYBIND_OPS_MODULE_NAME=decoder_core_pybind_ops
+    )
     # If we don't make sure this flag is set, we run into segfauls at import
     # time on Mac. See:
     #    https://github.com/pybind/pybind11/issues/3907#issuecomment-1170412764

src/torchcodec/_core/ops.py

@@ -15,6 +15,7 @@
 
 from torchcodec._internally_replaced_utils import (  # @manual=//pytorch/torchcodec/src:internally_replaced_utils
     _get_extension_path,
+    _get_pybind_ops_module_name,
     _load_pybind11_module,
 )
 
@@ -43,8 +44,8 @@ def load_torchcodec_shared_libraries():
     #      libraries do not meet those conditions.
 
     exceptions = []
-    pybind_ops_module_name = "decoder_core_pybind_ops"
     for ffmpeg_major_version in (7, 6, 5, 4):
+        pybind_ops_module_name = _get_pybind_ops_module_name(ffmpeg_major_version)
         decoder_library_name = f"libtorchcodec_decoder{ffmpeg_major_version}"
         custom_ops_library_name = f"libtorchcodec_custom_ops{ffmpeg_major_version}"
         pybind_ops_library_name = f"libtorchcodec_pybind_ops{ffmpeg_major_version}"

src/torchcodec/_core/pybind_ops.cpp

@@ -38,7 +38,11 @@ int64_t create_from_file_like(
   return reinterpret_cast<int64_t>(decoder);
 }
 
-PYBIND11_MODULE(decoder_core_pybind_ops, m) {
+#ifndef PYBIND_OPS_MODULE_NAME
+#error PYBIND_OPS_MODULE_NAME must be defined!
+#endif
+
+PYBIND11_MODULE(PYBIND_OPS_MODULE_NAME, m) {
   m.def("create_from_file_like", &create_from_file_like);
 }
 

src/torchcodec/_internally_replaced_utils.py

@@ -51,3 +51,14 @@ def _load_pybind11_module(module_name: str, library_path: str) -> ModuleType:
         )
 
     return importlib.util.module_from_spec(spec)
+
+
+def _get_pybind_ops_module_name(ffmpeg_major_version: str) -> str:
+    # Note that this value must match the value used as PYBIND_OPS_MODULE_NAME
+    # when we compile _core/pybind_ops.cpp. If the values do not match, we will
+    # not be able to import the C++ shared library as a Python module at
+    # runtime.
+    #
+    # The parameter ffmpeg_major_version is unused externally, but used
+    # internally.
+    return "decoder_core_pybind_ops"