(num=4)
| 745 | |
| 746 | |
| 747 | def wait_available_gpu(num=4): |
| 748 | pynvml.nvmlInit() |
| 749 | gpuDeviceCount = pynvml.nvmlDeviceGetCount() |
| 750 | available_gpus = [] |
| 751 | while True: |
| 752 | available_gpus = [] |
| 753 | for i in range(gpuDeviceCount): |
| 754 | handle = pynvml.nvmlDeviceGetHandleByIndex(i) |
| 755 | pidAllInfo = pynvml.nvmlDeviceGetComputeRunningProcesses(handle) |
| 756 | if len(pidAllInfo) == 0: |
| 757 | available_gpus.append(i) |
| 758 | if len(available_gpus) >= num: |
| 759 | print("Found available GPUs: {}!".format(available_gpus)) |
| 760 | break |
| 761 | else: |
| 762 | print("Not enough GPUs...") |
| 763 | time.sleep(600) |
| 764 | string = "" |
| 765 | for i in available_gpus: |
| 766 | string += "," |
| 767 | string += str(i) |
| 768 | |
| 769 | os.environ["CUDA_VISIBLE_DEVICES"] = string[1:] |
| 770 | os.system("export CUDA_VISIBLE_DEVICES={}".format(string[1:])) |
| 771 | |
| 772 | |
| 773 | def infer_local_model(model, dataset, output_path, device_num, prefix = None, suffix = None, chat = False): |
nothing calls this directly
no outgoing calls
no test coverage detected