Fix wheel relocation issues (#2777)

* [DEBUG] Check wheel relocation issues

* Call delocate on Mac

* Use yum instead of conda

* Do not copy ffmpeg dylibs

* Linux dry run

* Do not download FFmpeg on Mac

* Install bzip2

* Restore FFmpeg on Windows

* Remove ffmpeg temporarily

* Do not remove dependencies

* Disable FFmpeg temporarily on Linux wheels

* Test relocation in Linux

* Add docstring

* Call relocation script

* Minor error correction

* Import auditwheel only on Linux

* Restore ffmpeg again on Windows

* Return *device

* Try fix Windows

* Fix clang-format

* Start windows patchwork

* Copy all DLLs

* Disable FFmpeg on Windows for now

* Bugfix
Co-authored-by: Francisco Massa <fvsmassa@gmail.com>

packaging/build_wheel.sh

@@ -17,19 +17,20 @@ if [[ "$(uname)" == Darwin || "$OSTYPE" == "msys" ]]; then
     bin_path=$(dirname $python_exec)
     env_path=$(dirname $bin_path)
     if [[ "$(uname)" == Darwin ]]; then
-        # Include LibPNG
-        cp "$env_path/lib/libpng16.dylib" torchvision
-        # Include LibJPEG
-        cp "$env_path/lib/libjpeg.dylib" torchvision
+        # Install delocate to relocate the required binaries
+        pip_install delocate
     else
         cp "$bin_path/Library/bin/libpng16.dll" torchvision
         cp "$bin_path/Library/bin/libjpeg.dll" torchvision
     fi
 else
-    # Include LibPNG
-    cp "/usr/lib64/libpng.so" torchvision
-    # Include LibJPEG
-    cp "/usr/lib64/libjpeg.so" torchvision
+    # Install auditwheel to get some inspection utilities
+    pip_install auditwheel
+
+    # Point to custom libraries
+    export LD_LIBRARY_PATH=$(pwd)/ext_libraries/lib:$LD_LIBRARY_PATH
+    export TORCHVISION_INCLUDE=$(pwd)/ext_libraries/include
+    export TORCHVISION_LIBRARY=$(pwd)/ext_libraries/lib
 fi
 
 download_copy_ffmpeg
@@ -39,3 +40,20 @@ if [[ "$OSTYPE" == "msys" ]]; then
 else
     IS_WHEEL=1 python setup.py bdist_wheel
 fi
+
+
+if [[ "$(uname)" == Darwin ]]; then
+    pushd dist/
+    python_exec="$(which python)"
+    bin_path=$(dirname $python_exec)
+    env_path=$(dirname $bin_path)
+    for whl in *.whl; do
+        DYLD_LIBRARY_PATH="$env_path/lib/:$DYLD_LIBRARY_PATH" delocate-wheel -v $whl
+    done
+else
+    if [[ "$OSTYPE" == "msys" ]]; then
+        "$script_dir/windows/internal/vc_env_helper.bat" python $script_dir/wheel/relocate.py
+    else
+        LD_LIBRARY_PATH="/usr/local/lib:$LD_LIBRARY_PATH" python $script_dir/wheel/relocate.py
+    fi
+fi

packaging/pkg_helpers.bash

@@ -184,8 +184,8 @@ setup_wheel_python() {
     # Install libpng from Anaconda (defaults)
     conda install libpng jpeg -y
   else
-    # Install native CentOS libPNG
-    yum install -y libpng-devel libjpeg-turbo-devel
+    # Install native CentOS libJPEG, LAME, freetype and GnuTLS
+    yum install -y libjpeg-turbo-devel lame freetype gnutls
     case "$PYTHON_VERSION" in
       2.7)
         if [[ -n "$UNICODE_ABI" ]]; then
@@ -203,7 +203,13 @@ setup_wheel_python() {
         exit 1
         ;;
     esac
-    export PATH="/opt/python/$python_abi/bin:$PATH"
+    # Download all the dependencies required to compile image and video_reader
+    # extensions
+
+    mkdir -p ext_libraries
+    pushd ext_libraries
+    popd
+    export PATH="/opt/python/$python_abi/bin:$(pwd)/ext_libraries/bin:$PATH"
   fi
 }
 
@@ -353,10 +359,8 @@ setup_junit_results_folder() {
 
 
 download_copy_ffmpeg() {
-  mkdir ffmpeg_tmp
-  cd ffmpeg_tmp
   if [[ "$OSTYPE" == "msys" ]]; then
-    # conda install -yq ffmpeg -c pytorch
+    # conda install -yq ffmpeg=4.2 -c pytorch
     # curl -L -q https://anaconda.org/pytorch/ffmpeg/4.3/download/win-64/ffmpeg-4.3-ha925a31_0.tar.bz2 --output ffmpeg-4.3-ha925a31_0.tar.bz2
     # bzip2 --decompress --stdout ffmpeg-4.3-ha925a31_0.tar.bz2 | tar -x --file=-
     # cp Library/bin/*.dll ../torchvision
@@ -365,24 +369,15 @@ download_copy_ffmpeg() {
     if [[ "$(uname)" == Darwin ]]; then
       conda install -yq ffmpeg=4.2 -c pytorch
       conda install -yq wget
-      wget -q https://anaconda.org/pytorch/ffmpeg/4.2/download/osx-64/ffmpeg-4.2-h0a44026_0.tar.bz2
-      tar -xjvf ffmpeg-4.2-h0a44026_0.tar.bz2
-      for f in lib/*.dylib; do
-        if [[ $f =~ ([a-z])+\.dylib ]]; then
-          cp $f ../torchvision
-        fi
-      done
     else
-      wget -q https://anaconda.org/pytorch/ffmpeg/4.2/download/linux-64/ffmpeg-4.2-hf484d3e_0.tar.bz2
-      tar -xjvf ffmpeg-4.2-hf484d3e_0.tar.bz2
-      cp lib/*.so ../torchvision
-      cp -r lib/* /usr/lib
-      cp -r bin/* /usr/bin
-      cp -r include/* /usr/include
-      ldconfig
-      which ffmpeg
+      # pushd ext_libraries
+      # wget -q https://anaconda.org/pytorch/ffmpeg/4.2/download/linux-64/ffmpeg-4.2-hf484d3e_0.tar.bz2
+      # tar -xjvf ffmpeg-4.2-hf484d3e_0.tar.bz2
+      # rm -rf ffmpeg-4.2-hf484d3e_0.tar.bz2
+      # ldconfig
+      # which ffmpeg
+      # popd
+      echo "FFmpeg is disabled currently on Linux"
     fi
   fi
-  cd ..
-  rm -rf ffmpeg_tmp
 }

packaging/wheel/relocate.py

+# -*- coding: utf-8 -*-
+
+"""Helper script to package wheels and relocate binaries."""
+
+# Standard library imports
+import os
+import io
+import sys
+import glob
+import shutil
+import zipfile
+import hashlib
+import platform
+import subprocess
+import os.path as osp
+from base64 import urlsafe_b64encode
+
+# Third party imports
+if sys.platform == 'linux':
+    from auditwheel.lddtree import lddtree
+from wheel.bdist_wheel import get_abi_tag
+
+
+WHITELIST = {
+    'libgcc_s.so.1', 'libstdc++.so.6', 'libm.so.6',
+    'libdl.so.2', 'librt.so.1', 'libc.so.6',
+    'libnsl.so.1', 'libutil.so.1', 'libpthread.so.0',
+    'libresolv.so.2', 'libX11.so.6', 'libXext.so.6',
+    'libXrender.so.1', 'libICE.so.6', 'libSM.so.6',
+    'libGL.so.1', 'libgobject-2.0.so.0', 'libgthread-2.0.so.0',
+    'libglib-2.0.so.0', 'ld-linux-x86-64.so.2', 'ld-2.17.so'
+}
+
+WINDOWS_WHITELIST = {
+    'MSVCP140.dll', 'KERNEL32.dll',
+    'VCRUNTIME140_1.dll', 'VCRUNTIME140.dll',
+    'api-ms-win-crt-heap-l1-1-0.dll',
+    'api-ms-win-crt-runtime-l1-1-0.dll',
+    'api-ms-win-crt-stdio-l1-1-0.dll',
+    'api-ms-win-crt-filesystem-l1-1-0.dll',
+    'api-ms-win-crt-string-l1-1-0.dll',
+    'api-ms-win-crt-environment-l1-1-0.dll',
+    'api-ms-win-crt-math-l1-1-0.dll',
+    'api-ms-win-crt-convert-l1-1-0.dll'
+}
+
+
+HERE = osp.dirname(osp.abspath(__file__))
+PACKAGE_ROOT = osp.dirname(osp.dirname(HERE))
+PLATFORM_ARCH = platform.machine()
+PYTHON_VERSION = sys.version_info
+
+
+def read_chunks(file, size=io.DEFAULT_BUFFER_SIZE):
+    """Yield pieces of data from a file-like object until EOF."""
+    while True:
+        chunk = file.read(size)
+        if not chunk:
+            break
+        yield chunk
+
+
+def rehash(path, blocksize=1 << 20):
+    """Return (hash, length) for path using hashlib.sha256()"""
+    h = hashlib.sha256()
+    length = 0
+    with open(path, 'rb') as f:
+        for block in read_chunks(f, size=blocksize):
+            length += len(block)
+            h.update(block)
+    digest = 'sha256=' + urlsafe_b64encode(
+        h.digest()
+    ).decode('latin1').rstrip('=')
+    # unicode/str python2 issues
+    return (digest, str(length))  # type: ignore
+
+
+def unzip_file(file, dest):
+    """Decompress zip `file` into directory `dest`."""
+    with zipfile.ZipFile(file, 'r') as zip_ref:
+        zip_ref.extractall(dest)
+
+
+def is_program_installed(basename):
+    """
+    Return program absolute path if installed in PATH.
+    Otherwise, return None
+    On macOS systems, a .app is considered installed if
+    it exists.
+    """
+    if (sys.platform == 'darwin' and basename.endswith('.app') and
+            osp.exists(basename)):
+        return basename
+
+    for path in os.environ["PATH"].split(os.pathsep):
+        abspath = osp.join(path, basename)
+        if osp.isfile(abspath):
+            return abspath
+
+
+def find_program(basename):
+    """
+    Find program in PATH and return absolute path
+    Try adding .exe or .bat to basename on Windows platforms
+    (return None if not found)
+    """
+    names = [basename]
+    if os.name == 'nt':
+        # Windows platforms
+        extensions = ('.exe', '.bat', '.cmd', '.dll')
+        if not basename.endswith(extensions):
+            names = [basename + ext for ext in extensions] + [basename]
+    for name in names:
+        path = is_program_installed(name)
+        if path:
+            return path
+
+
+def patch_new_path(library_path, new_dir):
+    library = osp.basename(library_path)
+    name, *rest = library.split('.')
+    rest = '.'.join(rest)
+    hash_id = hashlib.sha256(library_path.encode('utf-8')).hexdigest()[:8]
+    new_name = '.'.join([name, hash_id, rest])
+    return osp.join(new_dir, new_name)
+
+
+def find_dll_dependencies(dumpbin, binary):
+    out = subprocess.run([dumpbin, "/dependents", binary],
+                         stdout=subprocess.PIPE)
+    out = out.stdout.strip().decode('utf-8')
+    start_index = out.find('dependencies:') + len('dependencies:')
+    end_index = out.find('Summary')
+    dlls = out[start_index:end_index].strip()
+    dlls = dlls.split(os.linesep)
+    dlls = [dll.strip() for dll in dlls]
+    return dlls
+
+
+def relocate_elf_library(patchelf, output_dir, output_library, binary):
+    """
+    Relocate an ELF shared library to be packaged on a wheel.
+
+    Given a shared library, find the transitive closure of its dependencies,
+    rename and copy them into the wheel while updating their respective rpaths.
+    """
+
+    print('Relocating {0}'.format(binary))
+    binary_path = osp.join(output_library, binary)
+
+    ld_tree = lddtree(binary_path)
+    tree_libs = ld_tree['libs']
+
+    binary_queue = [(n, binary) for n in ld_tree['needed']]
+    binary_paths = {binary: binary_path}
+    binary_dependencies = {}
+
+    while binary_queue != []:
+        library, parent = binary_queue.pop(0)
+        library_info = tree_libs[library]
+        print(library)
+
+        if library_info['path'] is None:
+            print('Omitting {0}'.format(library))
+            continue
+
+        if library in WHITELIST:
+            # Omit glibc/gcc/system libraries
+            print('Omitting {0}'.format(library))
+            continue
+
+        parent_dependencies = binary_dependencies.get(parent, [])
+        parent_dependencies.append(library)
+        binary_dependencies[parent] = parent_dependencies
+
+        if library in binary_paths:
+            continue
+
+        binary_paths[library] = library_info['path']
+        binary_queue += [(n, library) for n in library_info['needed']]
+
+    print('Copying dependencies to wheel directory')
+    new_libraries_path = osp.join(output_dir, 'torchvision.libs')
+    os.makedirs(new_libraries_path)
+
+    new_names = {binary: binary_path}
+
+    for library in binary_paths:
+        if library != binary:
+            library_path = binary_paths[library]
+            new_library_path = patch_new_path(library_path, new_libraries_path)
+            print('{0} -> {1}'.format(library, new_library_path))
+            shutil.copyfile(library_path, new_library_path)
+            new_names[library] = new_library_path
+
+    print('Updating dependency names by new files')
+    for library in binary_paths:
+        if library != binary:
+            if library not in binary_dependencies:
+                continue
+            library_dependencies = binary_dependencies[library]
+            new_library_name = new_names[library]
+            for dep in library_dependencies:
+                new_dep = osp.basename(new_names[dep])
+                print('{0}: {1} -> {2}'.format(library, dep, new_dep))
+                subprocess.check_output(
+                    [
+                        patchelf,
+                        '--replace-needed',
+                        dep,
+                        new_dep,
+                        new_library_name
+                    ],
+                    cwd=new_libraries_path)
+
+            print('Updating library rpath')
+            subprocess.check_output(
+                [
+                    patchelf,
+                    '--set-rpath',
+                    "$ORIGIN",
+                    new_library_name
+                ],
+                cwd=new_libraries_path)
+
+            subprocess.check_output(
+                [
+                    patchelf,
+                    '--print-rpath',
+                    new_library_name
+                ],
+                cwd=new_libraries_path)
+
+    print("Update library dependencies")
+    library_dependencies = binary_dependencies[binary]
+    for dep in library_dependencies:
+        new_dep = osp.basename(new_names[dep])
+        print('{0}: {1} -> {2}'.format(binary, dep, new_dep))
+        subprocess.check_output(
+            [
+                patchelf,
+                '--replace-needed',
+                dep,
+                new_dep,
+                binary
+            ],
+            cwd=output_library)
+
+    print('Update library rpath')
+    subprocess.check_output(
+        [
+            patchelf,
+            '--set-rpath',
+            "$ORIGIN:$ORIGIN/../torchvision.libs",
+            binary_path
+        ],
+        cwd=output_library
+    )
+
+
+def relocate_dll_library(dumpbin, output_dir, output_library, binary):
+    print('Relocating {0}'.format(binary))
+    binary_path = osp.join(output_library, binary)
+
+    library_dlls = find_dll_dependencies(dumpbin, binary_path)
+    binary_queue = [(dll, binary) for dll in library_dlls]
+    binary_paths = {binary: binary_path}
+    binary_dependencies = {}
+
+    while binary_queue != []:
+        library, parent = binary_queue.pop(0)
+        if library in WINDOWS_WHITELIST or library.startswith('api-ms-win'):
+            print('Omitting {0}'.format(library))
+            continue
+
+        library_path = find_program(library)
+        if library_path is None:
+            print('{0} not found'.format(library))
+            continue
+
+        if osp.basename(osp.dirname(library_path)) == 'system32':
+            continue
+
+        print('{0}: {1}'.format(library, library_path))
+        parent_dependencies = binary_dependencies.get(parent, [])
+        parent_dependencies.append(library)
+        binary_dependencies[parent] = parent_dependencies
+
+        if library in binary_paths:
+            continue
+
+        binary_paths[library] = library_path
+        downstream_dlls = find_dll_dependencies(dumpbin, library_path)
+        binary_queue += [(n, library) for n in downstream_dlls]
+
+    print('Copying dependencies to wheel directory')
+    package_dir = osp.join(output_dir, 'torchvision')
+    for library in binary_paths:
+        if library != binary:
+            library_path = binary_paths[library]
+            new_library_path = osp.join(package_dir, library)
+            print('{0} -> {1}'.format(library, new_library_path))
+            shutil.copyfile(library_path, new_library_path)
+
+
+def compress_wheel(output_dir, wheel, wheel_dir, wheel_name):
+    """Create RECORD file and compress wheel distribution."""
+    print('Update RECORD file in wheel')
+    dist_info = glob.glob(osp.join(output_dir, '*.dist-info'))[0]
+    record_file = osp.join(dist_info, 'RECORD')
+
+    with open(record_file, 'w') as f:
+        for root, _, files in os.walk(output_dir):
+            for this_file in files:
+                full_file = osp.join(root, this_file)
+                rel_file = osp.relpath(full_file, output_dir)
+                if full_file == record_file:
+                    f.write('{0},,\n'.format(rel_file))
+                else:
+                    digest, size = rehash(full_file)
+                    f.write('{0},{1},{2}\n'.format(rel_file, digest, size))
+
+    print('Compressing wheel')
+    base_wheel_name = osp.join(wheel_dir, wheel_name)
+    shutil.make_archive(base_wheel_name, 'zip', output_dir)
+    os.remove(wheel)
+    shutil.move('{0}.zip'.format(base_wheel_name), wheel)
+    shutil.rmtree(output_dir)
+
+
+def patch_linux():
+    # Get patchelf location
+    patchelf = find_program('patchelf')
+    if patchelf is None:
+        raise FileNotFoundError('Patchelf was not found in the system, please'
+                                ' make sure that is available on the PATH.')
+
+    # Find wheel
+    print('Finding wheels...')
+    wheels = glob.glob(osp.join(PACKAGE_ROOT, 'dist', '*.whl'))
+    output_dir = osp.join(PACKAGE_ROOT, 'dist', '.wheel-process')
+
+    image_binary = 'image.so'
+    video_binary = 'video_reader.so'
+    torchvision_binaries = [image_binary, video_binary]
+    for wheel in wheels:
+        if osp.exists(output_dir):
+            shutil.rmtree(output_dir)
+
+        os.makedirs(output_dir)
+
+        print('Unzipping wheel...')
+        wheel_file = osp.basename(wheel)
+        wheel_dir = osp.dirname(wheel)
+        print('{0}'.format(wheel_file))
+        wheel_name, _ = osp.splitext(wheel_file)
+        unzip_file(wheel, output_dir)
+
+        print('Finding ELF dependencies...')
+        output_library = osp.join(output_dir, 'torchvision')
+        for binary in torchvision_binaries:
+            if osp.exists(osp.join(output_library, binary)):
+                relocate_elf_library(
+                    patchelf, output_dir, output_library, binary)
+
+        compress_wheel(output_dir, wheel, wheel_dir, wheel_name)
+
+
+def patch_win():
+    # Get dumpbin location
+    dumpbin = find_program('dumpbin')
+    if dumpbin is None:
+        raise FileNotFoundError('Dumpbin was not found in the system, please'
+                                ' make sure that is available on the PATH.')
+
+    # Find wheel
+    print('Finding wheels...')
+    wheels = glob.glob(osp.join(PACKAGE_ROOT, 'dist', '*.whl'))
+    output_dir = osp.join(PACKAGE_ROOT, 'dist', '.wheel-process')
+
+    image_binary = 'image.pyd'
+    video_binary = 'video_reader.pyd'
+    torchvision_binaries = [image_binary, video_binary]
+    for wheel in wheels:
+        if osp.exists(output_dir):
+            shutil.rmtree(output_dir)
+
+        os.makedirs(output_dir)
+
+        print('Unzipping wheel...')
+        wheel_file = osp.basename(wheel)
+        wheel_dir = osp.dirname(wheel)
+        print('{0}'.format(wheel_file))
+        wheel_name, _ = osp.splitext(wheel_file)
+        unzip_file(wheel, output_dir)
+
+        print('Finding DLL/PE dependencies...')
+        output_library = osp.join(output_dir, 'torchvision')
+        for binary in torchvision_binaries:
+            if osp.exists(osp.join(output_library, binary)):
+                relocate_dll_library(
+                    dumpbin, output_dir, output_library, binary)
+
+        compress_wheel(output_dir, wheel, wheel_dir, wheel_name)
+
+
+if __name__ == '__main__':
+    if sys.platform == 'linux':
+        patch_linux()
+    elif sys.platform == 'win32':
+        patch_win()

torchvision/csrc/cpu/video/Video.cpp

@@ -49,8 +49,8 @@ size_t fillAudioTensor(DecoderOutputMessage& msgs, torch::Tensor& audioFrame) {
   return fillTensorList<float>(msgs, audioFrame);
 }
 
-std::pair<std::string, ffmpeg::MediaType> const* _parse_type(
-    const std::string& stream_string) {
+std::array<std::pair<std::string, ffmpeg::MediaType>, 4>::const_iterator
+_parse_type(const std::string& stream_string) {
   static const std::array<std::pair<std::string, MediaType>, 4> types = {{
       {"video", TYPE_VIDEO},
       {"audio", TYPE_AUDIO},