Get memory usage for a specific GPU.
(gpu_id)
| 95 | |
| 96 | |
| 97 | def get_gpu_memory_usage(gpu_id): |
| 98 | """Get memory usage for a specific GPU.""" |
| 99 | try: |
| 100 | output = ( |
| 101 | subprocess.check_output( |
| 102 | f"nvidia-smi --query-gpu=memory.used --format=csv,nounits,noheader -i {gpu_id}", shell=True |
| 103 | ) |
| 104 | .decode("utf-8") |
| 105 | .strip() |
| 106 | ) |
| 107 | return int(output) |
| 108 | except subprocess.CalledProcessError: |
| 109 | print(f"Failed to get memory usage for GPU {gpu_id}") |
| 110 | return None |
| 111 | |
| 112 | |
| 113 | def get_free_gpu(): |