()
| 54 | |
| 55 | |
| 56 | def get_gpu_memory(): |
| 57 | memory_info = [] |
| 58 | gpu_memory_info = nvidia_smi_memory_info() |
| 59 | |
| 60 | try: |
| 61 | gpu_index = [int(k) for k in os.environ['CUDA_VISIBLE_DEVICES'].split(',')] |
| 62 | except KeyError: |
| 63 | gpu_index = [x["id"] for x in gpu_memory_info] |
| 64 | |
| 65 | for gpu_id, i in enumerate( gpu_index): |
| 66 | gpu = gpu_memory_info[i] |
| 67 | total_memory = gpu["total_memory"] |
| 68 | used_memory = gpu["used_memory"] |
| 69 | memory_info.append((gpu_id, total_memory, used_memory)) |
| 70 | return memory_info |
| 71 | |
| 72 | |
| 73 | def get_lowest_occupied_gpu(wait_memory=1000): |
no test coverage detected