Get GPU type, GPU memory, and GPU count using nvidia-smi command.
()
| 214 | |
| 215 | |
| 216 | def _get_gpu_info(): |
| 217 | """ |
| 218 | Get GPU type, GPU memory, and GPU count using nvidia-smi command. |
| 219 | """ |
| 220 | try: |
| 221 | result = subprocess.run( |
| 222 | ["nvidia-smi", "--query-gpu=gpu_name,memory.total", "--format=csv,noheader,nounits"], |
| 223 | capture_output=True, |
| 224 | text=True, |
| 225 | check=True, |
| 226 | ) |
| 227 | gpu_lines = result.stdout.strip().split("\n") |
| 228 | gpu_count = len(gpu_lines) |
| 229 | gpu_info = [] |
| 230 | for line in gpu_lines: |
| 231 | gpu_name, gpu_memory = line.split(", ") |
| 232 | gpu_info.append( |
| 233 | { |
| 234 | "type": gpu_name, |
| 235 | "memory": float(gpu_memory) / 1024, # Convert to GB |
| 236 | } |
| 237 | ) |
| 238 | return gpu_count, gpu_info |
| 239 | except (subprocess.CalledProcessError, FileNotFoundError): |
| 240 | print("Failed to execute nvidia-smi command.") |
| 241 | return 0, [] |
| 242 | |
| 243 | |
| 244 | def _get_system_info(): |
no test coverage detected