Skip to content

Commit 8cda845

Browse files
committed
use importlib machinery to check for habana plugin
Not all environments have pip, and even fewer have grep (e.g. Windows platforms), and the shell command is relatively costly for larger environments. Instead, use the stdlib infrastructure with `importlib` to identify the version of `habana-torch-plugin` if it is installed.
1 parent e1dc75a commit 8cda845

1 file changed

Lines changed: 10 additions & 11 deletions

File tree

bitsandbytes/backends/utils.py

Lines changed: 10 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import subprocess
1+
import importlib
22

33
from packaging import version
44
import torch
@@ -67,17 +67,16 @@ def get_gaudi_sw_version():
6767
"""
6868
Returns the installed version of Gaudi SW.
6969
"""
70-
output = subprocess.run(
71-
"pip list | grep habana-torch-plugin",
72-
shell=True,
73-
text=True,
74-
capture_output=True,
75-
)
76-
# If grep return nothing
77-
if not output.stdout.strip():
78-
return None
70+
try:
71+
# if we find the spec, examine the installed version
72+
plugin_metadata = importlib.metadata.metadata("habana-torch-plugin")
73+
plugin_version = plugin_metadata.get("Version")
74+
if plugin_version:
75+
gaudi_version = version.parse(plugin_version)
76+
except Exception:
77+
gaudi_version = None
7978

80-
return version.parse(output.stdout.split("\n")[0].split()[-1])
79+
return gaudi_version
8180

8281

8382
GAUDI_SW_VER = get_gaudi_sw_version()

0 commit comments

Comments
 (0)