bool load_nvml()

in source/neuropod/internal/cuda_device_mapping.cc [93:136]

• 33 lines of code
• 5 McCabe index (conditional complexity)

bool load_nvml()
{
    void *nvml_handle = nullptr;

    // NVML suffixes in priority order
    const auto nvml_suffixes = {"", ".1"};
    for (const std::string &suffix : nvml_suffixes)
    {
        const auto sopath = "libnvidia-ml.so" + suffix;
        SPDLOG_TRACE("Trying to load NVML: {}", sopath);
        nvml_handle = dlopen(sopath.c_str(), RTLD_LAZY);
        if (nvml_handle)
        {
            break;
        }
    }

    if (!nvml_handle)
    {
        // Couldn't load NVML
        SPDLOG_DEBUG("Neuropod could not load NVML");
        return false;
    }

    // Load the functions we care about
    nvmlInit = reinterpret_cast<nvmlReturn_t (*)()>(dlsym(nvml_handle, "nvmlInit"));

    nvmlDeviceGetHandleByPciBusId = reinterpret_cast<nvmlReturn_t (*)(const char *, nvmlDevice_t *)>(
        dlsym(nvml_handle, "nvmlDeviceGetHandleByPciBusId"));

    nvmlDeviceGetUUID =
        reinterpret_cast<nvmlReturn_t (*)(nvmlDevice_t, char *, unsigned int)>(dlsym(nvml_handle, "nvmlDeviceGetUUID"));

    nvmlErrorString = reinterpret_cast<const char *(*) (nvmlReturn_t)>(dlsym(nvml_handle, "nvmlErrorString"));

    auto err = nvmlInit();
    if (err != 0 /* NVML_SUCCESS */)
    {
        SPDLOG_ERROR("Error when initializing NVML: {}", nvmlErrorString(err));
        return false;
    }

    return true;
}