diff --git a/lollms/server/endpoints/lollms_hardware_infos.py b/lollms/server/endpoints/lollms_hardware_infos.py index 3c88924..55e6803 100644 --- a/lollms/server/endpoints/lollms_hardware_infos.py +++ b/lollms/server/endpoints/lollms_hardware_infos.py @@ -17,6 +17,7 @@ from pathlib import Path from typing import List import psutil import subprocess +from ascii_colors import trace_exception # ----------------------- Defining router and main class ------------------------------ router = APIRouter() @@ -64,31 +65,85 @@ def ram_usage(): "percent_usage":ram.percent, "ram_usage": ram.used } - -@router.get("/vram_usage") +@router.get("/vram_usage") # Or @app.route("/vram_usage", methods=['GET']) for Flask def vram_usage(): + """ + Fetches GPU VRAM usage information using nvidia-smi. + + Returns: + dict: A dictionary containing a 'gpus' key. The value is a list of + dictionaries, each representing a GPU with its VRAM details. + Returns {"gpus": []} if nvidia-smi fails or no GPUs are detected. + """ + gpus_list = [] try: - output = subprocess.check_output(['nvidia-smi', '--query-gpu=memory.total,memory.used,gpu_name', '--format=csv,nounits,noheader']) - lines = output.decode().strip().split('\n') - vram_info = [line.split(',') for line in lines] - except (subprocess.CalledProcessError, FileNotFoundError): - return { - "nb_gpus": 0 - } - - ram_usage = { - "nb_gpus": len(vram_info) - } - - if vram_info is not None: - for i, gpu in enumerate(vram_info): - ram_usage[f"gpu_{i}_total_vram"] = int(gpu[0])*1024*1024 - ram_usage[f"gpu_{i}_used_vram"] = int(gpu[1])*1024*1024 - ram_usage[f"gpu_{i}_model"] = gpu[2].strip() - else: - # Set all VRAM-related entries to None - ram_usage["gpu_0_total_vram"] = None - ram_usage["gpu_0_used_vram"] = None - ram_usage["gpu_0_model"] = None - - return ram_usage \ No newline at end of file + # Execute nvidia-smi to get total memory, used memory, and GPU name + # The output is in MiB (Mebibytes) + output = subprocess.check_output( + ['nvidia-smi', '--query-gpu=memory.total,memory.used,gpu_name', '--format=csv,nounits,noheader'], + encoding='utf-8' # Decode the output directly + ) + + # Split the output into lines, one line per GPU + lines = output.strip().split('\n') + + # Process each line + for line in lines: + if not line.strip(): # Skip potential empty lines + continue + + parts = line.split(',') + if len(parts) != 3: + print(f"Warning: Unexpected nvidia-smi output line format: {line}") + continue # Skip this malformed line + + try: + # Parse values (nvidia-smi returns MiB) + total_vram_mib = int(parts[0].strip()) + used_vram_mib = int(parts[1].strip()) + gpu_model = parts[2].strip() + + # Convert MiB to Bytes (as expected by filesize library) + total_vram_bytes = total_vram_mib * 1024 * 1024 + used_vram_bytes = used_vram_mib * 1024 * 1024 + + # Calculate available space and percentage + available_vram_bytes = total_vram_bytes - used_vram_bytes + percentage_used = 0.0 + if total_vram_bytes > 0: + # Calculate percentage and round to 2 decimal places + percentage_used = round((used_vram_bytes / total_vram_bytes) * 100, 2) + + # Create the dictionary structure expected by the frontend + gpu_info = { + "total_vram": total_vram_bytes, + "used_vram": used_vram_bytes, + "available_space": available_vram_bytes, # Frontend uses this in details + "gpu_model": gpu_model, + "percentage": percentage_used # Frontend uses this directly + } + gpus_list.append(gpu_info) + + except ValueError as ve: + print(f"Warning: Could not parse VRAM numbers from line: {line}. Error: {ve}") + continue # Skip this line if parsing fails + + except FileNotFoundError: + # nvidia-smi command not found (likely no NVIDIA drivers or not in PATH) + print("nvidia-smi command not found. Assuming no compatible GPU.") + # Return empty list as expected by frontend when no GPUs detected + gpus_list = [] + except subprocess.CalledProcessError as e: + # nvidia-smi command failed to execute properly + print(f"nvidia-smi failed with error code {e.returncode}: {e.output}") + # Return empty list + gpus_list = [] + except Exception as e: + # Catch any other unexpected errors during processing + trace_exception(e) + # Return empty list + gpus_list = [] + + # Always return the structure the frontend expects, even if the list is empty + # If using Flask, use jsonify: return jsonify({"gpus": gpus_list}) + return {"gpus": gpus_list} \ No newline at end of file