Update TensorRT easyblock for TensorRT 10+

easybuilders · Jan 14, 2025 · f1923f5 · f1923f5
1 parent 3469151
commit f1923f5
Showing 1 changed file with 13 additions and 7 deletions.
diff --git 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."""