diff --git a/noxfile.py b/noxfile.py index aa9ad304382cdcf4d68e6073400ef76a8cc85b0f..d54815c43af42d3db393c0df87e912e2fe0f5a32 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] = ()):