(currently_active_devices)
| 66 | |
| 67 | |
| 68 | def auto_pick_devices(currently_active_devices): |
| 69 | global mem_free_threshold |
| 70 | |
| 71 | torch_platform_name = get_installed_torch_platform()[0] |
| 72 | |
| 73 | if is_cpu_device(torch_platform_name): |
| 74 | return [torch_platform_name] |
| 75 | |
| 76 | device_count = get_device_count() |
| 77 | log.debug("Autoselecting GPU. Using most free memory.") |
| 78 | devices = [] |
| 79 | for device_id in range(device_count): |
| 80 | device_id = f"{torch_platform_name}:{device_id}" if device_count > 1 else torch_platform_name |
| 81 | device = get_device(device_id) |
| 82 | |
| 83 | mem_free, mem_total = mem_get_info(device) |
| 84 | mem_free /= float(10**9) |
| 85 | mem_total /= float(10**9) |
| 86 | device_name = get_device_name(device) |
| 87 | log.debug( |
| 88 | f"{device_id} detected: {device_name} - Memory (free/total): {round(mem_free, 2)}Gb / {round(mem_total, 2)}Gb" |
| 89 | ) |
| 90 | devices.append({"device": device_id, "device_name": device_name, "mem_free": mem_free}) |
| 91 | |
| 92 | devices.sort(key=lambda x: x["mem_free"], reverse=True) |
| 93 | max_mem_free = devices[0]["mem_free"] |
| 94 | curr_mem_free_threshold = COMPARABLE_GPU_PERCENTILE * max_mem_free |
| 95 | mem_free_threshold = max(curr_mem_free_threshold, mem_free_threshold) |
| 96 | |
| 97 | # Auto-pick algorithm: |
| 98 | # 1. Pick the top 75 percentile of the GPUs, sorted by free_mem. |
| 99 | # 2. Also include already-running devices (GPU-only), otherwise their free_mem will |
| 100 | # always be very low (since their VRAM contains the model). |
| 101 | # These already-running devices probably aren't terrible, since they were picked in the past. |
| 102 | # Worst case, the user can restart the program and that'll get rid of them. |
| 103 | devices = [ |
| 104 | x["device"] for x in devices if x["mem_free"] >= mem_free_threshold or x["device"] in currently_active_devices |
| 105 | ] |
| 106 | return devices |
| 107 | |
| 108 | |
| 109 | def device_init(context, device_id): |
no test coverage detected