From ecade848e043f324dee535072b053c7db91ae081 Mon Sep 17 00:00:00 2001 From: Rok Mandeljc Date: Tue, 12 Dec 2023 12:56:54 +0100 Subject: [PATCH 01/35] tests: create separate test file for torch and its libraries Create a separate `test_pytorch.py` test for `torch` and its associated libraries. Move existing tests to the new file. --- .../tests/test_libraries.py | 45 -------------- .../tests/test_pytorch.py | 59 +++++++++++++++++++ 2 files changed, 59 insertions(+), 45 deletions(-) create mode 100644 src/_pyinstaller_hooks_contrib/tests/test_pytorch.py diff --git a/src/_pyinstaller_hooks_contrib/tests/test_libraries.py b/src/_pyinstaller_hooks_contrib/tests/test_libraries.py index 7841dbf6..b671b9e6 100644 --- a/src/_pyinstaller_hooks_contrib/tests/test_libraries.py +++ b/src/_pyinstaller_hooks_contrib/tests/test_libraries.py @@ -550,51 +550,6 @@ class User(pydantic.BaseModel): """) -def torch_onedir_only(test): - - def wrapped(pyi_builder): - if pyi_builder._mode != 'onedir': - pytest.skip('PyTorch tests support only onedir mode ' - 'due to potential distribution size.') - test(pyi_builder) - - return wrapped - - -@importorskip('torch') -@torch_onedir_only -def test_torch(pyi_builder): - pyi_builder.test_source(""" - import torch - - torch.rand((10, 10)) * torch.rand((10, 10)) - """) - - -@importorskip('torchvision') -@torch_onedir_only -def test_torchvision_nms(pyi_builder): - pyi_builder.test_source(""" - import torch - import torchvision - # boxes: Nx4 tensor (x1, y1, x2, y2) - boxes = torch.tensor([ - [0.0, 0.0, 1.0, 1.0], - [0.45, 0.0, 1.0, 1.0], - ]) - # scores: Nx1 tensor - scores = torch.tensor([ - 1.0, - 1.1 - ]) - keep = torchvision.ops.nms(boxes, scores, 0.5) - # The boxes have IoU of 0.55, and the second one has a slightly - # higher score, so we expect it to be kept while the first one - # is discarded. - assert keep == 1 - """) - - @requires('google-api-python-client >= 2.0.0') def test_googleapiclient(pyi_builder): pyi_builder.test_source(""" diff --git a/src/_pyinstaller_hooks_contrib/tests/test_pytorch.py b/src/_pyinstaller_hooks_contrib/tests/test_pytorch.py new file mode 100644 index 00000000..26474572 --- /dev/null +++ b/src/_pyinstaller_hooks_contrib/tests/test_pytorch.py @@ -0,0 +1,59 @@ +# ------------------------------------------------------------------ +# Copyright (c) 2020 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.GPL.txt, distributed with +# this software. +# +# SPDX-License-Identifier: GPL-2.0-or-later +# ------------------------------------------------------------------ + +import pytest + +from PyInstaller.utils.tests import importorskip + + +def torch_onedir_only(test): + + def wrapped(pyi_builder): + if pyi_builder._mode != 'onedir': + pytest.skip('PyTorch tests support only onedir mode due to potential distribution size.') + test(pyi_builder) + + return wrapped + + +@importorskip('torch') +@torch_onedir_only +def test_torch(pyi_builder): + pyi_builder.test_source(""" + import torch + + torch.rand((10, 10)) * torch.rand((10, 10)) + """) + + +@importorskip('torchvision') +@torch_onedir_only +def test_torchvision_nms(pyi_builder): + pyi_builder.test_source(""" + import torch + import torchvision + # boxes: Nx4 tensor (x1, y1, x2, y2) + boxes = torch.tensor([ + [0.0, 0.0, 1.0, 1.0], + [0.45, 0.0, 1.0, 1.0], + ]) + # scores: Nx1 tensor + scores = torch.tensor([ + 1.0, + 1.1 + ]) + keep = torchvision.ops.nms(boxes, scores, 0.5) + # The boxes have IoU of 0.55, and the second one has a slightly + # higher score, so we expect it to be kept while the first one + # is discarded. + assert keep == 1 + """) From 84bf7f1549e293d347fb61f93bc7898401e9d9af Mon Sep 17 00:00:00 2001 From: Rok Mandeljc Date: Tue, 12 Dec 2023 13:07:23 +0100 Subject: [PATCH 02/35] tests: improve the enforcement of onedir-only tests for torch Explicitly force the `pyi_builder` into onedir-only mode instead of skipping onefile tests. Reduces number of reported tests. --- src/_pyinstaller_hooks_contrib/tests/test_pytorch.py | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/src/_pyinstaller_hooks_contrib/tests/test_pytorch.py b/src/_pyinstaller_hooks_contrib/tests/test_pytorch.py index 26474572..4f94996d 100644 --- a/src/_pyinstaller_hooks_contrib/tests/test_pytorch.py +++ b/src/_pyinstaller_hooks_contrib/tests/test_pytorch.py @@ -15,14 +15,8 @@ from PyInstaller.utils.tests import importorskip -def torch_onedir_only(test): - - def wrapped(pyi_builder): - if pyi_builder._mode != 'onedir': - pytest.skip('PyTorch tests support only onedir mode due to potential distribution size.') - test(pyi_builder) - - return wrapped +# Run the tests in onedir mode only +torch_onedir_only = pytest.mark.parametrize('pyi_builder', ['onedir'], indirect=True) @importorskip('torch') From e2105a48451b3fafe192665d388073b5a14332b1 Mon Sep 17 00:00:00 2001 From: Rok Mandeljc Date: Tue, 12 Dec 2023 13:08:39 +0100 Subject: [PATCH 03/35] tests: torchvision: add test that uses torchscript Add a test that uses `torchvision.datasets`, `torchvision.transforms`, and torchscript. The latter demonstrates the need for collecting source .py files from `torchvision`. --- .../tests/test_pytorch.py | 46 +++++++++++++++++++ 1 file changed, 46 insertions(+) diff --git a/src/_pyinstaller_hooks_contrib/tests/test_pytorch.py b/src/_pyinstaller_hooks_contrib/tests/test_pytorch.py index 4f94996d..15fb055d 100644 --- a/src/_pyinstaller_hooks_contrib/tests/test_pytorch.py +++ b/src/_pyinstaller_hooks_contrib/tests/test_pytorch.py @@ -51,3 +51,49 @@ def test_torchvision_nms(pyi_builder): # is discarded. assert keep == 1 """) + + +# Advanced tests that uses torchvision.datasets and torchvision.transforms; +# the transforms are combined using torchscript, which requires access to +# transforms' sources. +@importorskip('torchvision') +@torch_onedir_only +def test_torchvision_scripted_transforms(pyi_builder): + pyi_builder.test_source(""" + import torch + import torch.nn + + import torchvision.transforms + import torchvision.datasets + + # Generate one image, and convert it from PIL to tensor + preproc = torchvision.transforms.Compose([ + torchvision.transforms.PILToTensor() + ]) + + dataset = torchvision.datasets.FakeData( + size=1, # 1 image + image_size=(3, 200, 200), + num_classes=2, + transform=preproc, + ) + + assert len(dataset) == 1 + image, category = dataset[0] + + assert image.size() == (3, 200, 200) + assert image.dtype == torch.uint8 + + # Create a composite transform that uses torchscript + transforms = torch.nn.Sequential( + torchvision.transforms.RandomCrop(100), + torchvision.transforms.RandomHorizontalFlip(p=0.3), + ) + scripted_transforms = torch.jit.script(transforms) + + # Transform image + transformed_image = scripted_transforms(image) + + assert transformed_image.size() == (3, 100, 100) + assert transformed_image.dtype == torch.uint8 + """) From f428a7489fcf77ee087ab3398f3d5c82c5d2cc99 Mon Sep 17 00:00:00 2001 From: Rok Mandeljc Date: Tue, 12 Dec 2023 13:44:52 +0100 Subject: [PATCH 04/35] hooks: torchvision: collect source .py files Rename `hook-torchvision.ops.py` to `hook-torchvision.py`, and add `module_collection_mode = 'pyz+py'` to collect source .py files for torch JIT/torchscript. --- news/676.update.rst | 2 ++ .../{hook-torchvision.ops.py => hook-torchvision.py} | 7 +++++-- 2 files changed, 7 insertions(+), 2 deletions(-) create mode 100644 news/676.update.rst rename src/_pyinstaller_hooks_contrib/hooks/stdhooks/{hook-torchvision.ops.py => hook-torchvision.py} (69%) diff --git a/news/676.update.rst b/news/676.update.rst new file mode 100644 index 00000000..144ebfc4 --- /dev/null +++ b/news/676.update.rst @@ -0,0 +1,2 @@ +Update ``torchvision`` hook to collect source .py files for TorchScript/JIT +(requires PyInstaller >= 5.3 to take effect). diff --git a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-torchvision.ops.py b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-torchvision.py similarity index 69% rename from src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-torchvision.ops.py rename to src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-torchvision.py index 31ce0ab3..546ad105 100644 --- a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-torchvision.ops.py +++ b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-torchvision.py @@ -10,6 +10,9 @@ # SPDX-License-Identifier: GPL-2.0-or-later # ------------------------------------------------------------------ -# Functions from torchvision.ops.* modules require torchvision._C -# extension module, which PyInstaller fails to pick up automatically... +# Functions from torchvision.ops.* modules require torchvision._C extension module, which PyInstaller fails to pick up +# automatically due to indirect load. hiddenimports = ['torchvision._C'] + +# Collect source .py files for JIT/torchscript. Requires PyInstaller >= 5.3, no-op in older versions. +module_collection_mode = 'pyz+py' From 5bf7e4fe53bd17dbb655ed015df309d3ff44748f Mon Sep 17 00:00:00 2001 From: Rok Mandeljc Date: Tue, 12 Dec 2023 14:21:29 +0100 Subject: [PATCH 05/35] hooks: torch: explicitly collect versioned .so files When collecting binaries in the PyInstaller >= 6.0 codepath, explicitly collect versioned .so files by adding '*.so.*' to the list of search patterns passed to `collect_dynamic_libs`. Just in case that any of those libs happens to be dynamically loaded at run-time... --- news/676.update.1.rst | 2 ++ .../hooks/stdhooks/hook-torch.py | 8 +++++++- 2 files changed, 9 insertions(+), 1 deletion(-) create mode 100644 news/676.update.1.rst diff --git a/news/676.update.1.rst b/news/676.update.1.rst new file mode 100644 index 00000000..7a64ec32 --- /dev/null +++ b/news/676.update.1.rst @@ -0,0 +1,2 @@ +(Linux) Update ``torch`` hook to explicitly collect versioned .so files +in the new PyInstaller >= 6.0 codepath. diff --git a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-torch.py b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-torch.py index b7ddadc7..b6104400 100644 --- a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-torch.py +++ b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-torch.py @@ -20,6 +20,8 @@ ) if is_module_satisfies("PyInstaller >= 6.0"): + from PyInstaller.utils.hooks import PY_DYLIB_PATTERNS + module_collection_mode = "pyz+py" warn_on_missing_hiddenimports = False @@ -35,8 +37,12 @@ "**/*.cmake", ], ) - binaries = collect_dynamic_libs("torch") hiddenimports = collect_submodules("torch") + binaries = collect_dynamic_libs( + "torch", + # Ensure we pick up fully-versioned .so files as well + search_patterns=PY_DYLIB_PATTERNS + ['*.so.*'], + ) else: datas = [(get_package_paths("torch")[1], "torch")] From c901c07dd4e220d68badb01429be8a2c259a4cdc Mon Sep 17 00:00:00 2001 From: Rok Mandeljc Date: Tue, 12 Dec 2023 15:16:02 +0100 Subject: [PATCH 06/35] hooks: torch: add support for external nvidia.* packages on linux The contemporary PyPI torch wheels for linux use CUDA libraries that are installed via `nvidia-*` packages. Therefore, attempt to convert the `nvidia-*` requirements from the `torch` metadata into hidden imports. This way, we can provide hooks for `nvidia.*` packages that collect the shared libs, in case any of them are dynamically loaded (which currently seems to be the case with some of the shared libraries from `nvidia.cudnn`). --- news/676.update.2.rst | 4 ++ .../hooks/stdhooks/hook-torch.py | 42 +++++++++++++++++++ 2 files changed, 46 insertions(+) create mode 100644 news/676.update.2.rst diff --git a/news/676.update.2.rst b/news/676.update.2.rst new file mode 100644 index 00000000..3112846f --- /dev/null +++ b/news/676.update.2.rst @@ -0,0 +1,4 @@ +(Linux) Extend ``torch`` hook to automatically collect CUDA libraries +distributed via ``nvidia-*`` packages (such as ``nvidia-cuda-runtime-cu12``) +if they are specified among the requirements in the ``torch`` distribution's +metadata. diff --git a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-torch.py b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-torch.py index b6104400..e57deba6 100644 --- a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-torch.py +++ b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-torch.py @@ -20,6 +20,7 @@ ) if is_module_satisfies("PyInstaller >= 6.0"): + from PyInstaller.compat import is_linux from PyInstaller.utils.hooks import PY_DYLIB_PATTERNS module_collection_mode = "pyz+py" @@ -43,6 +44,47 @@ # Ensure we pick up fully-versioned .so files as well search_patterns=PY_DYLIB_PATTERNS + ['*.so.*'], ) + + # On Linux, torch wheels built with non-default CUDA version bundle CUDA libraries themselves (and should be handled + # by the above `collect_dynamic_libs`). Wheels built with default CUDA version (which are available on PyPI), on the + # other hand, use CUDA libraries provided by nvidia-* packages. Due to all possible combinations (CUDA libs from + # nvidia-* packages, torch-bundled CUDA libs, CPU-only CUDA libs) we do not add hidden imports directly, but instead + # attempt to infer them from requirements listed in the `torch` metadata. + if is_linux: + def _infer_nvidia_hiddenimports(): + import re + import packaging.requirements # Requirement of PyInstaller >= 6.0, so guaranteed to be available here. + from PyInstaller.compat import importlib_metadata + + dist = importlib_metadata.distribution("torch") + requirements = [packaging.requirements.Requirement(req) for req in dist.requires or []] + requirements = [req.name for req in requirements if req.marker is None or req.marker.evaluate()] + + # 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.+)-cu[\d]+$' + nvidia_hiddenimports = [] + + for req in requirements: + m = re.match(_PATTERN, req) + if m is None: + if req.startswith("nvidia-"): + logger.warning("hook-torch: unhandled NVIDIA CUDA requirement: %r!", req) + else: + # Convert + package_name = "nvidia." + m.group('subpackage').replace('-', '_') + nvidia_hiddenimports.append(package_name) + + return nvidia_hiddenimports + + try: + nvidia_hiddenimports = _infer_nvidia_hiddenimports() + except Exception: + # Log the exception, but make it non-fatal + logger.warning("hook-torch: failed to infer NVIDIA CUDA hidden imports!", exc_info=True) + nvidia_hiddenimports = [] + logger.info("hook-torch: inferred hidden imports for CUDA libraries: %r", nvidia_hiddenimports) + hiddenimports += nvidia_hiddenimports else: datas = [(get_package_paths("torch")[1], "torch")] From d6a66a0e703c4d6310b2b0972a950e2f79f833cb Mon Sep 17 00:00:00 2001 From: Rok Mandeljc Date: Tue, 12 Dec 2023 23:48:29 +0100 Subject: [PATCH 07/35] hooks: add hooks for subpackages of nvidia package On Linux, NVIDIA CUDA 11.x and 12.x shared libraries can be installed via PyPI wheels (e.g., `nvidia-cuda-runtime-cu12`, `nvidia-cudnn-cu12`, `nvidia-cublas-cu12`). These all provide sub-packages under `nvidia` top level package (e.g., `nvidia.cuda_runtime`, `nvidia.cudnn`, `nvidia.cublas`). Add hooks for these sub-packages that ensure that all shared libraries from the sub-packages `lib` directory are collected, in case they are dynamically loaded. For example, `torch` PyPI wheels for linux do not bundle CUDA inside the `torch/lib` (whereas the wheels from their own server, built with "non-default" CUDA versions bundle them), and dynamically load `libcudnn_ops_infer.so.8` from `nvidia.cudnn.lib`. --- news/676.new.rst | 2 + .../hooks/stdhooks/hook-nvidia.cublas.py | 15 +++++++ .../hooks/stdhooks/hook-nvidia.cuda_cupti.py | 15 +++++++ .../hooks/stdhooks/hook-nvidia.cuda_nvcc.py | 23 ++++++++++ .../hooks/stdhooks/hook-nvidia.cuda_nvrtc.py | 15 +++++++ .../stdhooks/hook-nvidia.cuda_runtime.py | 15 +++++++ .../hooks/stdhooks/hook-nvidia.cudnn.py | 15 +++++++ .../hooks/stdhooks/hook-nvidia.cufft.py | 15 +++++++ .../hooks/stdhooks/hook-nvidia.curand.py | 15 +++++++ .../hooks/stdhooks/hook-nvidia.cusolver.py | 15 +++++++ .../hooks/stdhooks/hook-nvidia.cusparse.py | 15 +++++++ .../hooks/stdhooks/hook-nvidia.nccl.py | 15 +++++++ .../hooks/stdhooks/hook-nvidia.nvjitlink.py | 15 +++++++ .../hooks/stdhooks/hook-nvidia.nvtx.py | 15 +++++++ .../hooks/utils/__init__.py | 1 + .../hooks/utils/nvidia_cuda.py | 42 +++++++++++++++++++ 16 files changed, 248 insertions(+) create mode 100644 news/676.new.rst create mode 100644 src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.cublas.py create mode 100644 src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.cuda_cupti.py create mode 100644 src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.cuda_nvcc.py create mode 100644 src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.cuda_nvrtc.py create mode 100644 src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.cuda_runtime.py create mode 100644 src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.cudnn.py create mode 100644 src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.cufft.py create mode 100644 src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.curand.py create mode 100644 src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.cusolver.py create mode 100644 src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.cusparse.py create mode 100644 src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.nccl.py create mode 100644 src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.nvjitlink.py create mode 100644 src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.nvtx.py create mode 100644 src/_pyinstaller_hooks_contrib/hooks/utils/__init__.py create mode 100644 src/_pyinstaller_hooks_contrib/hooks/utils/nvidia_cuda.py diff --git a/news/676.new.rst b/news/676.new.rst new file mode 100644 index 00000000..787cba51 --- /dev/null +++ b/news/676.new.rst @@ -0,0 +1,2 @@ +Add hooks for ``nvidia.*`` packages, which provide a way of installing +CUDA via PyPI wheels (e.g., ``nvidia-cuda-runtime-cu12``). diff --git a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.cublas.py b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.cublas.py new file mode 100644 index 00000000..5b8a94dc --- /dev/null +++ b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.cublas.py @@ -0,0 +1,15 @@ +# ------------------------------------------------------------------ +# 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.GPL.txt, distributed with +# this software. +# +# SPDX-License-Identifier: GPL-2.0-or-later +# ------------------------------------------------------------------ + +from _pyinstaller_hooks_contrib.hooks.utils.nvidia_cuda import collect_nvidia_cuda_binaries + +binaries = collect_nvidia_cuda_binaries(__file__) diff --git a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.cuda_cupti.py b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.cuda_cupti.py new file mode 100644 index 00000000..5b8a94dc --- /dev/null +++ b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.cuda_cupti.py @@ -0,0 +1,15 @@ +# ------------------------------------------------------------------ +# 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.GPL.txt, distributed with +# this software. +# +# SPDX-License-Identifier: GPL-2.0-or-later +# ------------------------------------------------------------------ + +from _pyinstaller_hooks_contrib.hooks.utils.nvidia_cuda import collect_nvidia_cuda_binaries + +binaries = collect_nvidia_cuda_binaries(__file__) diff --git a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.cuda_nvcc.py b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.cuda_nvcc.py new file mode 100644 index 00000000..6ee509df --- /dev/null +++ b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.cuda_nvcc.py @@ -0,0 +1,23 @@ +# ------------------------------------------------------------------ +# 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.GPL.txt, distributed with +# this software. +# +# SPDX-License-Identifier: GPL-2.0-or-later +# ------------------------------------------------------------------ + +from PyInstaller.utils.hooks import collect_data_files +from _pyinstaller_hooks_contrib.hooks.utils.nvidia_cuda import collect_nvidia_cuda_binaries + +# Ensures that versioned .so files are collected +binaries = collect_nvidia_cuda_binaries(__file__) + +# Collect additional resources: +# - ptxas executable (which strictly speaking, should be collected as a binary) +# - nvvm/libdevice/libdevice.10.bc file +# - C headers; assuming ptxas requires them - if that is not the case, we could filter them out. +datas = collect_data_files('nvidia.cuda_nvcc') diff --git a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.cuda_nvrtc.py b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.cuda_nvrtc.py new file mode 100644 index 00000000..5b8a94dc --- /dev/null +++ b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.cuda_nvrtc.py @@ -0,0 +1,15 @@ +# ------------------------------------------------------------------ +# 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.GPL.txt, distributed with +# this software. +# +# SPDX-License-Identifier: GPL-2.0-or-later +# ------------------------------------------------------------------ + +from _pyinstaller_hooks_contrib.hooks.utils.nvidia_cuda import collect_nvidia_cuda_binaries + +binaries = collect_nvidia_cuda_binaries(__file__) diff --git a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.cuda_runtime.py b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.cuda_runtime.py new file mode 100644 index 00000000..5b8a94dc --- /dev/null +++ b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.cuda_runtime.py @@ -0,0 +1,15 @@ +# ------------------------------------------------------------------ +# 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.GPL.txt, distributed with +# this software. +# +# SPDX-License-Identifier: GPL-2.0-or-later +# ------------------------------------------------------------------ + +from _pyinstaller_hooks_contrib.hooks.utils.nvidia_cuda import collect_nvidia_cuda_binaries + +binaries = collect_nvidia_cuda_binaries(__file__) diff --git a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.cudnn.py b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.cudnn.py new file mode 100644 index 00000000..5b8a94dc --- /dev/null +++ b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.cudnn.py @@ -0,0 +1,15 @@ +# ------------------------------------------------------------------ +# 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.GPL.txt, distributed with +# this software. +# +# SPDX-License-Identifier: GPL-2.0-or-later +# ------------------------------------------------------------------ + +from _pyinstaller_hooks_contrib.hooks.utils.nvidia_cuda import collect_nvidia_cuda_binaries + +binaries = collect_nvidia_cuda_binaries(__file__) diff --git a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.cufft.py b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.cufft.py new file mode 100644 index 00000000..5b8a94dc --- /dev/null +++ b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.cufft.py @@ -0,0 +1,15 @@ +# ------------------------------------------------------------------ +# 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.GPL.txt, distributed with +# this software. +# +# SPDX-License-Identifier: GPL-2.0-or-later +# ------------------------------------------------------------------ + +from _pyinstaller_hooks_contrib.hooks.utils.nvidia_cuda import collect_nvidia_cuda_binaries + +binaries = collect_nvidia_cuda_binaries(__file__) diff --git a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.curand.py b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.curand.py new file mode 100644 index 00000000..5b8a94dc --- /dev/null +++ b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.curand.py @@ -0,0 +1,15 @@ +# ------------------------------------------------------------------ +# 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.GPL.txt, distributed with +# this software. +# +# SPDX-License-Identifier: GPL-2.0-or-later +# ------------------------------------------------------------------ + +from _pyinstaller_hooks_contrib.hooks.utils.nvidia_cuda import collect_nvidia_cuda_binaries + +binaries = collect_nvidia_cuda_binaries(__file__) diff --git a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.cusolver.py b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.cusolver.py new file mode 100644 index 00000000..5b8a94dc --- /dev/null +++ b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.cusolver.py @@ -0,0 +1,15 @@ +# ------------------------------------------------------------------ +# 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.GPL.txt, distributed with +# this software. +# +# SPDX-License-Identifier: GPL-2.0-or-later +# ------------------------------------------------------------------ + +from _pyinstaller_hooks_contrib.hooks.utils.nvidia_cuda import collect_nvidia_cuda_binaries + +binaries = collect_nvidia_cuda_binaries(__file__) diff --git a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.cusparse.py b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.cusparse.py new file mode 100644 index 00000000..5b8a94dc --- /dev/null +++ b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.cusparse.py @@ -0,0 +1,15 @@ +# ------------------------------------------------------------------ +# 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.GPL.txt, distributed with +# this software. +# +# SPDX-License-Identifier: GPL-2.0-or-later +# ------------------------------------------------------------------ + +from _pyinstaller_hooks_contrib.hooks.utils.nvidia_cuda import collect_nvidia_cuda_binaries + +binaries = collect_nvidia_cuda_binaries(__file__) diff --git a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.nccl.py b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.nccl.py new file mode 100644 index 00000000..5b8a94dc --- /dev/null +++ b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.nccl.py @@ -0,0 +1,15 @@ +# ------------------------------------------------------------------ +# 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.GPL.txt, distributed with +# this software. +# +# SPDX-License-Identifier: GPL-2.0-or-later +# ------------------------------------------------------------------ + +from _pyinstaller_hooks_contrib.hooks.utils.nvidia_cuda import collect_nvidia_cuda_binaries + +binaries = collect_nvidia_cuda_binaries(__file__) diff --git a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.nvjitlink.py b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.nvjitlink.py new file mode 100644 index 00000000..5b8a94dc --- /dev/null +++ b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.nvjitlink.py @@ -0,0 +1,15 @@ +# ------------------------------------------------------------------ +# 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.GPL.txt, distributed with +# this software. +# +# SPDX-License-Identifier: GPL-2.0-or-later +# ------------------------------------------------------------------ + +from _pyinstaller_hooks_contrib.hooks.utils.nvidia_cuda import collect_nvidia_cuda_binaries + +binaries = collect_nvidia_cuda_binaries(__file__) diff --git a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.nvtx.py b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.nvtx.py new file mode 100644 index 00000000..5b8a94dc --- /dev/null +++ b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-nvidia.nvtx.py @@ -0,0 +1,15 @@ +# ------------------------------------------------------------------ +# 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.GPL.txt, distributed with +# this software. +# +# SPDX-License-Identifier: GPL-2.0-or-later +# ------------------------------------------------------------------ + +from _pyinstaller_hooks_contrib.hooks.utils.nvidia_cuda import collect_nvidia_cuda_binaries + +binaries = collect_nvidia_cuda_binaries(__file__) diff --git a/src/_pyinstaller_hooks_contrib/hooks/utils/__init__.py b/src/_pyinstaller_hooks_contrib/hooks/utils/__init__.py new file mode 100644 index 00000000..792d6005 --- /dev/null +++ b/src/_pyinstaller_hooks_contrib/hooks/utils/__init__.py @@ -0,0 +1 @@ +# diff --git a/src/_pyinstaller_hooks_contrib/hooks/utils/nvidia_cuda.py b/src/_pyinstaller_hooks_contrib/hooks/utils/nvidia_cuda.py new file mode 100644 index 00000000..769b7d30 --- /dev/null +++ b/src/_pyinstaller_hooks_contrib/hooks/utils/nvidia_cuda.py @@ -0,0 +1,42 @@ +# ------------------------------------------------------------------ +# 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.GPL.txt, distributed with +# this software. +# +# SPDX-License-Identifier: GPL-2.0-or-later +# ------------------------------------------------------------------ + +import os + +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 From a9389f17754ce20767d0d20413ff722161ca8d92 Mon Sep 17 00:00:00 2001 From: Rok Mandeljc Date: Wed, 13 Dec 2023 16:54:57 +0100 Subject: [PATCH 08/35] tests: add test for torchaudio that uses torchscript Add a test for torchaudio that uses a transform to resample a signal. The test shows the need to collect binaries from the torchaudio package, and, due to use of torchscript, also shows the need to collect source .py files. --- .../tests/test_pytorch.py | 37 +++++++++++++++++++ 1 file changed, 37 insertions(+) diff --git a/src/_pyinstaller_hooks_contrib/tests/test_pytorch.py b/src/_pyinstaller_hooks_contrib/tests/test_pytorch.py index 15fb055d..d17cc7c7 100644 --- a/src/_pyinstaller_hooks_contrib/tests/test_pytorch.py +++ b/src/_pyinstaller_hooks_contrib/tests/test_pytorch.py @@ -29,6 +29,43 @@ def test_torch(pyi_builder): """) +# Test with torchaudio transform that uses torchcript, which requires +# access to transforms' sources. +@importorskip('torchaudio') +@torch_onedir_only +def test_torchaudio_scripted_transforms(pyi_builder): + pyi_builder.test_source(""" + import numpy as np + + import torch.nn + import torchaudio.transforms + + # Generate a sine waveform + volume = 0.5 # range [0.0, 1.0] + sampling_rate = 44100 # sampling rate, Hz + duration = 5.0 # seconds + freq = 500.0 # sine frequency, Hz + + points = np.arange(0, sampling_rate * duration) + signal = volume * np.sin(2 * np.pi * points * freq / sampling_rate) + + # Resample the signal using scripted transform + transforms = torch.nn.Sequential( + torchaudio.transforms.Resample( + orig_freq=sampling_rate, + new_freq=sampling_rate // 2 + ), + ) + scripted_transforms = torch.jit.script(transforms) + + signal_tensor = torch.from_numpy(signal).float() + resampled_tensor = scripted_transforms(signal_tensor) + + print("Result size:", resampled_tensor.size()) + assert len(resampled_tensor) == len(signal_tensor) / 2 + """) + + @importorskip('torchvision') @torch_onedir_only def test_torchvision_nms(pyi_builder): From ae2b12b58b1b6c25247ea0d56c8a6e157835f452 Mon Sep 17 00:00:00 2001 From: Rok Mandeljc Date: Wed, 13 Dec 2023 16:56:22 +0100 Subject: [PATCH 09/35] hooks: add hook for torchaudio Add hook for torchaudio that collects binaries and ensures that source .py files are collected. --- news/676.new.1.rst | 2 ++ .../hooks/stdhooks/hook-torchaudio.py | 21 +++++++++++++++++++ 2 files changed, 23 insertions(+) create mode 100644 news/676.new.1.rst create mode 100644 src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-torchaudio.py diff --git a/news/676.new.1.rst b/news/676.new.1.rst new file mode 100644 index 00000000..c0ac33d3 --- /dev/null +++ b/news/676.new.1.rst @@ -0,0 +1,2 @@ +Add hook for ``torchaudio`` that collects dynamically-loaded extensions, +as well as source .py files for TorchScript/JIT. diff --git a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-torchaudio.py b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-torchaudio.py new file mode 100644 index 00000000..435ff58a --- /dev/null +++ b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-torchaudio.py @@ -0,0 +1,21 @@ +# ------------------------------------------------------------------ +# 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.GPL.txt, distributed with +# this software. +# +# SPDX-License-Identifier: GPL-2.0-or-later +# ------------------------------------------------------------------ + +from PyInstaller.utils.hooks import collect_dynamic_libs, collect_submodules + +# Collect dynamic extensions from torchaudio/lib - some of them are loaded dynamically, and are thus not automatically +# collected. +binaries = collect_dynamic_libs('torchaudio') +hiddenimports = collect_submodules('torchaudio.lib') + +# Collect source .py files for JIT/torchscript. Requires PyInstaller >= 5.3, no-op in older versions. +module_collection_mode = 'pyz+py' From 061efe94e7dfd1bc363c3e2e72a7f06e30e51d04 Mon Sep 17 00:00:00 2001 From: Rok Mandeljc Date: Wed, 13 Dec 2023 19:29:16 +0100 Subject: [PATCH 10/35] tests: add test for torchtext that uses torchscript Add a test for torchext that uses a tokenization transform from Berta Encoder. The test shows the need to collect binaries from the torchtext package, and, due to use of torchscript, also shows the need to collect source .py files. We perform only tokenization part of processing, in order to avoid having to download the whole model (~240 MB). --- .../tests/test_pytorch.py | 39 +++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/src/_pyinstaller_hooks_contrib/tests/test_pytorch.py b/src/_pyinstaller_hooks_contrib/tests/test_pytorch.py index d17cc7c7..354167a7 100644 --- a/src/_pyinstaller_hooks_contrib/tests/test_pytorch.py +++ b/src/_pyinstaller_hooks_contrib/tests/test_pytorch.py @@ -66,6 +66,45 @@ def test_torchaudio_scripted_transforms(pyi_builder): """) +# Test with torchtext transform that uses torchcript, which requires +# access to transforms' sources. +@importorskip('torchtext') +@torch_onedir_only +def test_torchtext_scripted_berta_tokenizer_transform(pyi_builder): + pyi_builder.test_source(""" + import torch.nn + import torchtext.models + import torchtext.functional + + # Create Roberta Encoder with Base configuration + roberta_base = torchtext.models.ROBERTA_BASE_ENCODER + classifier_head = torchtext.models.RobertaClassificationHead(num_classes=2, input_dim=768) + transform = roberta_base.transform() + + # Create transform that uses torchscript + scripted_transform = torch.jit.script(transform) + print(scripted_transform) + + # Prepare test data + small_input_batch = [ + "Hello world", + "How are you!", + ] + + model_input = torchtext.functional.to_tensor(scripted_transform(small_input_batch), padding_value=1) + print("Tokenized input:", model_input) + + # Process + if False: + # Downloads the model (~ 240 MB), if necessary. + model = roberta_base.get_model(head=classifier_head) + + output = model(model_input) + print(output) + print(output.shape) + """) + + @importorskip('torchvision') @torch_onedir_only def test_torchvision_nms(pyi_builder): From dd3bed672040e66ae69f73875cfb519427611db2 Mon Sep 17 00:00:00 2001 From: Rok Mandeljc Date: Wed, 13 Dec 2023 19:30:39 +0100 Subject: [PATCH 11/35] hooks: add hook for torchtext Add hook for torchtext that collects binaries and ensures that source .py files are collected. --- news/676.new.2.rst | 2 ++ .../hooks/stdhooks/hook-torchtext.py | 21 +++++++++++++++++++ 2 files changed, 23 insertions(+) create mode 100644 news/676.new.2.rst create mode 100644 src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-torchtext.py diff --git a/news/676.new.2.rst b/news/676.new.2.rst new file mode 100644 index 00000000..1987d59e --- /dev/null +++ b/news/676.new.2.rst @@ -0,0 +1,2 @@ +Add hook for ``torchtext`` that collects dynamically-loaded extensions, +as well as source .py files for TorchScript/JIT. diff --git a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-torchtext.py b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-torchtext.py new file mode 100644 index 00000000..548dec89 --- /dev/null +++ b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-torchtext.py @@ -0,0 +1,21 @@ +# ------------------------------------------------------------------ +# 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.GPL.txt, distributed with +# this software. +# +# SPDX-License-Identifier: GPL-2.0-or-later +# ------------------------------------------------------------------ + +from PyInstaller.utils.hooks import collect_dynamic_libs, collect_submodules + +# Collect dynamic extensions from torchtext/lib - some of them are loaded dynamically, and are thus not automatically +# collected. +binaries = collect_dynamic_libs('torchtext') +hiddenimports = collect_submodules('torchtext.lib') + +# Collect source .py files for JIT/torchscript. Requires PyInstaller >= 5.3, no-op in older versions. +module_collection_mode = 'pyz+py' From 293a4582ea2808783679838bcbdf6a692652e060 Mon Sep 17 00:00:00 2001 From: Rok Mandeljc Date: Wed, 13 Dec 2023 23:27:15 +0100 Subject: [PATCH 12/35] tests: move tensorflow tests to their own file Move tensorflow tests into separate `test_tensorflow` file. Improve the `tensorflow_onedir_only` test mark (force pyi_builder into generating only onedir case instead of skipping the onefile case). --- .../tests/test_libraries.py | 46 ---------------- .../tests/test_tensorflow.py | 55 +++++++++++++++++++ 2 files changed, 55 insertions(+), 46 deletions(-) create mode 100644 src/_pyinstaller_hooks_contrib/tests/test_tensorflow.py diff --git a/src/_pyinstaller_hooks_contrib/tests/test_libraries.py b/src/_pyinstaller_hooks_contrib/tests/test_libraries.py index b671b9e6..a74d5ca6 100644 --- a/src/_pyinstaller_hooks_contrib/tests/test_libraries.py +++ b/src/_pyinstaller_hooks_contrib/tests/test_libraries.py @@ -80,52 +80,6 @@ def test_geopandas(pyi_builder): ) -def tensorflow_onedir_only(test): - def wrapped(pyi_builder): - if pyi_builder._mode != 'onedir': - pytest.skip('Tensorflow tests support only onedir mode ' - 'due to potential distribution size.') - test(pyi_builder) - - return wrapped - - -@importorskip('tensorflow') -@tensorflow_onedir_only -def test_tensorflow(pyi_builder): - pyi_builder.test_source( - """ - from tensorflow import * - """ - ) - - -# Test if tensorflow.keras imports properly result in tensorflow being collected. -# See https://github.com/pyinstaller/pyinstaller/discussions/6890 -@importorskip('tensorflow') -@tensorflow_onedir_only -def test_tensorflow_keras_import(pyi_builder): - pyi_builder.test_source( - """ - from tensorflow.keras.models import Sequential - from tensorflow.keras.layers import Dense, LSTM, Dropout - from tensorflow.keras.optimizers import Adam - """ - ) - - -@importorskip('tensorflow') -@tensorflow_onedir_only -def test_tensorflow_layer(pyi_builder): - pyi_builder.test_script('pyi_lib_tensorflow_layer.py') - - -@importorskip('tensorflow') -@tensorflow_onedir_only -def test_tensorflow_mnist(pyi_builder): - pyi_builder.test_script('pyi_lib_tensorflow_mnist.py') - - @importorskip('trimesh') def test_trimesh(pyi_builder): pyi_builder.test_source( diff --git a/src/_pyinstaller_hooks_contrib/tests/test_tensorflow.py b/src/_pyinstaller_hooks_contrib/tests/test_tensorflow.py new file mode 100644 index 00000000..e10f36d2 --- /dev/null +++ b/src/_pyinstaller_hooks_contrib/tests/test_tensorflow.py @@ -0,0 +1,55 @@ +# ------------------------------------------------------------------ +# Copyright (c) 2020 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.GPL.txt, distributed with +# this software. +# +# SPDX-License-Identifier: GPL-2.0-or-later +# ------------------------------------------------------------------ + +import pytest + +from PyInstaller.utils.tests import importorskip + + +# Run the tests in onedir mode only +tensorflow_onedir_only = pytest.mark.parametrize('pyi_builder', ['onedir'], indirect=True) + + +@importorskip('tensorflow') +@tensorflow_onedir_only +def test_tensorflow(pyi_builder): + pyi_builder.test_source( + """ + from tensorflow import * + """ + ) + + +# Test if tensorflow.keras imports properly result in tensorflow being collected. +# See https://github.com/pyinstaller/pyinstaller/discussions/6890 +@importorskip('tensorflow') +@tensorflow_onedir_only +def test_tensorflow_keras_import(pyi_builder): + pyi_builder.test_source( + """ + from tensorflow.keras.models import Sequential + from tensorflow.keras.layers import Dense, LSTM, Dropout + from tensorflow.keras.optimizers import Adam + """ + ) + + +@importorskip('tensorflow') +@tensorflow_onedir_only +def test_tensorflow_layer(pyi_builder): + pyi_builder.test_script('pyi_lib_tensorflow_layer.py') + + +@importorskip('tensorflow') +@tensorflow_onedir_only +def test_tensorflow_mnist(pyi_builder): + pyi_builder.test_script('pyi_lib_tensorflow_mnist.py') From ceff7b1fadf9ac99fb72a7a78d1433fc33ecd251 Mon Sep 17 00:00:00 2001 From: Rok Mandeljc Date: Thu, 14 Dec 2023 00:11:46 +0100 Subject: [PATCH 13/35] tests: add tests for transformers package Add a basic `transformers` pipeline test. Add a basic import test for `transformers.DebertaModel`, which shows that we need to collect source .py files for TorchScript. --- .../tests/test_deep_learning.py | 45 +++++++++++++++++++ 1 file changed, 45 insertions(+) create mode 100644 src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py diff --git a/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py b/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py new file mode 100644 index 00000000..fe728e63 --- /dev/null +++ b/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py @@ -0,0 +1,45 @@ +# ------------------------------------------------------------------ +# 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.GPL.txt, distributed with +# this software. +# +# SPDX-License-Identifier: GPL-2.0-or-later +# ------------------------------------------------------------------ + +import pytest + +from PyInstaller.utils.tests import importorskip + + +# Run the tests in onedir mode only +onedir_only = pytest.mark.parametrize('pyi_builder', ['onedir'], indirect=True) + + +# Basic transformers test with BERT-based unmasker +@importorskip('transformers') +@importorskip('torch') +@onedir_only +def test_transformers_bert_pipeline(pyi_builder): + pyi_builder.test_source(""" + import transformers + unmasker = transformers.pipeline('fill-mask', model='bert-base-uncased') + output = unmasker("Hello I'm a [MASK] model.") + print("Unmasked text:", output) + """) + + +# Trying to import DebertaModel triggers error about missing source files for TorchScript +@importorskip('transformers') +@importorskip('torch') +@onedir_only +def test_transformers_deberta_import(pyi_builder): + pyi_builder.test_source(""" + from transformers import DebertaConfig, DebertaModel + + configuration = DebertaConfig() + model = DebertaModel(configuration) + """) From 663df322a6f88e2fdced422cbab208943630e6f1 Mon Sep 17 00:00:00 2001 From: Rok Mandeljc Date: Thu, 14 Dec 2023 00:13:03 +0100 Subject: [PATCH 14/35] hooks: add hook for transformers Add hook for Hugging Face `transformers` package. Attempt to automatically collect metadata for all of package's dependencies (as declared in `deps` dictionary in the `transformers.dependency_versions_table` module). Collect source .py files as some of the functionality uses TorchScript. --- news/676.new.3.rst | 6 +++ .../hooks/stdhooks/hook-transformers.py | 37 +++++++++++++++++++ 2 files changed, 43 insertions(+) create mode 100644 news/676.new.3.rst create mode 100644 src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-transformers.py diff --git a/news/676.new.3.rst b/news/676.new.3.rst new file mode 100644 index 00000000..6acca5f0 --- /dev/null +++ b/news/676.new.3.rst @@ -0,0 +1,6 @@ +Add hook for Hugging Face ``transformers``. The hook attempts to +automatically collect the metadata of all dependencies (as declared +in `deps` dictionary in the `transformers.dependency_versions_table` +module), in order to make dependencies available at build time visible +to ``transformers`` at run time. The hook also collects source .py files +as some of the package's functionality uses TorchScript/JIT. diff --git a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-transformers.py b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-transformers.py new file mode 100644 index 00000000..936733b4 --- /dev/null +++ b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-transformers.py @@ -0,0 +1,37 @@ +from PyInstaller.utils.hooks import ( + copy_metadata, + get_module_attribute, + is_module_satisfies, + logger, +) + +datas = [] + +# At run-time, `transformers` queries the metadata of several packages to check for their presence. The list of required +# (core) packages is stored as `transformers.dependency_versions_check.pkgs_to_check_at_runtime˙. However, there is more +# comprehensive list of dependencies and their versions available in `transformers.dependency_versions_table.deps`, +# which includes non-core dependencies. Unfortunately, we cannot foresee which of those the user will actually require, +# so we collect metadata for all listed dists that are available in the build environment, in order to make them visible +# to `transformers` at run-time. +try: + dependencies = get_module_attribute( + 'transformers.dependency_versions_table', + 'deps', + ) +except Exception: + logger.warning( + "hook-transformers: failed to query dependency table (transformers.dependency_versions_table.deps)!", + exc_info=True, + ) + dependencies = {} + +for dependency_name, dependency_req in dependencies.items(): + if not is_module_satisfies(dependency_req): + continue + try: + datas += copy_metadata(dependency_name) + except Exception: + pass + +# Collect source .py files for JIT/torchscript. Requires PyInstaller >= 5.3, no-op in older versions. +module_collection_mode = 'pyz+py' From fc744e9e0d3ab5e6a4aaa90be26554e04519d026 Mon Sep 17 00:00:00 2001 From: Rok Mandeljc Date: Thu, 14 Dec 2023 11:08:47 +0100 Subject: [PATCH 15/35] hooks: add a hook for fastai Add a hook for fastai, which ensures that the package's source .py files are collected, as they are required for TorchScript. Add a test based on the "building models from tabular data" example from https://docs.fast.ai/quick_start.html, which demonstrates the need for the source .py files. --- news/676.new.4.rst | 1 + .../hooks/stdhooks/hook-fastai.py | 14 +++++++ .../tests/test_deep_learning.py | 40 +++++++++++++++++++ 3 files changed, 55 insertions(+) create mode 100644 news/676.new.4.rst create mode 100644 src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-fastai.py diff --git a/news/676.new.4.rst b/news/676.new.4.rst new file mode 100644 index 00000000..4f73d9f4 --- /dev/null +++ b/news/676.new.4.rst @@ -0,0 +1 @@ +Add hook for ``fastai`` to collect its source .py files for TorchScript/JIT. diff --git a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-fastai.py b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-fastai.py new file mode 100644 index 00000000..6f47b491 --- /dev/null +++ b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-fastai.py @@ -0,0 +1,14 @@ +# ------------------------------------------------------------------ +# 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.GPL.txt, distributed with +# this software. +# +# SPDX-License-Identifier: GPL-2.0-or-later +# ------------------------------------------------------------------ + +# Collect source .py files for JIT/torchscript. Requires PyInstaller >= 5.3, no-op in older versions. +module_collection_mode = 'pyz+py' diff --git a/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py b/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py index fe728e63..97b6f5a0 100644 --- a/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py +++ b/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py @@ -43,3 +43,43 @@ def test_transformers_deberta_import(pyi_builder): configuration = DebertaConfig() model = DebertaModel(configuration) """) + + +# Building models from tabular data example from https://docs.fast.ai/quick_start.html +@importorskip('fastai') +@onedir_only +def test_fastai_tabular_data(pyi_builder): + pyi_builder.test_source(""" + from fastai.tabular.all import * + + path = untar_data(URLs.ADULT_SAMPLE) + print(f"Dataset path: {path}") + + dls = TabularDataLoaders.from_csv( + path/'adult.csv', + path=path, + y_names="salary", + cat_names = [ + 'workclass', + 'education', + 'marital-status', + 'occupation', + 'relationship', + 'race', + ], + cont_names = [ + 'age', + 'fnlwgt', + 'education-num', + ], + procs = [ + Categorify, + FillMissing, + Normalize, + ], + ) + + learn = tabular_learner(dls, metrics=accuracy) + learn.fit_one_cycle(2) + learn.show_results() + """) From bbb98725f0bd94dd65938dec59b4b71d7cf9e37e Mon Sep 17 00:00:00 2001 From: Rok Mandeljc Date: Thu, 14 Dec 2023 11:48:42 +0100 Subject: [PATCH 16/35] hooks: torchvision: ensure torchvision.io.image works `torchvision.io.image` attempts to dynamically load `torchvision.image` extension, so add a hook that ensures the latter is collected. Add a basic image encoding/decoding test. --- news/676.new.5.rst | 2 ++ .../hooks/stdhooks/hook-torchvision.io.image.py | 14 ++++++++++++++ .../tests/test_pytorch.py | 16 ++++++++++++++++ 3 files changed, 32 insertions(+) create mode 100644 news/676.new.5.rst create mode 100644 src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-torchvision.io.image.py diff --git a/news/676.new.5.rst b/news/676.new.5.rst new file mode 100644 index 00000000..f0f31e2f --- /dev/null +++ b/news/676.new.5.rst @@ -0,0 +1,2 @@ +Add hook for ``torchvision.io.image`` to ensure that dynamically-loaded +extension, required by this module, is collected. diff --git a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-torchvision.io.image.py b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-torchvision.io.image.py new file mode 100644 index 00000000..b16dfa2f --- /dev/null +++ b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-torchvision.io.image.py @@ -0,0 +1,14 @@ +# ------------------------------------------------------------------ +# Copyright (c) 2020 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.GPL.txt, distributed with +# this software. +# +# SPDX-License-Identifier: GPL-2.0-or-later +# ------------------------------------------------------------------ + +# torchivison.io.image attempts to dynamically load the torchvision.image extension. +hiddenimports = ['torchvision.image'] diff --git a/src/_pyinstaller_hooks_contrib/tests/test_pytorch.py b/src/_pyinstaller_hooks_contrib/tests/test_pytorch.py index 354167a7..bfa7c25f 100644 --- a/src/_pyinstaller_hooks_contrib/tests/test_pytorch.py +++ b/src/_pyinstaller_hooks_contrib/tests/test_pytorch.py @@ -129,6 +129,22 @@ def test_torchvision_nms(pyi_builder): """) +# Ensure that torchvision.io.image manages to load torchvision.image extension for its ops. +@importorskip('torchvision') +@torch_onedir_only +def test_torchvision_image_io(pyi_builder): + pyi_builder.test_source(""" + import torch + import torchvision.io.image + + image = torch.zeros((3, 100, 100), dtype=torch.uint8) + png_data = torchvision.io.image.encode_png(image) + decoded_image = torchvision.io.image.decode_png(png_data) + + assert torch.equal(image, decoded_image), "Original and decoded image are not identical!" + """) + + # Advanced tests that uses torchvision.datasets and torchvision.transforms; # the transforms are combined using torchscript, which requires access to # transforms' sources. From 030e9bb65cb293804ce5d4b067c50b36deafa32e Mon Sep 17 00:00:00 2001 From: Rok Mandeljc Date: Thu, 14 Dec 2023 12:00:25 +0100 Subject: [PATCH 17/35] hooks: add hook for timm (Hugging face torch image models) Add hook for timm, which ensures that the package's source .py files are collected, as they are required for TorchScript. Add a basic model listing and creation test. --- news/676.new.6.rst | 2 ++ .../hooks/stdhooks/hook-timm.py | 14 ++++++++++++++ .../tests/test_deep_learning.py | 17 +++++++++++++++++ 3 files changed, 33 insertions(+) create mode 100644 news/676.new.6.rst create mode 100644 src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-timm.py diff --git a/news/676.new.6.rst b/news/676.new.6.rst new file mode 100644 index 00000000..66cc48fd --- /dev/null +++ b/news/676.new.6.rst @@ -0,0 +1,2 @@ +Add hook for ``timm`` (Hugging Face PyTorch Image Models) to collect its +source .py files for TorchScript/JIT. diff --git a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-timm.py b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-timm.py new file mode 100644 index 00000000..6f47b491 --- /dev/null +++ b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-timm.py @@ -0,0 +1,14 @@ +# ------------------------------------------------------------------ +# 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.GPL.txt, distributed with +# this software. +# +# SPDX-License-Identifier: GPL-2.0-or-later +# ------------------------------------------------------------------ + +# Collect source .py files for JIT/torchscript. Requires PyInstaller >= 5.3, no-op in older versions. +module_collection_mode = 'pyz+py' diff --git a/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py b/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py index 97b6f5a0..28e2b939 100644 --- a/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py +++ b/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py @@ -83,3 +83,20 @@ def test_fastai_tabular_data(pyi_builder): learn.fit_one_cycle(2) learn.show_results() """) + + +@importorskip('timm') +@onedir_only +def test_timm_model_creation(pyi_builder): + pyi_builder.test_source(""" + import timm + + # List available models + pretrained_models = timm.list_models(pretrained=True) + print(f"Pre-trained models: {len(pretrained_models)}") + assert len(pretrained_models) > 0 + + # Create a model (non-trained version, to avoid downloading weights) + model = timm.create_model("resnet50d", pretrained=False) + print(model) + """) From 64751d9f6dc586d753c7fd09cee1709addfc473f Mon Sep 17 00:00:00 2001 From: Rok Mandeljc Date: Thu, 14 Dec 2023 22:43:25 +0100 Subject: [PATCH 18/35] tests: add test for lightning Add test for `lightning`, based on their "Getting started" example with autoencoder trained on MNIST dataset. Requires `torchivsion` for dataset. --- .../tests/test_deep_learning.py | 70 +++++++++++++++++++ 1 file changed, 70 insertions(+) diff --git a/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py b/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py index 28e2b939..defd5d5c 100644 --- a/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py +++ b/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py @@ -100,3 +100,73 @@ def test_timm_model_creation(pyi_builder): model = timm.create_model("resnet50d", pretrained=False) print(model) """) + + +@importorskip('lightning') +@importorskip('torchvision') +@importorskip('torch') +@onedir_only +def test_lightning_mnist_autoencoder(pyi_builder): + pyi_builder.test_source(""" + import os + + import torch + import torchvision + import lightning + + + class LitAutoEncoder(lightning.LightningModule): + def __init__(self): + super().__init__() + self.encoder = torch.nn.Sequential( + torch.nn.Linear(28 * 28, 128), + torch.nn.ReLU(), + torch.nn.Linear(128, 3), + ) + self.decoder = torch.nn.Sequential( + torch.nn.Linear(3, 128), + torch.nn.ReLU(), + torch.nn.Linear(128, 28 * 28), + ) + + def forward(self, x): + embedding = self.encoder(x) + return embedding + + def training_step(self, batch, batch_idx): + x, y = batch + x = x.view(x.size(0), -1) + z = self.encoder(x) + x_hat = self.decoder(z) + loss = torch.nn.functional.mse_loss(x_hat, x) + return loss + + def configure_optimizers(self): + optimizer = torch.optim.Adam( + self.parameters(), + lr=1e-3, + ) + return optimizer + + + # Dataset + dataset = torchvision.datasets.MNIST( + os.path.dirname(__file__), + download=True, + transform=torchvision.transforms.ToTensor(), + ) + dataset_size = len(dataset) + num_samples = 100 + train, val = torch.utils.data.random_split( + dataset, + [num_samples, dataset_size - num_samples], + ) + + # Train + autoencoder = LitAutoEncoder() + trainer = lightning.Trainer(max_epochs=1, logger=False) + trainer.fit( + autoencoder, + torch.utils.data.DataLoader(train), + ) + """) From aa493db5412e9440aa534f83d452a6c9c330436f Mon Sep 17 00:00:00 2001 From: Rok Mandeljc Date: Sat, 16 Dec 2023 16:46:46 +0100 Subject: [PATCH 19/35] hooks: add hook for lightning Add hook for (PyTorch) `lightning`. Currently, the main functionality is to ensure that the `version.info` file from the package is collected. We do not collect source .py files, as it seems that even if `lightning.LightningModule.to_torchscript()` is used, it requires the source where the model inheriting from `lightning.LightningModule` is defined, rather than `lightning`'s own sources. We can always add source .py files collection later, if it proves to be necessary. --- news/676.new.7.rst | 2 ++ .../hooks/stdhooks/hook-lightning.py | 21 +++++++++++++++++++ 2 files changed, 23 insertions(+) create mode 100644 news/676.new.7.rst create mode 100644 src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-lightning.py diff --git a/news/676.new.7.rst b/news/676.new.7.rst new file mode 100644 index 00000000..a2259737 --- /dev/null +++ b/news/676.new.7.rst @@ -0,0 +1,2 @@ +Add hook for ``lightning`` (PyTorch Lightning) to ensure that its +``version.info`` data file is collected. diff --git a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-lightning.py b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-lightning.py new file mode 100644 index 00000000..40833f60 --- /dev/null +++ b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-lightning.py @@ -0,0 +1,21 @@ +# ------------------------------------------------------------------ +# 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.GPL.txt, distributed with +# this software. +# +# SPDX-License-Identifier: GPL-2.0-or-later +# ------------------------------------------------------------------ + +from PyInstaller.utils.hooks import collect_data_files + +# Collect version.info (which is read during package import at run-time). Avoid collecting data from `lightning.app`, +# which likely does not work with PyInstaller without additional tricks (if we need to collect that data, it should +# be done in separate `lightning.app` hook). +datas = collect_data_files( + 'lightning', + includes=['version.info'], +) From a81ba9a52610b68a2f3f1d782579e4b2e0c2fd10 Mon Sep 17 00:00:00 2001 From: Rok Mandeljc Date: Sat, 16 Dec 2023 18:03:04 +0100 Subject: [PATCH 20/35] hooks: add hooks for bitsandbytes and triton Add hooks for `bitsandbytes`, and its dependency `triton`. Both packages have dynamically-loaded extension libraries, and both require collection of source .py files for (`triton`'s) JIT module. With `triton`, some of the submodules need to be collected only as source .py files (no PYZ), because the code naively tries to read the files pointed to by `__file__` attribute under assumption that they are source .py files. So we must not collect these modules into the PYZ. --- news/676.new.8.rst | 7 ++++ .../hooks/stdhooks/hook-bitsandbytes.py | 23 +++++++++++++ .../hooks/stdhooks/hook-triton.py | 32 +++++++++++++++++++ .../tests/test_deep_learning.py | 14 ++++++++ 4 files changed, 76 insertions(+) create mode 100644 news/676.new.8.rst create mode 100644 src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-bitsandbytes.py create mode 100644 src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-triton.py diff --git a/news/676.new.8.rst b/news/676.new.8.rst new file mode 100644 index 00000000..d6a4054d --- /dev/null +++ b/news/676.new.8.rst @@ -0,0 +1,7 @@ +Add hooks for ``bitsandbytes``, and its dependency ``triton``. Both +packages have dynamically-loaded extension libraries that need to be +collected, and both require collection of source .py files for +(``triton``'s) JIT module. Some submodules of ``triton`` need to be +collected only as source .py files (bypassing PYZ archive), because the +code naively assumes that ``__file__`` attribute points to the source +.py file. diff --git a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-bitsandbytes.py b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-bitsandbytes.py new file mode 100644 index 00000000..45acc97d --- /dev/null +++ b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-bitsandbytes.py @@ -0,0 +1,23 @@ +# ------------------------------------------------------------------ +# 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.GPL.txt, distributed with +# this software. +# +# SPDX-License-Identifier: GPL-2.0-or-later +# --------------------------------------------------- + +from PyInstaller.utils.hooks import collect_dynamic_libs + +# bitsandbytes contains several extensions for CPU and different CUDA versions: libbitsandbytes_cpu.so, +# libbitsandbytes_cuda110_nocublaslt.so, libbitsandbytes_cuda110.so, etc. At build-time, we could query the +# `bitsandbytes.cextension.setup` and its `binary_name˙ attribute for the extension that is in use. However, if the +# build system does not have CUDA available, this would automatically mean that we will not collect any of the CUDA +# libs. So for now, we collect them all. +binaries = collect_dynamic_libs("bitsandbytes") + +# bitsandbytes uses triton's JIT module, which requires access to source .py files. +module_collection_mode = 'pyz+py' diff --git a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-triton.py b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-triton.py new file mode 100644 index 00000000..ef9cebaf --- /dev/null +++ b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-triton.py @@ -0,0 +1,32 @@ +# ------------------------------------------------------------------ +# 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.GPL.txt, distributed with +# this software. +# +# SPDX-License-Identifier: GPL-2.0-or-later +# --------------------------------------------------- + +from PyInstaller.utils.hooks import collect_data_files, collect_dynamic_libs + +# Ensure that triton/_C/libtriton.so is collected +binaries = collect_dynamic_libs('triton') + +# triton has a JIT module that requires its source .py files. For some god-forsaken reason, this JIT module +# (`triton.runtime.jit` attempts to directly read the contents of file pointed to by its `__file__` attribute (assuming +# it is a source file). Therefore, `triton.runtime.jit` must not be collected into PYZ. Same goes for `compiler` and +# `language` sub-packages. +module_collection_mode = { + 'triton': 'pyz+py', + 'triton.runtime.jit': 'py', + 'triton.compiler': 'py', + 'triton.language': 'py', +} + +# Collect ptxas compiler files from triton/third_party/cuda directory. Strictly speaking, the ptxas executable from bin +# directory should be collected as a binary, but in this case, it makes no difference (plus, PyInstaller >= 6.0 has +# automatic binary-vs-data reclassification). +datas = collect_data_files('triton.third_party.cuda') diff --git a/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py b/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py index defd5d5c..d4c9a562 100644 --- a/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py +++ b/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py @@ -170,3 +170,17 @@ def configure_optimizers(self): torch.utils.data.DataLoader(train), ) """) + + +@importorskip('bitsandbytes') +@onedir_only +def test_bitsandbytes(pyi_builder): + pyi_builder.test_source(""" + import bitsandbytes + + # Instantiate a model and optimizer + dim1 = 256 + dim2 = 256 + linear = bitsandbytes.nn.Linear8bitLt(dim1, dim2, bias=True, has_fp16_weights=False, threshold=6.0) + adam = bitsandbytes.optim.Adam8bit(linear.parameters(), lr=0.001, betas=(0.9, 0.995)) + """) From 7f579890c717e5e1d61a62fe4fb84a3f2d5fa9de Mon Sep 17 00:00:00 2001 From: Rok Mandeljc Date: Sat, 16 Dec 2023 18:21:22 +0100 Subject: [PATCH 21/35] hooks: add hook for linear_operator Add hook and basic test for `linear_operator`. The package uses torchscript/JIT, so we need to collect its source .py files. --- news/676.new.9.rst | 2 ++ .../hooks/stdhooks/hook-linear_operator.py | 14 ++++++++++++++ .../tests/test_deep_learning.py | 17 +++++++++++++++++ 3 files changed, 33 insertions(+) create mode 100644 news/676.new.9.rst create mode 100644 src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-linear_operator.py diff --git a/news/676.new.9.rst b/news/676.new.9.rst new file mode 100644 index 00000000..b1d947b8 --- /dev/null +++ b/news/676.new.9.rst @@ -0,0 +1,2 @@ +Add hook for ``linear_operator`` to collect its source .py files for +TorchScript/JIT. diff --git a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-linear_operator.py b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-linear_operator.py new file mode 100644 index 00000000..962e0beb --- /dev/null +++ b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-linear_operator.py @@ -0,0 +1,14 @@ +# ------------------------------------------------------------------ +# 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.GPL.txt, distributed with +# this software. +# +# SPDX-License-Identifier: GPL-2.0-or-later +# --------------------------------------------------- + +# Collect source .py files for JIT/torchscript. Requires PyInstaller >= 5.3, no-op in older versions. +module_collection_mode = 'pyz+py' diff --git a/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py b/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py index d4c9a562..6a666d4e 100644 --- a/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py +++ b/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py @@ -184,3 +184,20 @@ def test_bitsandbytes(pyi_builder): linear = bitsandbytes.nn.Linear8bitLt(dim1, dim2, bias=True, has_fp16_weights=False, threshold=6.0) adam = bitsandbytes.optim.Adam8bit(linear.parameters(), lr=0.001, betas=(0.9, 0.995)) """) + + +@importorskip('linear_operator') +@onedir_only +def test_linear_operator(pyi_builder): + pyi_builder.test_source(""" + import torch + from linear_operator.operators import DiagLinearOperator, LowRankRootLinearOperator + + diag1 = 0.1 + torch.rand(100) + diag2 = 0.1 + torch.rand(100) + + mat1 = DiagLinearOperator(diag1) + mat2 = DiagLinearOperator(diag2) + + result = (mat1 + mat2).diagonal() + """) From 85bf2ceb716f91b86d1168086f4af19271a69f50 Mon Sep 17 00:00:00 2001 From: Rok Mandeljc Date: Sat, 16 Dec 2023 18:34:07 +0100 Subject: [PATCH 22/35] tests: add test for gpytorch Add a basic test for GPyTorch, based on their "Simple GP Regression" example. --- .../tests/test_deep_learning.py | 77 +++++++++++++++++++ 1 file changed, 77 insertions(+) diff --git a/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py b/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py index 6a666d4e..f462d7eb 100644 --- a/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py +++ b/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py @@ -201,3 +201,80 @@ def test_linear_operator(pyi_builder): result = (mat1 + mat2).diagonal() """) + + +# Based on https://docs.gpytorch.ai/en/latest/examples/01_Exact_GPs/Simple_GP_Regression.html +@importorskip('gpytorch') +@onedir_only +def test_gpytorch_simple_gp_regression(pyi_builder): + pyi_builder.test_source(""" + import math + + import torch + import gpytorch + + ## Training + # Training data is 100 points in [0,1] inclusive regularly spaced + train_x = torch.linspace(0, 1, 100) + + # True function is sin(2*pi*x) with Gaussian noise + train_y = torch.sin(train_x * (2 * math.pi)) + torch.randn(train_x.size()) * math.sqrt(0.04) + + # We will use the simplest form of GP model, exact inference + class ExactGPModel(gpytorch.models.ExactGP): + def __init__(self, train_x, train_y, likelihood): + super().__init__(train_x, train_y, likelihood) + self.mean_module = gpytorch.means.ConstantMean() + self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel()) + + def forward(self, x): + mean_x = self.mean_module(x) + covar_x = self.covar_module(x) + return gpytorch.distributions.MultivariateNormal(mean_x, covar_x) + + # Initialize likelihood and model + likelihood = gpytorch.likelihoods.GaussianLikelihood() + model = ExactGPModel(train_x, train_y, likelihood) + + # Find optimal model hyperparameters + training_iter = 2 + + model.train() + likelihood.train() + + # Use the adam optimizer + optimizer = torch.optim.Adam(model.parameters(), lr=0.1) # Includes GaussianLikelihood parameters + + # "Loss" for GPs - the marginal log likelihood + mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model) + + print("Training the model...") + for i in range(training_iter): + # Zero gradients from previous iteration + optimizer.zero_grad() + # Output from model + output = model(train_x) + # Calc loss and backprop gradients + loss = -mll(output, train_y) + loss.backward() + print('Iter %d/%d - Loss: %.3f lengthscale: %.3f noise: %.3f' % ( + i + 1, training_iter, loss.item(), + model.covar_module.base_kernel.lengthscale.item(), + model.likelihood.noise.item() + )) + optimizer.step() + + ## Inference + # Get into evaluation (predictive posterior) mode + model.eval() + likelihood.eval() + + # Test points are regularly spaced along [0,1] + # Make predictions by feeding model through likelihood + with torch.no_grad(), gpytorch.settings.fast_pred_var(): + test_x = torch.linspace(0, 1, 51) + observed_pred = likelihood(model(test_x)) + + print("Test X:", test_x.numpy()) + print("Predicted Y:", observed_pred.mean.numpy()) + """) From 368fc1177ea47a7a61f6cc89c8b8f8929a081d82 Mon Sep 17 00:00:00 2001 From: Rok Mandeljc Date: Sat, 16 Dec 2023 19:02:27 +0100 Subject: [PATCH 23/35] hooks: add hook for fvcore.nn Add hook for `fvcore.nn` to collect its source .py files for torchscript/JIT. Add a basic import test that demonstrates the need for that. --- news/676.new.10.rst | 2 ++ .../hooks/stdhooks/hook-fvcore.nn.py | 14 ++++++++++++++ .../tests/test_deep_learning.py | 9 +++++++++ 3 files changed, 25 insertions(+) create mode 100644 news/676.new.10.rst create mode 100644 src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-fvcore.nn.py diff --git a/news/676.new.10.rst b/news/676.new.10.rst new file mode 100644 index 00000000..376795b3 --- /dev/null +++ b/news/676.new.10.rst @@ -0,0 +1,2 @@ +Add hook for ``fvcore.nn`` to collect its source .py files for +TorchScript/JIT. diff --git a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-fvcore.nn.py b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-fvcore.nn.py new file mode 100644 index 00000000..6f47b491 --- /dev/null +++ b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-fvcore.nn.py @@ -0,0 +1,14 @@ +# ------------------------------------------------------------------ +# 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.GPL.txt, distributed with +# this software. +# +# SPDX-License-Identifier: GPL-2.0-or-later +# ------------------------------------------------------------------ + +# Collect source .py files for JIT/torchscript. Requires PyInstaller >= 5.3, no-op in older versions. +module_collection_mode = 'pyz+py' diff --git a/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py b/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py index f462d7eb..4e54d556 100644 --- a/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py +++ b/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py @@ -278,3 +278,12 @@ def forward(self, x): print("Test X:", test_x.numpy()) print("Predicted Y:", observed_pred.mean.numpy()) """) + + +# Basic import test for fvcore.nn, which shows that we need to collect its source.py files for TorchScript/JIT. +@importorskip('fvcore') +@onedir_only +def test_fvcore(pyi_builder): + pyi_builder.test_source(""" + import fvcore.nn + """) From 5b51e853c2de5e684c1d7381516685a09db07fe3 Mon Sep 17 00:00:00 2001 From: Rok Mandeljc Date: Sat, 16 Dec 2023 19:17:10 +0100 Subject: [PATCH 24/35] hooks: add hook for detectron2 Add hook for `detectron` to collect its source .py files for torchscript/JIT. Add a basic import test that demonstrates the need for that. --- news/676.new.11.rst | 2 ++ .../hooks/stdhooks/hook-detectron2.py | 14 ++++++++++++++ .../tests/test_deep_learning.py | 18 ++++++++++++++++++ 3 files changed, 34 insertions(+) create mode 100644 news/676.new.11.rst create mode 100644 src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-detectron2.py diff --git a/news/676.new.11.rst b/news/676.new.11.rst new file mode 100644 index 00000000..d0c005d3 --- /dev/null +++ b/news/676.new.11.rst @@ -0,0 +1,2 @@ +Add hook for ``detectron2`` to collect its source .py files for +TorchScript/JIT. diff --git a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-detectron2.py b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-detectron2.py new file mode 100644 index 00000000..6f47b491 --- /dev/null +++ b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-detectron2.py @@ -0,0 +1,14 @@ +# ------------------------------------------------------------------ +# 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.GPL.txt, distributed with +# this software. +# +# SPDX-License-Identifier: GPL-2.0-or-later +# ------------------------------------------------------------------ + +# Collect source .py files for JIT/torchscript. Requires PyInstaller >= 5.3, no-op in older versions. +module_collection_mode = 'pyz+py' diff --git a/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py b/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py index 4e54d556..9094ba64 100644 --- a/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py +++ b/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py @@ -287,3 +287,21 @@ def test_fvcore(pyi_builder): pyi_builder.test_source(""" import fvcore.nn """) + + +# Basic test for detectron2, which shows that we need to collect its source.py files for TorchScript/JIT. +@importorskip('detectron2') +@onedir_only +def test_detectron2(pyi_builder): + pyi_builder.test_source(""" + from detectron2 import model_zoo + from detectron2.config import get_cfg + from detectron2.engine import DefaultTrainer + + cfg = get_cfg() + print("Config:", cfg) + + # We cannot instantiate DefaultTrainer without specifying training datasets in config... + #trainer = DefaultTrainer(cfg) + #print(trainer) + """) From 4ba8a67df691d67ce5468d3c9fc9c37debd74935 Mon Sep 17 00:00:00 2001 From: Rok Mandeljc Date: Sat, 16 Dec 2023 19:36:33 +0100 Subject: [PATCH 25/35] hooks: add hook for Hugging Face datasets Add hook for `datasets` to collect its source .py files for torchscript/JIT. Add a basic dataset loading test that demonstrates the need for that. --- news/676.new.12.rst | 2 ++ .../hooks/stdhooks/hook-datasets.py | 14 ++++++++++++++ .../tests/test_deep_learning.py | 19 +++++++++++++++++++ 3 files changed, 35 insertions(+) create mode 100644 news/676.new.12.rst create mode 100644 src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-datasets.py diff --git a/news/676.new.12.rst b/news/676.new.12.rst new file mode 100644 index 00000000..aad40d42 --- /dev/null +++ b/news/676.new.12.rst @@ -0,0 +1,2 @@ +Add hook for Hugging Face ``datasets`` to collect its source .py files for +TorchScript/JIT. diff --git a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-datasets.py b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-datasets.py new file mode 100644 index 00000000..6f47b491 --- /dev/null +++ b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-datasets.py @@ -0,0 +1,14 @@ +# ------------------------------------------------------------------ +# 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.GPL.txt, distributed with +# this software. +# +# SPDX-License-Identifier: GPL-2.0-or-later +# ------------------------------------------------------------------ + +# Collect source .py files for JIT/torchscript. Requires PyInstaller >= 5.3, no-op in older versions. +module_collection_mode = 'pyz+py' diff --git a/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py b/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py index 9094ba64..45822aa2 100644 --- a/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py +++ b/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py @@ -305,3 +305,22 @@ def test_detectron2(pyi_builder): #trainer = DefaultTrainer(cfg) #print(trainer) """) + + +# Hugging Face datasets: Download squad dataset (76 MB train, 10 MB validation) +@importorskip('datasets') +@onedir_only +def test_datasets_download_squad(pyi_builder): + pyi_builder.test_source(""" + from datasets import load_dataset + from huggingface_hub import list_datasets + + # Print all the available datasets + available_datasets = [dataset.id for dataset in list_datasets()] + print("Available datasets:", len(available_datasets)) + + # Load a dataset and print the first example in the training set + print("Loading squad dataset...") + squad_dataset = load_dataset('squad') + print("First sample:", squad_dataset['train'][0]) + """) From fac988fe890c3898fe8f09cbf30647467a9362ba Mon Sep 17 00:00:00 2001 From: Rok Mandeljc Date: Sun, 17 Dec 2023 18:12:39 +0100 Subject: [PATCH 26/35] tests: add a basic test for Hugging Face accelerate Add basic test for Hugging Face `accelerate`; demonstrates that for the tested subset of functionality, no hook is necessary. --- .../tests/test_deep_learning.py | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py b/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py index 45822aa2..0008beee 100644 --- a/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py +++ b/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py @@ -324,3 +324,24 @@ def test_datasets_download_squad(pyi_builder): squad_dataset = load_dataset('squad') print("First sample:", squad_dataset['train'][0]) """) + + +# Basic test for Hugging Face accelerate framework +@importorskip('accelerate') +@onedir_only +def test_accelerate(pyi_builder): + pyi_builder.test_source(""" + import torch + from accelerate import Accelerator + + accelerator = Accelerator() + device = accelerator.device + print("Accelerator device:", device) + + model = torch.nn.Transformer().to(device) + optimizer = torch.optim.Adam(model.parameters()) + + model, optimizer = accelerator.prepare(model, optimizer) + print("Model:", model) + print("Optimizer:", optimizer) + """) From 4319060e7dfc41ee6719a5d3a559d3b164bc67cf Mon Sep 17 00:00:00 2001 From: Rok Mandeljc Date: Sun, 17 Dec 2023 18:14:05 +0100 Subject: [PATCH 27/35] hook: tensorflow: reformat line wrapping Use 120-character lines to reduce amount of line wrapping and make the hook easier to read. --- .../hooks/stdhooks/hook-tensorflow.py | 27 +++++++++---------- 1 file changed, 12 insertions(+), 15 deletions(-) diff --git a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-tensorflow.py b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-tensorflow.py index faf220c0..6a49efe9 100644 --- a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-tensorflow.py +++ b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-tensorflow.py @@ -10,8 +10,11 @@ # SPDX-License-Identifier: GPL-2.0-or-later # ------------------------------------------------------------------ -from PyInstaller.utils.hooks import is_module_satisfies, \ - collect_submodules, collect_data_files +from PyInstaller.utils.hooks import ( + collect_data_files, + collect_submodules, + is_module_satisfies, +) tf_pre_1_15_0 = is_module_satisfies("tensorflow < 1.15.0") tf_post_1_15_0 = is_module_satisfies("tensorflow >= 1.15.0") @@ -31,11 +34,9 @@ "**/*.lib", ] -# Under tensorflow 2.3.0 (the most recent version at the time of writing), -# _pywrap_tensorflow_internal extension module ends up duplicated; once -# as an extension, and once as a shared library. In addition to increasing -# program size, this also causes problems on macOS, so we try to prevent -# the extension module "variant" from being picked up. +# Under tensorflow 2.3.0 (the most recent version at the time of writing), _pywrap_tensorflow_internal extension module +# ends up duplicated; once as an extension, and once as a shared library. In addition to increasing program size, this +# also causes problems on macOS, so we try to prevent the extension module "variant" from being picked up. # # See pyinstaller/pyinstaller-hooks-contrib#49 for details. excluded_submodules = ['tensorflow.python._pywrap_tensorflow_internal'] @@ -47,24 +48,20 @@ def _submodules_filter(x): if tf_pre_1_15_0: # 1.14.x and earlier: collect everything from tensorflow - hiddenimports = collect_submodules('tensorflow', - filter=_submodules_filter) + hiddenimports = collect_submodules('tensorflow', filter=_submodules_filter) datas = collect_data_files('tensorflow', excludes=data_excludes) elif tf_post_1_15_0 and tf_pre_2_2_0: # 1.15.x - 2.1.x: collect everything from tensorflow_core - hiddenimports = collect_submodules('tensorflow_core', - filter=_submodules_filter) + hiddenimports = collect_submodules('tensorflow_core', filter=_submodules_filter) datas = collect_data_files('tensorflow_core', excludes=data_excludes) # Under 1.15.x, we seem to fail collecting a specific submodule, # and need to add it manually... if tf_post_1_15_0 and tf_pre_2_0_0: - hiddenimports += \ - ['tensorflow_core._api.v1.compat.v2.summary.experimental'] + hiddenimports += ['tensorflow_core._api.v1.compat.v2.summary.experimental'] else: # 2.2.0 and newer: collect everything from tensorflow again - hiddenimports = collect_submodules('tensorflow', - filter=_submodules_filter) + hiddenimports = collect_submodules('tensorflow', filter=_submodules_filter) datas = collect_data_files('tensorflow', excludes=data_excludes) # From 2.6.0 on, we also need to explicitly collect keras (due to From 2b4eeaea0ccbc2826ec29ea09269590e85fec360 Mon Sep 17 00:00:00 2001 From: Rok Mandeljc Date: Sun, 17 Dec 2023 18:17:16 +0100 Subject: [PATCH 28/35] hook: tensorflow: revise the _pywrap_tensorflow_internal hack Remove the `tensorflow.python._pywrap_tensorflow_internal` hack (adding it to excluded modules to avoid duplication) for PyInstaller >= 6.0, where the issue is alleviated thanks to the binary dependency analysis preserving the parent directory layout of discovered/collected shared libraries. This should fix the problem with `tensorflow` builds where the `_pywrap_tensorflow_internal` module is not used as a shared library, as seen in `tensorflow` builds for Raspberry Pi: https://github.com/pyinstaller/pyinstaller-hooks-contrib/issues/121 --- news/676.update.3.rst | 8 ++++++++ .../hooks/stdhooks/hook-tensorflow.py | 9 ++++++++- 2 files changed, 16 insertions(+), 1 deletion(-) create mode 100644 news/676.update.3.rst diff --git a/news/676.update.3.rst b/news/676.update.3.rst new file mode 100644 index 00000000..45bf9e3f --- /dev/null +++ b/news/676.update.3.rst @@ -0,0 +1,8 @@ +(Linux) Remove the ``tensorflow.python._pywrap_tensorflow_internal`` +hack in the ``tensorflow`` hook (i.e., adding it to excluded modules +to avoid duplication) when using PyInstaller >= 6.0, where the +duplication issue is alleviated thanks to the binary dependency analysis +preserving the parent directory layout of discovered/collected shared +libraries. This should fix the problem with ``tensorflow`` builds where +the ``_pywrap_tensorflow_internal`` module is not used as a shared +library, as seen in ``tensorflow`` builds for Raspberry Pi. diff --git a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-tensorflow.py b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-tensorflow.py index 6a49efe9..d3fa7d97 100644 --- a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-tensorflow.py +++ b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-tensorflow.py @@ -39,7 +39,14 @@ # also causes problems on macOS, so we try to prevent the extension module "variant" from being picked up. # # See pyinstaller/pyinstaller-hooks-contrib#49 for details. -excluded_submodules = ['tensorflow.python._pywrap_tensorflow_internal'] +# +# With PyInstaller >= 6.0, this issue is alleviated, because the binary dependency analysis (which picks up the +# extension in question as a shared library that other extensions are linked against) now preserves the parent directory +# layout, and creates a symbolic link to the top-level application directory. +if is_module_satisfies('PyInstaller >= 6.0'): + excluded_submodules = [] +else: + excluded_submodules = ['tensorflow.python._pywrap_tensorflow_internal'] def _submodules_filter(x): From 8103da54e71b1a7e3d6fd63da3851ded24667081 Mon Sep 17 00:00:00 2001 From: Rok Mandeljc Date: Sun, 17 Dec 2023 22:23:28 +0100 Subject: [PATCH 29/35] hooks: tensorflow: collect sources for tensorflow.python.autograph Collect sources for `tensorflow.python.autograph` to avoid run-time warning about AutoGraph being unavailable. Not sure if we need to collect sources for other parts of `tensorflow`, though; if that proves to be the case, we can adjust the collection mode later. --- .../hooks/stdhooks/hook-tensorflow.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-tensorflow.py b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-tensorflow.py index d3fa7d97..687af661 100644 --- a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-tensorflow.py +++ b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-tensorflow.py @@ -85,3 +85,12 @@ def _submodules_filter(x): # Suppress warnings for missing hidden imports generated by this hook. # Requires PyInstaller > 5.1 (with pyinstaller/pyinstaller#6914 merged); no-op otherwise. warn_on_missing_hiddenimports = False + +# Collect the AutoGraph part of `tensorflow` code, to avoid a run-time warning about AutoGraph being unavailable: +# `WARNING:tensorflow:AutoGraph is not available in this environment: functions lack code information. ...` +# The warning is emitted if source for `log` function from `tensorflow.python.autograph.utils.ag_logging` cannot be +# looked up. Not sure if we need sources for other parts of `tesnorflow`, though. +# Requires PyInstaller >= 5.3, no-op in older versions. +module_collection_mode = { + 'tensorflow.python.autograph': 'py+pyz', +} From c85cc87cac5be32b2c631b3d6364587585fd95da Mon Sep 17 00:00:00 2001 From: Rok Mandeljc Date: Sun, 17 Dec 2023 22:33:12 +0100 Subject: [PATCH 30/35] Add _pyinstaller_hooks_contrib.compat module Add `_pyinstaller_hooks_contrib.compat` module to hide the gory details of stdlib `importlib.metadata` vs. `importlib_metadata`. Import the `importlib_metadata` from `PyInstaller.compat` if available (PyInstaller >= 6.0), otherwise duplicate the logic. Copy `importlib_metadata` and `packaging` requirements from PyInstaller to pyinstaller-hooks-contrib. --- news/676.update.4.rst | 3 ++ setup.cfg | 4 +++ src/_pyinstaller_hooks_contrib/compat.py | 42 ++++++++++++++++++++++++ 3 files changed, 49 insertions(+) create mode 100644 news/676.update.4.rst create mode 100644 src/_pyinstaller_hooks_contrib/compat.py diff --git a/news/676.update.4.rst b/news/676.update.4.rst new file mode 100644 index 00000000..ad34bc8c --- /dev/null +++ b/news/676.update.4.rst @@ -0,0 +1,3 @@ +Update ``tensorflow`` hook to collect source .py files for +``tensorflow.python.autograph`` in order to silence a run-time warning +about AutoGraph not being available. diff --git a/setup.cfg b/setup.cfg index f9d58fa5..d09edac6 100644 --- a/setup.cfg +++ b/setup.cfg @@ -32,6 +32,10 @@ package_dir= =src packages=find: python_requires = >=3.7 +install_requires = + setuptools >= 42.0.0 + importlib_metadata >= 4.6 ; python_version < "3.10" + packaging >= 22.0 [options.packages.find] where=src diff --git a/src/_pyinstaller_hooks_contrib/compat.py b/src/_pyinstaller_hooks_contrib/compat.py new file mode 100644 index 00000000..2e18c022 --- /dev/null +++ b/src/_pyinstaller_hooks_contrib/compat.py @@ -0,0 +1,42 @@ +# ------------------------------------------------------------------ +# 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.GPL.txt, distributed with +# this software. +# +# SPDX-License-Identifier: GPL-2.0-or-later +# ------------------------------------------------------------------ + +import sys + +from PyInstaller.utils.hooks import is_module_satisfies + + +if is_module_satisfies("PyInstaller >= 6.0"): + # PyInstaller >= 6.0 imports importlib_metadata in its compat module + from PyInstaller.compat import importlib_metadata +else: + # Older PyInstaller version - duplicate logic from PyInstaller 6.0 + class ImportlibMetadataError(SystemExit): + def __init__(self): + super().__init__( + "pyinstaller-hooks-contrib requires importlib.metadata from python >= 3.10 stdlib or " + "importlib_metadata from importlib-metadata >= 4.6" + ) + + if sys.version_info >= (3, 10): + import importlib.metadata as importlib_metadata + else: + try: + import importlib_metadata + except ImportError as e: + raise ImportlibMetadataError() from e + + import packaging.version # For importlib_metadata version check + + # Validate the version + if packaging.version.parse(importlib_metadata.version("importlib-metadata")) < packaging.version.parse("4.6"): + raise ImportlibMetadataError() From 6ef93046071b8831b12e3ac3107c6e3de9a3bf6a Mon Sep 17 00:00:00 2001 From: Rok Mandeljc Date: Sun, 17 Dec 2023 23:19:24 +0100 Subject: [PATCH 31/35] hooks: tensorflow: rework the tensorflow version check Determine the tensorflow's dist name using list of potential candidates, and if a dist is found, retrieve the version from it. Otherwise, fall back to reading `tensorflow.__version__`. The `tensorflow` package is available in several variants, and sometimes the `tensorflow` dist installs a separate dist (e.g., `tensorflow-intel` on Windows); but the user can install this separate dist without installing the "top-level" `tensorflow` one. In PyInstaller v5 and earlier, the `is_module_satisfies` fell back to querying `tensorflow.__version__` if the dist could not be found - in v6, the implementation of `is_module_satisfies` (or rather, `check_requirement`) checks only the metadata. As an added bonus, the direct version comparisons are nicer to read than the comparisons against `tf_pre_1_15_0`, `tf_post_1_15_0`, etc. --- news/676.update.5.rst | 5 ++ .../hooks/stdhooks/hook-tensorflow.py | 70 +++++++++++++++---- 2 files changed, 62 insertions(+), 13 deletions(-) create mode 100644 news/676.update.5.rst diff --git a/news/676.update.5.rst b/news/676.update.5.rst new file mode 100644 index 00000000..a7eb6a2d --- /dev/null +++ b/news/676.update.5.rst @@ -0,0 +1,5 @@ +Update ``tensorflow`` hook to attempt to resolve the top-level distribution +name and infer the package version from it, in order to improve version +handling when the "top-level" ``tensorflow`` dist is not installed (for +example, user installs only ``tensorflow-intel`` or ``tensorflow-macos``) +or has a different name (e.g., ``tf-nightly``). diff --git a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-tensorflow.py b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-tensorflow.py index 687af661..6496baf6 100644 --- a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-tensorflow.py +++ b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-tensorflow.py @@ -10,16 +10,62 @@ # SPDX-License-Identifier: GPL-2.0-or-later # ------------------------------------------------------------------ +from _pyinstaller_hooks_contrib.compat import importlib_metadata +from packaging.version import Version + from PyInstaller.utils.hooks import ( collect_data_files, collect_submodules, + get_module_attribute, is_module_satisfies, + logger, +) + +# Determine the name of `tensorflow` dist; this is available under different names (releases vs. nightly, plus build +# variants). We need to determine the dist that we are dealing with, so we can query its version and metadata. +_CANDIDATE_DIST_NAMES = ( + "tensorflow", + "tensorflow-cpu", + "tensorflow-gpu", + "tensorflow-intel", + "tensorflow-rocm", + "tensorflow-macos", + "tensorflow-aarch64", + "tensorflow-cpu-aws", + "tf-nightly", + "tf-nightly-cpu", + "tf-nightly-gpu", + "tf-nightly-rocm", + "intel-tensorflow", + "intel-tensorflow-avx512", ) +dist = None +for candidate_dist_name in _CANDIDATE_DIST_NAMES: + try: + dist = importlib_metadata.distribution(candidate_dist_name) + break + except importlib_metadata.PackageNotFoundError: + continue + +version = None +if dist is None: + logger.warning( + "hook-tensorflow: failed to determine tensorflow dist name! Reading version from tensorflow.__version__!" + ) + try: + version = get_module_attribute("tensorflow", "__version__") + except Exception as e: + raise Exception("Failed to read tensorflow.__version__") from e +else: + logger.info("hook-tensorflow: tensorflow dist name: %s", dist.name) + version = dist.version -tf_pre_1_15_0 = is_module_satisfies("tensorflow < 1.15.0") -tf_post_1_15_0 = is_module_satisfies("tensorflow >= 1.15.0") -tf_pre_2_0_0 = is_module_satisfies("tensorflow < 2.0.0") -tf_pre_2_2_0 = is_module_satisfies("tensorflow < 2.2.0") +# Parse version +logger.info("hook-tensorflow: tensorflow version: %s", version) +try: + version = Version(version) +except Exception as e: + raise Exception("Failed to parse tensorflow version!") from e # Exclude from data collection: # - development headers in include subdirectory @@ -53,31 +99,29 @@ def _submodules_filter(x): return x not in excluded_submodules -if tf_pre_1_15_0: +if version < Version("1.15.0a0"): # 1.14.x and earlier: collect everything from tensorflow hiddenimports = collect_submodules('tensorflow', filter=_submodules_filter) datas = collect_data_files('tensorflow', excludes=data_excludes) -elif tf_post_1_15_0 and tf_pre_2_2_0: +elif version >= Version("1.15.0a0") and version < Version("2.2.0a0"): # 1.15.x - 2.1.x: collect everything from tensorflow_core hiddenimports = collect_submodules('tensorflow_core', filter=_submodules_filter) datas = collect_data_files('tensorflow_core', excludes=data_excludes) - # Under 1.15.x, we seem to fail collecting a specific submodule, - # and need to add it manually... - if tf_post_1_15_0 and tf_pre_2_0_0: + # Under 1.15.x, we seem to fail collecting a specific submodule, and need to add it manually... + if version < Version("2.0.0a0"): hiddenimports += ['tensorflow_core._api.v1.compat.v2.summary.experimental'] else: # 2.2.0 and newer: collect everything from tensorflow again hiddenimports = collect_submodules('tensorflow', filter=_submodules_filter) datas = collect_data_files('tensorflow', excludes=data_excludes) - # From 2.6.0 on, we also need to explicitly collect keras (due to - # lazy mapping of tensorflow.keras.xyz -> keras.xyz) - if is_module_satisfies("tensorflow >= 2.6.0"): + # From 2.6.0 on, we also need to explicitly collect keras (due to lazy mapping of tensorflow.keras.xyz -> keras.xyz) + if version >= Version("2.6.0a0"): hiddenimports += collect_submodules('keras') # Starting with 2.14.0, we need `ml_dtypes` among hidden imports. - if is_module_satisfies("tensorflow >= 2.14.0"): + if version >= Version("2.14.0"): hiddenimports += ['ml_dtypes'] excludedimports = excluded_submodules From 39f7dee1f22323c82037ce405f16e5f3654231aa Mon Sep 17 00:00:00 2001 From: Rok Mandeljc Date: Wed, 20 Dec 2023 13:47:28 +0100 Subject: [PATCH 32/35] hooks: tensorflow: generate hidden imports for nvidia.* modules Linux builds of `tensorflow` can install CUDA via nvidia-* packages that are enabled by `and-cuda` extra marker. So parse the dist metadata for requirements, and turn the `nvidia-*` requirements into `nvidia.*` hidden imports. Consolidate the shared code for conversion of `nvidia-*` dist name into `nvidia.*` module in a utility function, and use it in both `torch` and `tensorflow` hooks. --- news/676.update.6.rst | 4 +++ .../hooks/stdhooks/hook-tensorflow.py | 27 +++++++++++++++++++ .../hooks/stdhooks/hook-torch.py | 23 +++------------- .../hooks/utils/nvidia_cuda.py | 20 ++++++++++++++ 4 files changed, 55 insertions(+), 19 deletions(-) create mode 100644 news/676.update.6.rst diff --git a/news/676.update.6.rst b/news/676.update.6.rst new file mode 100644 index 00000000..7011b9e3 --- /dev/null +++ b/news/676.update.6.rst @@ -0,0 +1,4 @@ +(Linux) Extend ``tensorflow`` hook to automatically collect CUDA libraries +distributed via ``nvidia-*`` packages (such as ``nvidia-cuda-runtime-cu12``) +if they are specified among the requirements in the ``tensorflow`` +distribution's metadata. diff --git a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-tensorflow.py b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-tensorflow.py index 6496baf6..97dfe6b5 100644 --- a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-tensorflow.py +++ b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-tensorflow.py @@ -13,6 +13,7 @@ from _pyinstaller_hooks_contrib.compat import importlib_metadata from packaging.version import Version +from PyInstaller.compat import is_linux from PyInstaller.utils.hooks import ( collect_data_files, collect_submodules, @@ -138,3 +139,29 @@ def _submodules_filter(x): module_collection_mode = { 'tensorflow.python.autograph': 'py+pyz', } + +# Linux builds of tensorflow can optionally use CUDA from nvidia-* packages. If we managed to obtain dist, query the +# requirements from metadata (the `and-cuda` extra marker), and convert them to module names. +# +# NOTE: while the installation of nvidia-* packages via `and-cuda` extra marker is not gated by the OS version check, +# it is effectively available only on Linux (last Windows-native build that supported GPU is v2.10.0, and assumed that +# CUDA is externally available). +if is_linux and dist is not None: + def _infer_nvidia_hiddenimports(): + import packaging.requirements + from _pyinstaller_hooks_contrib.hooks.utils import nvidia_cuda as cudautils + + requirements = [packaging.requirements.Requirement(req) for req in dist.requires or []] + env = {'extra': 'and-cuda'} + requirements = [req.name for req in requirements if req.marker is None or req.marker.evaluate(env)] + + return cudautils.infer_hiddenimports_from_requirements(requirements) + + try: + nvidia_hiddenimports = _infer_nvidia_hiddenimports() + except Exception: + # Log the exception, but make it non-fatal + logger.warning("hook-tensorflow: failed to infer NVIDIA CUDA hidden imports!", exc_info=True) + nvidia_hiddenimports = [] + logger.info("hook-tensorflow: inferred hidden imports for CUDA libraries: %r", nvidia_hiddenimports) + hiddenimports += nvidia_hiddenimports diff --git a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-torch.py b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-torch.py index e57deba6..01b6b3ae 100644 --- a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-torch.py +++ b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-torch.py @@ -52,30 +52,15 @@ # attempt to infer them from requirements listed in the `torch` metadata. if is_linux: def _infer_nvidia_hiddenimports(): - import re - import packaging.requirements # Requirement of PyInstaller >= 6.0, so guaranteed to be available here. - from PyInstaller.compat import importlib_metadata + import packaging.requirements + from _pyinstaller_hooks_contrib.compat import importlib_metadata + from _pyinstaller_hooks_contrib.hooks.utils import nvidia_cuda as cudautils dist = importlib_metadata.distribution("torch") requirements = [packaging.requirements.Requirement(req) for req in dist.requires or []] requirements = [req.name for req in requirements if req.marker is None or req.marker.evaluate()] - # 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.+)-cu[\d]+$' - nvidia_hiddenimports = [] - - for req in requirements: - m = re.match(_PATTERN, req) - if m is None: - if req.startswith("nvidia-"): - logger.warning("hook-torch: unhandled NVIDIA CUDA requirement: %r!", req) - else: - # Convert - package_name = "nvidia." + m.group('subpackage').replace('-', '_') - nvidia_hiddenimports.append(package_name) - - return nvidia_hiddenimports + return cudautils.infer_hiddenimports_from_requirements(requirements) try: nvidia_hiddenimports = _infer_nvidia_hiddenimports() diff --git a/src/_pyinstaller_hooks_contrib/hooks/utils/nvidia_cuda.py b/src/_pyinstaller_hooks_contrib/hooks/utils/nvidia_cuda.py index 769b7d30..3b697c92 100644 --- a/src/_pyinstaller_hooks_contrib/hooks/utils/nvidia_cuda.py +++ b/src/_pyinstaller_hooks_contrib/hooks/utils/nvidia_cuda.py @@ -11,6 +11,7 @@ # ------------------------------------------------------------------ import os +import re from PyInstaller.utils.hooks import ( logger, @@ -40,3 +41,22 @@ def collect_nvidia_cuda_binaries(hook_file): 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.+)-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 From cce102169a2a033dd949fb2d12f3c0545c3d5813 Mon Sep 17 00:00:00 2001 From: Rok Mandeljc Date: Wed, 20 Dec 2023 14:00:08 +0100 Subject: [PATCH 33/35] pytest: consolidate pytest.ini into setup.cfg Consolidate `pytest` configuration from `pytest.ini` into `setup.cfg`, to match what we have in the main PyInstaller repository. Add -v, -rsxXfE, and ----doctest-glob= to test flags. The addition of -v ensures that in manual (local) pytest runs, the test names are displayed as they are ran (the CI workflows seem to explicitly specify -v as part of their pytest commnads). --- pytest.ini | 8 -------- setup.cfg | 11 +++++++++++ 2 files changed, 11 insertions(+), 8 deletions(-) delete mode 100644 pytest.ini diff --git a/pytest.ini b/pytest.ini deleted file mode 100644 index 422d0a6e..00000000 --- a/pytest.ini +++ /dev/null @@ -1,8 +0,0 @@ -[pytest] -addopts=--maxfail=3 -m "not slow" - -markers = - darwin: only run on macOS - linux: only runs on GNU/Linux - win32: only runs on Windows - slow: Long tests are disabled by default. Re-enable with -m slow diff --git a/setup.cfg b/setup.cfg index d09edac6..4c366457 100644 --- a/setup.cfg +++ b/setup.cfg @@ -47,3 +47,14 @@ where=src # E265 - block comment should start with '# ' extend-ignore = E265 max-line-length=120 + +[tool:pytest] +# Display summary info for (s)skipped, (x)xfailed, (X)xpassed, (f)failed and (E)errored tests +# Skip doctest text files +addopts=--maxfail=3 -m "not slow" -v -rsxXfE --doctest-glob= + +markers = + darwin: only run on macOS + linux: only runs on GNU/Linux + win32: only runs on Windows + slow: Long tests are disabled by default. Re-enable with -m slow From 0772f24dc5ffe7db6246b2784121cfddf1561f92 Mon Sep 17 00:00:00 2001 From: Rok Mandeljc Date: Wed, 20 Dec 2023 15:28:38 +0100 Subject: [PATCH 34/35] tests: add multiprocessing.freeze_support() call to lightning test When running the `test_lightning_mnist_autoencoder` under arm64 macOS, `multiprocessing` seems to be activated at some point, and the test gets stuck due to lack of `multiprocessing.freeze_support` call. --- src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py b/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py index 0008beee..c5e7b54a 100644 --- a/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py +++ b/src/_pyinstaller_hooks_contrib/tests/test_deep_learning.py @@ -110,6 +110,11 @@ def test_lightning_mnist_autoencoder(pyi_builder): pyi_builder.test_source(""" import os + # On macOS, multiprocessing seems to be used at some point... + if __name__ == '__main__': + import multiprocessing + multiprocessing.freeze_support() + import torch import torchvision import lightning From 5ec43b8a2802008abd8ca1581124d11a857ba55f Mon Sep 17 00:00:00 2001 From: Rok Mandeljc Date: Thu, 21 Dec 2023 12:09:47 +0100 Subject: [PATCH 35/35] hooks: tensorflow: collect plugins from tensorflow-plugins Have the `tensorflow` standard hook collect binaries from the `tensorflow-plugins` package; this contains plugins for tensorflow's pluggable device architecture (such as `tensorflow-metal` for macOS and `tensorflow-directml-plugin` for Windows). Have the `tensorflow` run-time hook override the `site.getsitepackages()` with custom implementation that allows us to trick `tensorflow` into loading the plugins. --- news/676.update.7.rst | 5 ++ .../hooks/rthooks/pyi_rth_tensorflow.py | 50 ++++++++++++++++--- .../hooks/stdhooks/hook-tensorflow.py | 7 +++ 3 files changed, 54 insertions(+), 8 deletions(-) create mode 100644 news/676.update.7.rst diff --git a/news/676.update.7.rst b/news/676.update.7.rst new file mode 100644 index 00000000..9b6b456d --- /dev/null +++ b/news/676.update.7.rst @@ -0,0 +1,5 @@ +Extend ``tensorflow`` hook to collect plugins installed in the +``tensorflow-plugins`` directory/package. Have the run-time ``tensorflow`` +hook provide an override for ``site.getsitepackages()`` that allows us +to work around a broken module file location check and trick ``tensorflow`` +into loading the collected plugins. diff --git a/src/_pyinstaller_hooks_contrib/hooks/rthooks/pyi_rth_tensorflow.py b/src/_pyinstaller_hooks_contrib/hooks/rthooks/pyi_rth_tensorflow.py index c967f489..5941013e 100644 --- a/src/_pyinstaller_hooks_contrib/hooks/rthooks/pyi_rth_tensorflow.py +++ b/src/_pyinstaller_hooks_contrib/hooks/rthooks/pyi_rth_tensorflow.py @@ -9,11 +9,45 @@ # SPDX-License-Identifier: Apache-2.0 #----------------------------------------------------------------------------- -# `tensorflow` versions prior to 2.3.0 attempt to use `site.USER_SITE` in path/string manipulation functions. -# As frozen application runs with disabled `site`, the value of this variable is `None`, and causes path/string -# manipulation functions to raise an error. As a work-around, we set `site.USER_SITE` to an empty string, which is -# also what the fake `site` module available in PyInstaller prior to v5.5 did. -import site - -if site.USER_SITE is None: - site.USER_SITE = '' +def _pyi_rthook(): + import sys + + # `tensorflow` versions prior to 2.3.0 attempt to use `site.USER_SITE` in path/string manipulation functions. + # As frozen application runs with disabled `site`, the value of this variable is `None`, and causes path/string + # manipulation functions to raise an error. As a work-around, we set `site.USER_SITE` to an empty string, which is + # also what the fake `site` module available in PyInstaller prior to v5.5 did. + import site + + if site.USER_SITE is None: + site.USER_SITE = '' + + # The issue described about with site.USER_SITE being None has largely been resolved in contemporary `tensorflow` + # versions, which now check that `site.ENABLE_USER_SITE` is set and that `site.USER_SITE` is not None before + # trying to use it. + # + # However, `tensorflow` will attempt to search and load its plugins only if it believes that it is running from + # "a pip-based installation" - if the package's location is rooted in one of the "site-packages" directories. See + # https://github.com/tensorflow/tensorflow/blob/6887368d6d46223f460358323c4b76d61d1558a8/tensorflow/api_template.__init__.py#L110C76-L156 + # Unfortunately, they "cleverly" infer the module's location via `inspect.getfile(inspect.currentframe())`, which + # in the frozen application returns anonymized relative source file name (`tensorflow/__init__.py`) - so we need one + # of the "site directories" to be just "tensorflow" (to fool the `_running_from_pip_package()` check), and we also + # need `sys._MEIPASS` to be among them (to load the plugins from the actual `sys._MEIPASS/tensorflow-plugins`). + # Therefore, we monkey-patch `site.getsitepackages` to add those two entries to the list of "site directories". + + _orig_getsitepackages = getattr(site, 'getsitepackages') + + def _pyi_getsitepackages(): + return [ + sys._MEIPASS, + "tensorflow", + *(_orig_getsitepackages() if _orig_getsitepackages is not None else []), + ] + + site.getsitepackages = _pyi_getsitepackages + + # NOTE: instead of the above override, we could also set TF_PLUGGABLE_DEVICE_LIBRARY_PATH, but that works only + # for tensorflow >= 2.12. + + +_pyi_rthook() +del _pyi_rthook diff --git a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-tensorflow.py b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-tensorflow.py index 97dfe6b5..64da3b94 100644 --- a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-tensorflow.py +++ b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-tensorflow.py @@ -16,6 +16,7 @@ from PyInstaller.compat import is_linux from PyInstaller.utils.hooks import ( collect_data_files, + collect_dynamic_libs, collect_submodules, get_module_attribute, is_module_satisfies, @@ -125,6 +126,7 @@ def _submodules_filter(x): if version >= Version("2.14.0"): hiddenimports += ['ml_dtypes'] +binaries = [] excludedimports = excluded_submodules # Suppress warnings for missing hidden imports generated by this hook. @@ -165,3 +167,8 @@ def _infer_nvidia_hiddenimports(): nvidia_hiddenimports = [] logger.info("hook-tensorflow: inferred hidden imports for CUDA libraries: %r", nvidia_hiddenimports) hiddenimports += nvidia_hiddenimports + + +# Collect the tensorflow-plugins (pluggable device plugins) +hiddenimports += ['tensorflow-plugins'] +binaries += collect_dynamic_libs('tensorflow-plugins')