/
githubmirror
/
faceswap
Обзор
Документация
Войти
/
githubmirror
/
faceswap
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
master
lib/system/ml_libs.py
957 строк
35 KB
torzdf
Initializers to Torch (#1549)
29 май 2026, 13:40
Не верифицирован
29 май 2026, 13:40
fa79d4d
Код
Авторство
О чём код?
#! /usr/env/bin/python """Queries information about system installed Machine Learning Libraries. NOTE: Only packages from Python's Standard Library should be imported in this module """ from __future__ import annotations import json import logging import os import platform import re import typing as T from abc import ABC, abstractmethod from shutil import which from lib.utils import get_module_objects from .system import _lines_from_command if platform.system() == "Windows": import winreg # pylint:disable=import-error else: winreg = None # type:ignore[assignment] # pylint:disable=invalid-name if T.TYPE_CHECKING: from winreg import HKEYType # type:ignore[attr-defined] logger = logging.getLogger(__name__) _TORCH_ROCM_REQUIREMENTS = {">=2.2.1,<2.4.0": ((6, 0), (6, 0))} """Minimum and maximum ROCm versions""" def _check_dynamic_linker(lib: str) -> list[str]: """Locate the folders that contain a given library in ldconfig and $LD_LIBRARY_PATH Parameters ---------- lib The library to locate Returns ------- All real existing folders from ldconfig or $LD_LIBRARY_PATH that contain the given lib """ paths: set[str] = set() ldconfig = which("ldconfig") if ldconfig: paths.update({os.path.realpath(os.path.dirname(line.split("=>")[-1].strip())) for line in _lines_from_command([ldconfig, "-p"]) if lib in line and "=>" in line}) if not os.environ.get("LD_LIBRARY_PATH"): return list(paths) paths.update({os.path.realpath(path) for path in os.environ["LD_LIBRARY_PATH"].split(":") if path and os.path.exists(path) for fname in os.listdir(path) if lib in fname}) return list(paths) def _files_from_folder(folder: str, prefix: str) -> list[str]: """Obtain all filenames from the given folder that start with the given prefix Parameters ---------- folder The folder to search for files in prefix The filename prefix to search for Returns ------- All filenames that exist in the given folder with the given prefix """ if not os.path.exists(folder): return [] return [f for f in os.listdir(folder) if f.startswith(prefix)] class _Alternatives: """Holds output from the update-alternatives command for the given package Parameters ---------- package The package to query update-alternatives for information """ def __init__(self, package: str) -> None: self._package = package self._bin = which("update-alternatives") self._default_marker = "link currently points to" self._alternatives_marker = "priority" self._output: list[str] | None = None @property def alternatives(self) -> list[str]: """Full path to alternatives listed for the given package""" if self._output is None: self._query() if not self._output: return [] retval = [line.rsplit(" - ", maxsplit=1)[0] for line in self._output if self._alternatives_marker in line.lower()] logger.debug("Versions from 'update-alternatives' for '%s': %s", self._package, retval) return retval @property def default(self) -> str: """Full path to the default package""" if self._output is None: self._query() if not self._output: return "" retval = next((x for x in self._output if x.startswith(self._default_marker)), "").replace(self._default_marker, "").strip() logger.debug("Default from update-alternatives for '%s': %s", self._package, retval) return retval def _query(self) -> None: """Query update-alternatives for the given package and place stripped output into :attr:`_output`""" if not self._bin: self._output = [] return cmd = [self._bin, "--display", self._package] retval = [line.strip() for line in _lines_from_command(cmd)] logger.debug("update-alternatives output for command %s: %s", cmd, retval) self._output = retval class _Cuda(ABC): """Find the location of system installed Cuda and cuDNN on Windows and Linux.""" def __init__(self) -> None: self.versions: list[tuple[int, int]] = [] """All detected globally installed Cuda versions""" self.version: tuple[int, int] = (0, 0) """Default installed Cuda version. (0, 0) if not detected""" self.cudnn_versions: dict[tuple[int, int], tuple[int, int, int]] = {} """Detected cuDNN version for each installed Cuda. key (0, 0) denotes globally installed cudnn""" self._paths: list[str] = [] """list of path to Cuda install folders relating to :attr:`versions`""" self._version_file = "version.json" self._lib = "libcudart.so" self._cudnn_header = "cudnn_version.h" self._alternatives = _Alternatives("cuda") self._re_cudnn = re.compile(r"#define CUDNN_(MAJOR|MINOR|PATCHLEVEL)\s+(\d+)") if platform.system() in ("Windows", "Linux"): self._get_versions() self._get_version() self._get_cudnn_versions() def __repr__(self) -> str: """Pretty representation of this class""" attrs = ", ".join(f"{k}={repr(v)}" for k, v in self.__dict__.items() if not k.startswith("_")) return f"{self.__class__.__name__}({attrs})" @classmethod def _tuple_from_string(cls, version: str) -> tuple[int, int] | None: """Convert a Cuda version string to a version tuple Parameters ---------- version The Cuda version string to convert Returns ------- The converted Cuda version string. ``None`` if not a valid version string """ if version.startswith("."): version = version[1:] split = version.split(".") if len(split) not in (2, 3): return None split = split[:2] if not all(x.isdigit() for x in split): return None return (int(split[0]), int(split[1])) @abstractmethod def get_versions(self) -> dict[tuple[int, int], str]: """Override to Attempt to detect all installed Cuda versions on Linux or Windows systems Returns ------- The Cuda versions to the folder path on the system """ @abstractmethod def get_version(self) -> tuple[int, int] | None: """Override to attempt to locate the default Cuda version on Linux or Windows Returns ------- The Default global Cuda version or ``None`` if not found """ @abstractmethod def get_cudnn_versions(self) -> dict[tuple[int, int], tuple[int, int, int]]: """Override to attempt to locate any installed cuDNN versions Returns ------- Detected cuDNN version for each installed Cuda. key (0, 0) denotes globally installed cudnn """ def version_from_version_file(self, folder: str) -> tuple[int, int] | None: """Attempt to get an installed Cuda version from its version.json file Parameters ---------- folder Full path to the folder to check for a version file Returns ------- The detected Cuda version or ``None`` if not detected """ vers_file = os.path.join(folder, self._version_file) if not os.path.exists(vers_file): return None with open(vers_file, "r", encoding="utf-8", errors="replace") as f: vers = json.load(f) retval = self._tuple_from_string(vers.get("cuda_cudart", {}).get("version")) logger.debug("Version from '%s': %s", vers_file, retval) return retval def _version_from_nvcc(self) -> tuple[int, int] | None: """Obtain the version from NVCC output if it is on PATH Returns ------- The detected default Cuda version. ``None`` if not version detected """ retval = None nvcc = which("nvcc") if not nvcc: return retval for line in _lines_from_command([nvcc, "-V"]): vers = re.match(r".*release (\d+\.\d+)", line) if vers is not None: retval = self._tuple_from_string(vers.group(1)) break logger.debug("Version from NVCC '%s': %s", nvcc, retval) return retval def _get_versions(self) -> None: """Attempt to detect all installed Cuda versions and populate to :attr:`versions`""" versions = self.get_versions() if versions: logger.debug("Cuda Versions: %s", versions) self.versions = list(versions) self._paths = list(versions.values()) return logger.debug("Could not locate any Cuda versions") def _get_version(self) -> None: """Attempt to detect the default Cuda version and populate to :attr:`version`""" version: tuple[int, int] | None = None if len(self.versions) == 1: version = self.versions[0] logger.debug("Only 1 installed Cuda version: %s", version) if not version: version = self._version_from_nvcc() if not version: version = self.get_version() if version: self.version = version logger.debug("Cuda version: %s", self.version if version else "not detected") def _get_cudnn_versions(self) -> None: """Attempt to locate any installed cuDNN versions and add to :attr`cudnn_versions`""" versions = self.get_cudnn_versions() if versions: logger.debug("cudnn versions: %s", versions) self.cudnn_versions = versions return logger.debug("No cudnn versions found") def cudnn_version_from_header(self, folder: str) -> tuple[int, int, int] | None: """Attempt to detect the cuDNN version from the version header file Parameters ---------- folder The folder to check for the cuDNN header file Returns ------- The cuDNN version found from the given folder or ``None`` if not detected """ path = os.path.join(folder, self._cudnn_header) if not os.path.exists(path): logger.debug("cudnn file '%s' does not exist", path) return None with open(path, "r", encoding="utf-8", errors="ignore") as f: file = f.read() version = {v[0]: int(v[1]) if v[1].isdigit() else 0 for v in self._re_cudnn.findall(file)} if not version: logger.debug("cudnn version could not be found in '%s'", path) return None logger.debug("cudnn version from '%s': %s", path, version) retval = (version.get("MAJOR", 0), version.get("MINOR", 0), version.get("PATCHLEVEL", 0)) logger.debug("cudnn versions: %s", retval) return retval class CudaLinux(_Cuda): """Find the location of system installed Cuda and cuDNN on Linux.""" def __init__(self) -> None: self._folder_prefix = "cuda-" super().__init__() def _version_from_lib(self, folder: str) -> tuple[int, int] | None: """Attempt to locate the version from the existence of libcudart.so within a Cuda targets/x86_64-linux/lib folder Parameters ---------- folder Full file path to the Cuda folder Returns ------- The Cuda version identified by the existence of the libcudart.so file. ``None`` if not detected """ lib_folder = os.path.join(folder, "targets", "x86_64-linux", "lib") lib_versions = [f.replace(self._lib, "") for f in _files_from_folder(lib_folder, self._lib)] if not lib_versions: return None versions = [self._tuple_from_string(f[1:]) for f in lib_versions if f and f.startswith(".")] valid = [v for v in versions if v is not None] if not valid or not len(set(valid)) == 1: return None retval = valid[0] logger.debug("Version from '%s': %s", os.path.join(lib_folder, self._lib), retval) return retval def _versions_from_usr(self) -> dict[tuple[int, int], str]: """Attempt to detect all installed Cuda versions from the /usr/local folder Scan /usr/local for cuda-x.x folders containing either a version.json file or include/lib/libcudart.so.x. Returns ------- A dictionary of detected Cuda versions to their install paths """ retval: dict[tuple[int, int], str] = {} usr = os.path.join(os.sep, "usr", "local") for folder in _files_from_folder(usr, self._folder_prefix): path = os.path.join(usr, folder) if os.path.islink(path): continue version = self.version_from_version_file(path) or self._version_from_lib(path) if version is not None: retval[version] = path return retval def _versions_from_alternatives(self) -> dict[tuple[int, int], str]: """Attempt to detect all installed Cuda versions from update-alternatives Returns ------- A dictionary of detected Cuda versions to their install paths found in update-alternatives """ retval: dict[tuple[int, int], str] = {} alts = self._alternatives.alternatives for path in alts: vers = self.version_from_version_file(path) or self._version_from_lib(path) if vers is not None: retval[vers] = path logger.debug("Versions from 'update-alternatives': %s", retval) return retval def _parent_from_targets(self, folder: str) -> str: """Obtain the Cuda parent folder from a path obtained from child targets folder Parameters ---------- folder Full path to a folder that has a 'targets' folder in its path Returns ------- The potential parent Cuda folder, or an empty string if not detected """ split = folder.split(os.sep) return os.sep.join(split[:split.index("targets")]) if "targets" in split else "" def _versions_from_dynamic_linker(self) -> dict[tuple[int, int], str]: """Attempt to detect all installed Cuda versions from ldconfig Returns ------- The Cuda version to the folder path found from ldconfig """ retval: dict[tuple[int, int], str] = {} folders = _check_dynamic_linker(self._lib) cuda_roots = [self._parent_from_targets(f) for f in folders] for path in cuda_roots: if not path: continue version = self.version_from_version_file(path) or self._version_from_lib(path) if version is not None: retval[version] = path logger.debug("Versions from 'ld_config': %s", retval) return retval def get_versions(self) -> dict[tuple[int, int], str]: """Attempt to detect all installed Cuda versions on Linux systems Returns ------- The Cuda version to the folder path on Linux """ versions = (self._versions_from_usr() | self._versions_from_alternatives() | self._versions_from_dynamic_linker()) return {k: versions[k] for k in sorted(versions)} def _version_from_alternatives(self) -> tuple[int, int] | None: """Attempt to get the default Cuda version from update-alternatives Returns ------- The detected default Cuda version. ``None`` if not version detected """ default = self._alternatives.default if not default: return None retval = self.version_from_version_file(default) or self._version_from_lib(default) logger.debug("Version from update-alternatives: %s", retval) return retval def _version_from_link(self) -> tuple[int, int] | None: """Attempt to get the default Cuda version from the /usr/local/cuda file Returns ------- The detected default Cuda version. ``None`` if not version detected """ path = os.path.join(os.sep, "usr", "local", "cuda") if not os.path.exists(path): return None real_path = os.path.abspath(os.path.realpath(path)) if os.path.islink(path) else path retval = self.version_from_version_file(real_path) or self._version_from_lib(real_path) logger.debug("Version from symlink: %s", retval) return retval def _version_from_dynamic_linker(self) -> tuple[int, int] | None: """Attempt to get the default version from ldconfig or $LD_LIBRARY_PATH Returns ------- The detected default ROCm version. ``None`` if not version detected """ paths = _check_dynamic_linker(self._lib) if len(paths) != 1: # Multiple or None return None root = self._parent_from_targets(paths[0]) retval = self.version_from_version_file(root) or self._version_from_lib(root) logger.debug("Version from ld_config: %s", retval) return retval def get_version(self) -> tuple[int, int] | None: """Attempt to locate the default Cuda version on Linux Checks, in order: update-alternatives, /usr/local/cuda, ldconfig, nvcc Returns ------- The Default global Cuda version or ``None`` if not found """ return (self._version_from_alternatives() or self._version_from_link() or self._version_from_dynamic_linker()) def get_cudnn_versions(self) -> dict[tuple[int, int], tuple[int, int, int]]: """Attempt to locate any installed cuDNN versions on Linux Returns ------- Detected cuDNN version for each installed Cuda. key (0, 0) denotes globally installed cudnn """ retval: dict[tuple[int, int], tuple[int, int, int]] = {} gbl = ["/usr/include", "/usr/local/include"] lcl = [os.path.join(f, "include") for f in self._paths] for root in gbl + lcl: for folder, _, filenames in os.walk(root): if self._cudnn_header not in filenames: continue version = self.cudnn_version_from_header(folder) if not version: continue cuda_vers = ((0, 0) if root in gbl else self.versions[self._paths.index(os.path.dirname(root))]) retval[cuda_vers] = version return retval class CudaWindows(_Cuda): """Find the location of system installed Cuda and cuDNN on Windows.""" @classmethod def _enum_sub_keys(cls, key: HKEYType) -> T.Generator[str, None, None]: """Iterate through a Registry key's sub-keys Parameters ---------- key The Registry key to iterate Yields ------ A sub-key name from the given registry key """ assert winreg is not None i = 0 while True: try: yield winreg.EnumKey(key, i) # type:ignore[attr-defined] except OSError: break i += 1 def get_versions(self) -> dict[tuple[int, int], str]: """Attempt to detect all installed Cuda versions on Windows systems from the registry Returns ------- The Cuda version to the folder path on Windows """ retval: dict[tuple[int, int], str] = {} assert winreg is not None reg_key = r"SOFTWARE\NVIDIA Corporation\GPU Computing Toolkit\CUDA" paths = {k.lower().replace("cuda_path_", "").replace("_", "."): v for k, v in os.environ.items() if "cuda_path_v" in k.lower()} try: with winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, # type:ignore[attr-defined] reg_key) as key: for version in self._enum_sub_keys(key): vers_tuple = self._tuple_from_string(version[1:]) if vers_tuple is not None: retval[vers_tuple] = paths.get(version, "") except FileNotFoundError: logger.debug("Could not find Windows Registry key '%s'", reg_key) return {k: retval[k] for k in sorted(retval)} def get_version(self) -> tuple[int, int] | None: """Attempt to get the default Cuda version from the Environment Variable Returns ------- The Default global Cuda version or ``None`` if not found """ path = os.environ.get("CUDA_PATH") if not path or path not in self._paths: return None retval = self.versions[self._paths.index(path)] logger.debug("Version from CUDA_PATH Environment Variable: %s", path) return retval def _get_cudnn_paths(self) -> list[str]: # noqa[C901] """Attempt to locate the locations of cuDNN installs for Windows Returns ------- Full path to existing cuDNN installs under Windows """ assert winreg is not None paths: set[str] = set() cudnn_key = "cudnn_cuda" reg_key = r"SOFTWARE\Microsoft\Windows\CurrentVersion\Uninstall" lookups = (winreg.HKEY_LOCAL_MACHINE, # type:ignore[attr-defined] winreg.HKEY_CURRENT_USER) # type:ignore[attr-defined] for lookup in lookups: try: key = winreg.OpenKey(lookup, reg_key) # type:ignore[attr-defined] except FileNotFoundError: continue for name in self._enum_sub_keys(key): if cudnn_key not in name.lower(): logger.debug("Skipping sub_keys '%s'", name) continue try: sub_keys = winreg.OpenKey(key, name) # type:ignore[attr-defined] logger.debug("Skipping sub_keys not found '%s'", name) except FileNotFoundError: continue logger.debug("Parsing cudnn key '%s'", cudnn_key) try: path, _ = winreg.QueryValueEx(sub_keys, # type:ignore[attr-defined] "InstallLocation") except (FileNotFoundError, OSError): logger.debug("Skipping missing InstallLocation for sub-key '%s'", sub_keys) continue if not os.path.isdir(path): logger.debug("Skipping non-existent path '%s'", path) continue paths.add(path) retval = list(paths) logger.debug("cudnn install paths: %s", retval) return retval def get_cudnn_versions(self) -> dict[tuple[int, int], tuple[int, int, int]]: """Attempt to locate any installed cuDNN versions on Windows Returns ------- Detected cuDNN version for each installed Cuda. key (0, 0) denotes globally installed cudnn """ retval: dict[tuple[int, int], tuple[int, int, int]] = {} gbl = self._get_cudnn_paths() lcl = [os.path.join(f, "include") for f in self._paths] for root in gbl + lcl: for folder, _, filenames in os.walk(root): if self._cudnn_header not in filenames: continue version = self.cudnn_version_from_header(folder) if not version: continue cuda_vers = ((0, 0) if root in gbl else self.versions[self._paths.index(os.path.dirname(root))]) retval[cuda_vers] = version return retval def get_cuda_finder() -> type[_Cuda]: """Create a platform-specific CUDA object. Returns ------- The OS specific finder for system-wide Cuda """ if platform.system().lower() == "windows": return CudaWindows return CudaLinux Cuda = get_cuda_finder() class ROCm(): """Find the location of system installed ROCm on Linux""" def __init__(self) -> None: self.version_min = min(v[0] for v in _TORCH_ROCM_REQUIREMENTS.values()) self.version_max = max(v[1] for v in _TORCH_ROCM_REQUIREMENTS.values()) self.versions: list[tuple[int, int, int]] = [] """All detected ROCm installed versions""" self.version: tuple[int, int, int] = (0, 0, 0) """Default ROCm installed version. (0, 0, 0) if not detected""" self._folder_prefix = "rocm-" self._version_files = ["version-rocm", "version"] self._lib = "librocm-core.so" self._alternatives = _Alternatives("rocm") self._re_version = re.compile(r"(\d+\.\d+\.\d+)(?=$|[-.])") self._re_config = re.compile(r"\sroc-(\d+\.\d+\.\d+)(?=\s|[-.])") if platform.system() == "Linux": self._rocm_check() def __repr__(self) -> str: """Pretty representation of this class""" attrs = ", ".join(f"{k}={repr(v)}" for k, v in self.__dict__.items() if not k.startswith("_")) return f"{self.__class__.__name__}({attrs})" @property def valid_versions(self) -> list[tuple[int, int, int]]: """Valid ROCm versions""" return [v for v in self.versions if self.version_min <= v[:2] <= self.version_max] @property def valid_installed(self) -> bool: """``True`` if a valid version of ROCm is installed""" return any(self.valid_versions) @property def is_valid(self): """``True`` if the default ROCm version is valid""" return self.version_min <= self.version[:2] <= self.version_max @classmethod def _tuple_from_string(cls, version: str) -> tuple[int, int, int] | None: """Convert a ROCm version string to a version tuple Parameters ---------- version The ROCm version string to convert Returns ------- The converted ROCm version string. ``None`` if not a valid version string """ split = version.split(".") if len(split) != 3: return None if not all(x.isdigit() for x in split): return None return (int(split[0]), int(split[1]), int(split[2])) def _version_from_string(self, string: str) -> tuple[int, int, int] | None: """Obtain the ROCm version from the end of a string Parameters ---------- string The string to test for a valid ROCm version Returns ------- The ROCm version from the end of the string or ``None`` if not detected """ re_vers = self._re_version.search(string) if re_vers is None: return None return self._tuple_from_string(re_vers.group(1)) def _version_from_info(self, folder: str) -> tuple[int, int, int] | None: """Attempt to locate the version from a version file within a ROCm .info folder Parameters ---------- file_path Full path to the ROCm .info folder Returns ------- The ROCm version extracted from a version file within the .info folder. ``None`` if not detected """ info_loc = [os.path.join(folder, ".info", v) for v in self._version_files] for info_file in info_loc: if not os.path.exists(info_file): continue with open(info_file, "r", encoding="utf-8") as f: vers_string = f.read().strip() if not vers_string: continue retval = self._tuple_from_string(vers_string.split("-", maxsplit=1)[0]) if retval is None: continue logger.debug("Version from '%s': %s", info_file, retval) return retval return None def _version_from_lib(self, folder: str) -> tuple[int, int, int] | None: """Attempt to locate the version from the existence of librocm-core.so within a ROCm lib folder Parameters ---------- folder Full file path to the ROCm folder Returns ------- The ROCm version identified by the existence of the librocm-core.so file. ``None`` if not detected """ lib_folder = os.path.join(folder, "lib") lib_files = _files_from_folder(lib_folder, self._lib) if not lib_files: return None # librocm-core naming is librocm-core.so.1.0.##### which is ambiguous. Get from folder rocm_folder = os.path.basename(folder) if not rocm_folder.startswith(self._folder_prefix): return None retval = self._version_from_string(rocm_folder) logger.debug("Version from '%s': %s", os.path.join(lib_folder, self._lib), retval) return retval def _versions_from_opt(self) -> list[tuple[int, int, int]]: """Attempt to detect all installed ROCm versions from the /opt folder Scan /opt for rocm.x.x.x folders containing either .info or lib/librocm-core.so.x Returns ------- Any ROCm versions found in the /opt folder """ retval: list[tuple[int, int, int]] = [] opt = os.path.join(os.sep, "opt") for folder in _files_from_folder(opt, self._folder_prefix): path = os.path.join(opt, folder) version = self._version_from_info(path) or self._version_from_lib(path) if version is not None: retval.append(version) return retval def _versions_from_alternatives(self) -> list[tuple[int, int, int]]: """Attempt to detect all installed ROCm versions from update-alternatives Returns ------- Any ROCm versions found in update-alternatives """ alts = self._alternatives.alternatives if not alts: return [] versions = [self._version_from_string(c) for c in alts] retval = list(set(v for v in versions if v is not None)) logger.debug("Versions from 'update-alternatives': %s", retval) return retval def _versions_from_dynamic_linker(self) -> list[tuple[int, int, int]]: """Attempt to detect all installed ROCm versions from ldconfig Returns ------- The ROCm versions found from ldconfig """ retval: list[tuple[int, int, int]] = [] folders = _check_dynamic_linker(self._lib) for folder in folders: path = os.path.dirname(folder) version = self._version_from_info(path) or self._version_from_lib(path) if version is not None: retval.append(version) logger.debug("Versions from 'ld_config': %s", retval) return retval def _get_versions(self) -> None: """Attempt to detect all installed ROCm versions and populate to :attr:`rocm_versions`""" versions = list(sorted(set(self._versions_from_opt()) | set(self._versions_from_alternatives()) | set(self._versions_from_dynamic_linker()))) if versions: logger.debug("ROCm Versions: %s", versions) self.versions = versions return logger.debug("Could not locate any ROCm versions") def _version_from_hipconfig(self) -> tuple[int, int, int] | None: """Attempt to get the default version from hipconfig Returns ------- The detected default ROCm version. ``None`` if not version detected """ retval: tuple[int, int, int] | None = None exe = which("hipconfig") if not exe: return retval lines = _lines_from_command([exe, "--full"]) if not lines: return retval for line in lines: line = line.strip() if line.startswith("ROCM_PATH"): path = line.split(":", maxsplit=1)[-1] retval = self._version_from_info(path) or self._version_from_lib(path) match = self._re_config.search(line) if match is not None: retval = self._tuple_from_string(match.group(1)) logger.debug("Version from hipconfig: %s", retval) return retval def _version_from_alternatives(self) -> tuple[int, int, int] | None: """Attempt to get the default version from update-alternatives Returns ------- The detected default ROCm version. ``None`` if not version detected """ default = self._alternatives.default if not default: return None retval = self._version_from_string(default.rsplit(os.sep, maxsplit=1)[-1]) logger.debug("Version from update-alternatives: %s", retval) return retval def _version_from_link(self) -> tuple[int, int, int] | None: """Attempt to get the default version from the /opt/rocm file Returns ------- The detected default ROCm version. ``None`` if not version detected """ path = os.path.join(os.sep, "opt", "rocm") if not os.path.exists(path): return None real_path = os.path.abspath(os.path.realpath(path)) if os.path.islink(path) else path retval = self._version_from_info(real_path) or self._version_from_lib(real_path) logger.debug("Version from symlink: %s", retval) return retval def _version_from_dynamic_linker(self) -> tuple[int, int, int] | None: """Attempt to get the default version from ldconfig or $LD_LIBRARY_PATH Returns ------- The detected default ROCm version. ``None`` if not version detected """ paths = _check_dynamic_linker("librocm-core.so.") if len(paths) != 1: # Multiple or None return None path = os.path.dirname(paths[0]) retval = self._version_from_info(path) or self._version_from_lib(path) logger.debug("Version from ld_config: %s", retval) return retval def _get_version(self) -> None: """Attempt to detect the default ROCm version""" version = (self._version_from_hipconfig() or self._version_from_alternatives() or self._version_from_link() or self._version_from_dynamic_linker()) if version is not None: logger.debug("ROCm default version: %s", version) self.version = version return logger.debug("Could not locate default ROCm version") def _rocm_check(self) -> None: """Attempt to locate the installed ROCm versions and the default ROCm version""" self._get_versions() self._get_version() logger.debug("ROCm Versions: %s, Version: %s", self.versions, self.version) __all__ = get_module_objects(__name__) if __name__ == "__main__": print(Cuda()) print(ROCm())