diff --git a/noxfile.py b/noxfile.py index a4b11019186d58afacd456bec1a81c2af2cffd38..db8398b00cbf832a428e3dc3aa2f9acf3a318df5 100644 --- a/noxfile.py +++ b/noxfile.py @@ -4,21 +4,26 @@ from typing import Sequence import os import nox import subprocess +import re nox.options.sessions = ["lint", "typecheck", "testsuite"] def get_cuda_version() -> None | tuple[int, ...]: - smi_args = ["nvidia-smi", "--version"] - + query_args = ["nvcc", "--version"] + try: - result = subprocess.run(smi_args, capture_output=True) + query_result = subprocess.run(query_args, capture_output=True) except FileNotFoundError: return None - smi_output = str(result.stdout).splitlines() - cuda_version = smi_output[-1].split(":")[1].strip() - return tuple(int(v) for v in cuda_version.split(".")) + matches = re.findall(r"release \d+\.\d+", str(query_result.stdout)) + if matches: + match = matches[0] + version_string = match.split()[-1] + return tuple(int(v) for v in version_string.split(".")) + else: + return None def editable_install(session: nox.Session, opts: Sequence[str] = ()):