diff --git a/cuda_bindings/examples/4_CUDA_Libraries/nvidia_smi.py b/cuda_bindings/examples/4_CUDA_Libraries/nvidia_smi.py new file mode 100644 index 0000000000..8290e491c6 --- /dev/null +++ b/cuda_bindings/examples/4_CUDA_Libraries/nvidia_smi.py @@ -0,0 +1,241 @@ +# Copyright 2026 NVIDIA Corporation. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE + + +# ################################################################################ +# +# This example demonstrates the core cuda.bindings.nvml functionality by +# implementing a subset of the NVIDIA System Management Interface (nvidia-smi) +# command line tool in Python. +# +# ################################################################################ + + +import sys + +from cuda.bindings import nvml + +################################################################################## +# FORMATTING HELPERS + +# Utilities to help format the output table. See below for NVML usage. + + +def format_size(bytes_val: int) -> str: + """Formats bytes to MiB.""" + return f"{bytes_val / (1024 * 1024):.0f}MiB" + + +LINES = [[[4, 27, 6], [18, 3], [20]], [[4, 6, 13, 13], [22], [9, 10]]] + + +class TableFormatter: + def __init__(self, lines): + self.formats, self.sizes, self.counts = zip(*[self._create_line_format(line) for line in lines]) + + def _create_line_format(self, descriptor): + parts = [] + sizes = [] + for section in descriptor: + parts.append("| ") + sizes.append(1) + for i, align in enumerate(section): + if i == len(section) - 1: + direct = ">" + else: + direct = "<" + parts.append(f"{{:{direct}{align}}} ") + sizes[-1] += align + 1 + parts.append("|") + return "".join(parts), sizes, sum(len(x) for x in descriptor) + + def print_line(self, char="-"): + parts = ["+"] + for size in self.sizes[0]: + parts.append(char * size) + parts.append("+") + print("".join(parts)) + + def print_values(self, *args): + for line_format, count in zip(self.formats, self.counts): + print(line_format.format(*args[:count])) + args = args[count:] + + +def print_table(metadata, devices): + formatter = TableFormatter(LINES) + + print("+-----------------------------------------------------------------------------------------+") + print( + f"| NVIDIA-MINI-SMI {metadata['driver_version']:<16} Driver Version: {metadata['driver_version']:<15} CUDA Version: {metadata['cuda_version']:<9}|" + ) + formatter.print_line() + print("| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC |") + print("| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. |") + formatter.print_line("=") + + for device in devices: + formatter.print_values( + str(device["index"]), + device["name"], + device["persistence"], + device["bus_id"], + device["display_active"], + device["ecc_mode"], + device["fan_speed"], + device["temperature"], + device["performance_state"], + device["power"], + device["memory"], + device["utilization"], + device["compute_mode"], + ) + formatter.print_line() + + +################################################################################## +# NVML USAGE EXAMPLES + + +def collect_info(): + metadata = {} + metadata["driver_version"] = nvml.system_get_driver_version() + cuda_version_int = nvml.system_get_cuda_driver_version() + cuda_major = cuda_version_int // 1000 + cuda_minor = (cuda_version_int % 1000) // 10 + metadata["cuda_version"] = f"{cuda_major}.{cuda_minor}" + + devices = [] + + device_count = nvml.device_get_count_v2() + + for i in range(device_count): + device = {} + device["index"] = i + + handle = nvml.device_get_handle_by_index_v2(i) + + name = nvml.device_get_name(handle) + device["name"] = name + + try: + persistence = nvml.device_get_persistence_mode(handle) + persistence_str = "On" if persistence == nvml.EnableState.FEATURE_ENABLED else "Off" + except nvml.NvmlError: + persistence_str = "N/A" + device["persistence"] = persistence_str + + try: + pci_info = nvml.device_get_pci_info_v3(handle) + bus_id = pci_info.bus_id + except nvml.NvmlError: + bus_id = "N/A" + device["bus_id"] = bus_id + + try: + display_active = nvml.device_get_display_active(handle) + disp_str = "On" if display_active == nvml.EnableState.FEATURE_ENABLED else "Off" + except nvml.NvmlError: + disp_str = "N/A" + device["display_active"] = disp_str + + try: + current, _ = nvml.device_get_ecc_mode(handle) + ecc_str = "On" if current == nvml.EnableState.FEATURE_ENABLED else "Off" + except nvml.NvmlError: + ecc_str = "N/A" + device["ecc_mode"] = ecc_str + + try: + fan = nvml.device_get_fan_speed(handle) + fan_str = f"{fan: >3}%" + except nvml.NvmlError: + fan_str = "N/A" + device["fan_speed"] = fan_str + + try: + temp = nvml.device_get_temperature_v(handle, nvml.TemperatureSensors.TEMPERATURE_GPU) + temp_str = f"{temp}C" + except nvml.NvmlError: + temp_str = "N/A" + device["temperature"] = temp_str + + try: + perf_state = nvml.device_get_performance_state(handle) + perf_str = f"P{perf_state}" + except nvml.NvmlError: + perf_str = "N/A" + device["performance_state"] = perf_str + + try: + power_usage = nvml.device_get_power_usage(handle) # mW + usage_str = f"{power_usage // 1000}W" + except nvml.NvmlError: + usage_str = "N/A" + + try: + power_cap = nvml.device_get_power_management_limit(handle) # mW + cap_str = f"{power_cap // 1000}W" + except nvml.NvmlError: + cap_str = "N/A" + + pwr_str = f"{usage_str} / {cap_str}" + device["power"] = pwr_str + + try: + mem_info = nvml.device_get_memory_info_v2(handle) + except nvml.NvmlError: + mem_str = "N/A" + else: + mem_used = format_size(mem_info.used) + mem_total = format_size(mem_info.total) + mem_str = f"{mem_used} / {mem_total}" + + device["memory"] = mem_str + + try: + util_rates = nvml.device_get_utilization_rates(handle) + except nvml.NvmlError: + util_str = "N/A" + else: + gpu_util = util_rates.gpu + util_str = f"{gpu_util: >3}%" + + device["utilization"] = util_str + + try: + compute_mode = nvml.device_get_compute_mode(handle) + except nvml.NvmlError: + comp_str = "N/A" + else: + if compute_mode == nvml.ComputeMode.COMPUTEMODE_DEFAULT: + comp_str = "Default" + elif compute_mode == nvml.ComputeMode.COMPUTEMODE_EXCLUSIVE_PROCESS: + comp_str = "E. Process" + elif compute_mode == nvml.ComputeMode.COMPUTEMODE_PROHIBITED: + comp_str = "Prohibited" + else: + comp_str = "Unknown" + device["compute_mode"] = comp_str + + devices.append(device) + + return metadata, devices + + +def main(): + try: + nvml.init_v2() + except nvml.NvmlError as e: + print(f"Failed to initialize NVML: {e}") + sys.exit(1) + + try: + metadata, devices = collect_info() + print_table(metadata, devices) + finally: + nvml.shutdown() + + +if __name__ == "__main__": + main()