/
githubmirror
/
faceswap
Обзор
Документация
Войти
/
githubmirror
/
faceswap
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
master
lib/gpu_stats/nvidia.py
202 строки
7 KB
torzdf
Backported code to support future updates
27 фев 2026, 22:30
27 фев 2026, 22:30
aa40efa
Код
Авторство
О чём код?
#!/usr/bin/env python3 """Collects and returns Information on available Nvidia GPUs. """ import os import pynvml # pylint:disable=import-error from lib.utils import FaceswapError, get_module_objects from ._base import _GPUStats, _EXCLUDE_DEVICES class NvidiaStats(_GPUStats): """Holds information and statistics about Nvidia GPU(s) available on the currently running system. Notes ----- PyNVML is used for hooking in to Nvidia's Machine Learning Library and allows for pulling fairly extensive statistics for Nvidia GPUs Parameters ---------- log Whether the class should output information to the logger. There may be occasions where the logger has not yet been set up when this class is queried. Attempting to log in these instances will raise an error. If GPU stats are being queried prior to the logger being available then this parameter should be set to ``False``. Otherwise set to ``True``. Default: ``True`` """ def _initialize(self) -> None: """Initialize PyNVML for Nvidia GPUs. If :attr:`_is_initialized` is ``True`` then this function just returns performing no action. Otherwise :attr:`is_initialized` is set to ``True`` after successfully initializing NVML. Raises ------ FaceswapError If the NVML library could not be successfully loaded """ if self._is_initialized: return try: self._log("debug", "Initializing PyNVML for Nvidia GPU.") pynvml.nvmlInit() except (pynvml.NVMLError_LibraryNotFound, # pylint:disable=no-member pynvml.NVMLError_DriverNotLoaded, # pylint:disable=no-member pynvml.NVMLError_NoPermission) as err: # pylint:disable=no-member msg = ("There was an error reading from the Nvidia Machine Learning Library. The most " "likely cause is incorrectly installed drivers. If this is the case, Please " "remove and reinstall your Nvidia drivers before reporting. Original " f"Error: {str(err)}") raise FaceswapError(msg) from err except Exception as err: # pylint:disable=broad-except msg = ("An unhandled exception occurred reading from the Nvidia Machine Learning " f"Library. Original error: {str(err)}") raise FaceswapError(msg) from err os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" super()._initialize() def _shutdown(self) -> None: """Cleanly close access to NVML and set :attr:`_is_initialized` back to ``False``. """ self._log("debug", "Shutting down NVML") pynvml.nvmlShutdown() super()._shutdown() def _get_device_count(self) -> int: """Detect the number of GPUs attached to the system. Returns ------- The total number of GPUs connected to the PC """ try: retval = pynvml.nvmlDeviceGetCount() except pynvml.NVMLError as err: self._log("debug", "Error obtaining device count. Setting to 0. " f"Original error: {str(err)}") retval = 0 self._log("debug", f"GPU Device count: {retval}") return retval def _get_active_devices(self) -> list[int]: """Obtain the indices of active GPUs (those that have not been explicitly excluded by CUDA_VISIBLE_DEVICES environment variable or explicitly excluded in the command line arguments). Returns ------- The list of device indices that are available for Faceswap to use """ # pylint:disable=duplicate-code devices = super()._get_active_devices() env_devices = os.environ.get("CUDA_VISIBLE_DEVICES") if env_devices: new_devices = [int(i) for i in env_devices.split(",")] devices = [idx for idx in devices if idx in new_devices] self._log("debug", f"Active GPU Devices: {devices}") return devices def _get_handles(self) -> list: """Obtain the device handles for all connected Nvidia GPUs. Returns ------- The list of pointers for connected Nvidia GPUs """ handles = [pynvml.nvmlDeviceGetHandleByIndex(i) for i in range(self._device_count)] self._log("debug", f"GPU Handles found: {len(handles)}") return handles def _get_driver(self) -> str: """Obtain the Nvidia driver version currently in use. Returns ------- The current GPU driver version """ try: driver = pynvml.nvmlSystemGetDriverVersion() except pynvml.NVMLError as err: self._log("debug", f"Unable to obtain driver. Original error: {str(err)}") driver = "No Nvidia driver found" self._log("debug", f"GPU Driver: {driver}") return driver def _get_device_names(self) -> list[str]: """Obtain the list of names of connected Nvidia GPUs as identified in :attr:`_handles`. Returns ------- The list of connected Nvidia GPU names """ names = [pynvml.nvmlDeviceGetName(handle) for handle in self._handles] self._log("debug", f"GPU Devices: {names}") return names def _get_vram(self) -> list[int]: """Obtain the VRAM in Megabytes for each connected Nvidia GPU as identified in :attr:`_handles`. Returns ------- The VRAM in Megabytes for each connected Nvidia GPU """ vram = [pynvml.nvmlDeviceGetMemoryInfo(handle).total / (1024 * 1024) for handle in self._handles] self._log("debug", f"GPU VRAM: {vram}") return vram def _get_free_vram(self) -> list[int]: """Obtain the amount of VRAM that is available, in Megabytes, for each connected Nvidia GPU. Returns ------- List of `float`s containing the amount of VRAM available, in Megabytes, for each connected GPU as corresponding to the values in :attr:`_handles """ is_initialized = self._is_initialized if not is_initialized: self._initialize() self._handles = self._get_handles() vram = [pynvml.nvmlDeviceGetMemoryInfo(handle).free / (1024 * 1024) for handle in self._handles] if not is_initialized: self._shutdown() self._log("debug", f"GPU VRAM free: {vram}") return vram def exclude_devices(self, devices: list[int]) -> None: """Exclude GPU devices from being used by Faceswap. Sets the CUDA_VISIBLE_DEVICES environment variable. This must be called before Torch/Keras are imported Parameters ---------- The GPU device IDS to be excluded """ # pylint:disable=duplicate-code if not devices: return self._log("debug", f"Excluding GPU indices: {devices}") _EXCLUDE_DEVICES.extend(devices) active = self._get_active_devices() os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(str(d) for d in active if d not in _EXCLUDE_DEVICES) env_vars = [f"{k}: {v}" for k, v in os.environ.items() if k.lower().startswith("cuda")] self._log("debug", f"Cuda environment variables: {env_vars}") __all__ = get_module_objects(__name__)