Return the GPUs devices to manage based on the flags and resource scalars.
| 77 | // Return the GPUs devices to manage |
| 78 | // based on the flags and resource scalars. |
| 79 | static Try<set<Gpu>> enumerateGpus( |
| 80 | const Flags& flags, |
| 81 | const Resources& resources) |
| 82 | { |
| 83 | vector<unsigned int> indices; |
| 84 | |
| 85 | if (flags.nvidia_gpu_devices.isSome()) { |
| 86 | indices = flags.nvidia_gpu_devices.get(); |
| 87 | } else { |
| 88 | for (size_t i = 0; i < resources.gpus().getOrElse(0); ++i) { |
| 89 | indices.push_back(i); |
| 90 | } |
| 91 | } |
| 92 | |
| 93 | set<Gpu> gpus; |
| 94 | |
| 95 | foreach (unsigned int index, indices) { |
| 96 | Try<nvmlDevice_t> handle = nvml::deviceGetHandleByIndex(index); |
| 97 | if (handle.isError()) { |
| 98 | return Error("Failed to nvml::deviceGetHandleByIndex: " + handle.error()); |
| 99 | } |
| 100 | |
| 101 | Try<unsigned int> minor = nvml::deviceGetMinorNumber(handle.get()); |
| 102 | if (minor.isError()) { |
| 103 | return Error("Failed to nvml::deviceGetMinorNumber: " + minor.error()); |
| 104 | } |
| 105 | |
| 106 | Gpu gpu; |
| 107 | gpu.major = NVIDIA_MAJOR_DEVICE; |
| 108 | gpu.minor = minor.get(); |
| 109 | |
| 110 | gpus.insert(gpu); |
| 111 | } |
| 112 | |
| 113 | return gpus; |
| 114 | } |
| 115 | |
| 116 | |
| 117 | // To determine the proper number of GPU resources to return, we |