# ------------------------------------------------------------------
|
# Copyright (c) 2023 PyInstaller Development Team.
|
#
|
# This file is distributed under the terms of the GNU General Public
|
# License (version 2.0 or later).
|
#
|
# The full license is available in LICENSE, distributed with
|
# this software.
|
#
|
# SPDX-License-Identifier: GPL-2.0-or-later
|
# ------------------------------------------------------------------
|
|
import os
|
import re
|
|
from PyInstaller import compat
|
from PyInstaller.utils.hooks import (
|
logger,
|
is_module_satisfies,
|
)
|
|
|
# Helper for collecting shared libraries from NVIDIA CUDA packages on linux.
|
def collect_nvidia_cuda_binaries(hook_file):
|
# Find the module underlying this nvidia.something hook; i.e., change ``/path/to/hook-nvidia.something.py`` to
|
# ``nvidia.something``.
|
hook_name, hook_ext = os.path.splitext(os.path.basename(hook_file))
|
assert hook_ext.startswith('.py')
|
assert hook_name.startswith('hook-')
|
module_name = hook_name[5:]
|
|
# `search_patterns` was added to `collect_dynamic_libs` in PyInstaller 5.8, so that is the minimum required version.
|
binaries = []
|
if is_module_satisfies('PyInstaller >= 5.8'):
|
from PyInstaller.utils.hooks import collect_dynamic_libs, PY_DYLIB_PATTERNS
|
binaries = collect_dynamic_libs(
|
module_name,
|
# Collect fully-versioned .so files (not included in default search patterns).
|
search_patterns=PY_DYLIB_PATTERNS + ["lib*.so.*"],
|
)
|
else:
|
logger.warning("hook-%s: this hook requires PyInstaller >= 5.8!", module_name)
|
|
return binaries
|
|
|
# Helper to turn list of requirements (e.g., ['nvidia-cublas-cu12', 'nvidia-nccl-cu12', 'nvidia-cudnn-cu12']) into
|
# list of corresponding nvidia.* module names (e.g., ['nvidia.cublas', 'nvidia.nccl', 'nvidia-cudnn']), while ignoring
|
# unrecognized requirements. Intended for use in hooks for frameworks, such as `torch` and `tensorflow`.
|
def infer_hiddenimports_from_requirements(requirements):
|
# All nvidia-* packages install to nvidia top-level package, so we cannot query top-level module via
|
# metadata. Instead, we manually translate them from dist name to package name.
|
_PATTERN = r'^nvidia-(?P<subpackage>.+)-cu[\d]+$'
|
nvidia_hiddenimports = []
|
|
for req in requirements:
|
m = re.match(_PATTERN, req)
|
if m is not None:
|
# Convert
|
package_name = "nvidia." + m.group('subpackage').replace('-', '_')
|
nvidia_hiddenimports.append(package_name)
|
|
return nvidia_hiddenimports
|
|
|
def create_symlink_suppression_patterns(hook_file):
|
hook_name, hook_ext = os.path.splitext(os.path.basename(hook_file))
|
assert hook_ext.startswith('.py')
|
assert hook_name.startswith('hook-')
|
module_name = hook_name[5:]
|
|
# Applicable only to Linux
|
if not compat.is_linux:
|
return []
|
|
# Pattern: **/{module_dir}/lib/lib*.so*
|
return [os.path.join('**', *module_name.split('.'), 'lib', 'lib*.so*')]
|
