From f1923f53c20c69521ad4eae61243dbd1d019b78d Mon Sep 17 00:00:00 2001 From: Alexander Grund Date: Tue, 14 Jan 2025 15:41:34 +0100 Subject: [PATCH] Update TensorRT easyblock for TensorRT 10+ --- easybuild/easyblocks/t/tensorrt.py | 20 +++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/easybuild/easyblocks/t/tensorrt.py b/easybuild/easyblocks/t/tensorrt.py index 2e066e510c..fb3fbbaf8e 100644 --- a/easybuild/easyblocks/t/tensorrt.py +++ b/easybuild/easyblocks/t/tensorrt.py @@ -1,5 +1,5 @@ ## -# Copyright 2017-2025 Ghent University +# Copyright 2017-2024 Ghent University # # This file is part of EasyBuild, # originally created by the HPC team of Ghent University (http://ugent.be/hpc/en), @@ -96,11 +96,15 @@ def extensions_step(self): super(EB_TensorRT, self).extensions_step() pyver = ''.join(get_software_version('Python').split('.')[:2]) - whls = [ - os.path.join('graphsurgeon', 'graphsurgeon-*-py2.py3-none-any.whl'), - os.path.join('uff', 'uff-*-py2.py3-none-any.whl'), - os.path.join('python', 'tensorrt-%s-cp%s-*-linux_x86_64.whl' % (self.version, pyver)), - ] + whls = [] + # graphsurgeon and uff removed in 10.0.1 + if self.version < LooseVersion('10.0.1'): + whls.extend([ + os.path.join('graphsurgeon', 'graphsurgeon-*-py2.py3-none-any.whl'), + os.path.join('uff', 'uff-*-py2.py3-none-any.whl'), + ]) + whls.append(os.path.join('python', 'tensorrt-%s-cp%s-*-linux_x86_64.whl' % (self.version, pyver))) + for whl in whls: whl_paths = glob.glob(os.path.join(self.installdir, whl)) if len(whl_paths) == 1: @@ -118,8 +122,10 @@ def extensions_step(self): cmd += " --ignore-installed --no-deps" run_cmd(cmd, log_all=True, simple=True, log_ok=True) + elif whl_paths: + raise EasyBuildError("Failed to isolate .whl in %s: %s", self.installdir, whl_paths) else: - raise EasyBuildError("Failed to isolate .whl in %s: %s", whl_paths, self.installdir) + raise EasyBuildError("No .whl found in %s for patter %s", self.installdir, whl) def sanity_check_step(self): """Custom sanity check for TensorRT."""