From a724c05f31f0bf39b8da30e0beaca5c05385197f Mon Sep 17 00:00:00 2001 From: Aarni Koskela Date: Tue, 30 Jan 2024 13:23:45 +0200 Subject: [PATCH] CUDA setup: make version number replacement logic more obvious --- bitsandbytes/cuda_setup/main.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/bitsandbytes/cuda_setup/main.py b/bitsandbytes/cuda_setup/main.py index 760d557a4..1669b08e1 100644 --- a/bitsandbytes/cuda_setup/main.py +++ b/bitsandbytes/cuda_setup/main.py @@ -112,9 +112,13 @@ def manual_override(self): if not override_value: return - binary_name = self.binary_name.rsplit(".", 1)[0] - # TODO: what's the magic value `-3` here? - self.binary_name = binary_name[:-3] + f'{override_value}{DYNAMIC_LIBRARY_SUFFIX}' + binary_name_stem, _, binary_name_ext = self.binary_name.rpartition(".") + # `binary_name_stem` will now be e.g. `/foo/bar/libbitsandbytes_cuda118`; + # let's remove any trailing numbers: + binary_name_stem = binary_name_stem.rstrip("0123456789") + # `binary_name_stem` will now be e.g. `/foo/bar/libbitsandbytes_cuda`; + # let's tack the new version number and the original extension back on. + self.binary_name = f"{binary_name_stem}{override_value}.{binary_name_ext}" warn( f'\n\n{"=" * 80}\n'