nvml_process_name.c

#include <stdio.h>
#include <string.h>
#include <unistd.h>
#include <cuda.h>
#include <nvml.h>

#define BUF_LEN 1024
#define MAX_PROCS 64

static int check_cu(CUresult r, const char *what) {
    if (r != CUDA_SUCCESS) {
        const char *msg = NULL;
        cuGetErrorString(r, &msg);
        fprintf(stderr, "%s failed: %s\n", what, msg ? msg : "unknown");
        return 1;
    }
    return 0;
}

static int check_nvml(nvmlReturn_t r, const char *what) {
    if (r != NVML_SUCCESS) {
        fprintf(stderr, "%s failed: %s\n", what, nvmlErrorString(r));
        return 1;
    }
    return 0;
}

int main(void) {
    // Step 1: get a CUDA context and allocate something on device 0 so
    // this process is registered with the NVIDIA driver as a GPU user.
    if (check_cu(cuInit(0), "cuInit")) return 1;

    CUdevice dev;
    if (check_cu(cuDeviceGet(&dev, 0), "cuDeviceGet")) return 1;

    CUcontext ctx;
    if (check_cu(cuDevicePrimaryCtxRetain(&ctx, dev), "cuDevicePrimaryCtxRetain")) return 1;
    if (check_cu(cuCtxSetCurrent(ctx), "cuCtxSetCurrent")) return 1;

    CUdeviceptr dptr;
    if (check_cu(cuMemAlloc(&dptr, 16), "cuMemAlloc")) return 1;

    // Step 2: initialise NVML and prime its PID->name cache by walking
    // the device's running-process list. Without this enumeration step,
    // nvmlSystemGetProcessName returns NVML_ERROR_NOT_FOUND even for a
    // PID that the driver clearly knows about.
    if (check_nvml(nvmlInit_v2(), "nvmlInit_v2")) return 1;

    nvmlDevice_t nvml_dev;
    if (check_nvml(nvmlDeviceGetHandleByIndex_v2(0, &nvml_dev), "nvmlDeviceGetHandleByIndex_v2")) {
        nvmlShutdown();
        return 1;
    }

    nvmlProcessInfo_t procs[MAX_PROCS];
    unsigned int count = MAX_PROCS;
    if (check_nvml(nvmlDeviceGetComputeRunningProcesses_v3(nvml_dev, &count, procs),
                   "nvmlDeviceGetComputeRunningProcesses_v3")) {
        nvmlShutdown();
        return 1;
    }

    // Step 3: now nvmlSystemGetProcessName will resolve our own PID.
    unsigned int pid = (unsigned int)getpid();
    char name[BUF_LEN];
    memset(name, 0, sizeof(name));

    nvmlReturn_t ret = nvmlSystemGetProcessName(pid, name, BUF_LEN);
    if (ret != NVML_SUCCESS) {
        fprintf(stderr, "nvmlSystemGetProcessName failed: %s\n", nvmlErrorString(ret));
        nvmlShutdown();
        cuMemFree(dptr);
        cuDevicePrimaryCtxRelease(dev);
        return 1;
    }

    size_t len = strnlen(name, BUF_LEN);
    printf("pid: %u\n", pid);
    printf("length: %zu\n", len);
    printf("bytes:");
    for (size_t i = 0; i < len; ++i) {
        printf(" %02x", (unsigned char)name[i]);
    }
    printf("\n");
    printf("string: %.*s\n", (int)len, name);

    nvmlShutdown();
    cuMemFree(dptr);
    cuDevicePrimaryCtxRelease(dev);
    return 0;
}