HAMi源码解析——HAMi-Core-1

原创

DifficultWork

修改于 2025-07-09 14:54:52

8380

文章被收录于专栏：阶梯计划阶梯计划

背景

这里不交代了，有机会再补。

本系列内容均基于HAMi开源项目v2.5.1的代码进行分析。

1 libvgpu源码结构

libvgpu.so（HAMi-Core）是 HAMi (异构计算资源管理系统) 中负责 GPU 虚拟化的核心动态库，它提供了 GPU 资源虚拟化、隔离和管理的底层功能。其核心是通过重写 dlsym 函数，以劫持 NVIDIA 动态链接库（如 CUDA 和 NVML）的调用，特别是针对以 cu 和 nvml 开头的函数进行拦截。

1.1 libvgpu.so 生效原理

device plugin 在 Allocate 方法中使用 hostPath 方式将宿主机上的 libvgpu.so 挂载到 Pod 中
通过 LD_PRELOAD 方式实现优先加载上一步中挂载的 libvgpu.so
通过注入环境变量 CUDA_DEVICE_MEMORY_LIMIT_X 和 CUDA_DEVICE_SM_LIMIT 指定 Memory 和 Core 的阈值

1.2 libvgpu.so 代码结构

src
├── cuda
│   ├── context.c
│   ├── device.c
│   ├── event.c
│   ├── graph.c
│   ├── hook.c
│   ├── memory.c
│   ├── stream.c
├── include
│   ├── cuda_addition_func.h
│   ├── libcuda_hook.h
│   ├── libnvml_hook.h
│   ├── libvgpu.h
│   ├── log_utils.h
│   ├── memory_limit.h
│   ├── multi_func_hook.h
│   ├── nvml_override.h
│   ├── nvml_prefix.h
│   ├── nvml-subset.h
│   ├── process_utils.h
│   ├── utils.h
├── multiprocess
│   ├── common_header.h
│   ├── multiprocess_memory_limit.c
│   ├── multiprocess_memory_limit.h
│   ├── multiprocess_utilization_watcher.c
│   ├── multiprocess_utilization_watcher.h
│   ├── shrreg_tool.c
├── nvml
│   ├── hook.c
│   ├── nvml_entry.c
├── libvgpu.c
├── utils.c

2 CUDA API 拦截

2.1 dlsym

dlsym 是一个用于符号解析的函数，声明在 dlfcn.h 头文件中，适用于 Linux 系统和其他符合 POSIX 标准的系统，允许程序在运行时动态地加载和使用共享库中的符号。

2.2 重写 dlsym

HAMi-core 中重写了 dlsym 函数，以劫持 NVIDIA 动态链接库（如 CUDA 和 NVML）的调用，特别是针对以 cu 和 nvml 开头的函数进行拦截。

src/libvgpu.c

// 拦截加载
FUNC_ATTR_VISIBLE void* dlsym(void* handle, const char* symbol) {
    LOG_DEBUG("into dlsym %s",symbol);
    pthread_once(&dlsym_init_flag,init_dlsym);
    // 先初始化dlsym（可以理解为CUDA Driver）
    if (real_dlsym == NULL) {
        real_dlsym = dlvsym(RTLD_NEXT,"dlsym","GLIBC_2.2.5");
        vgpulib = dlopen("/usr/local/vgpu/libvgpu.so",RTLD_LAZY);
        if (real_dlsym == NULL) {
            LOG_ERROR("real dlsym not found");
            real_dlsym = _dl_sym(RTLD_NEXT, "dlsym", dlsym);
            if (real_dlsym == NULL)
                LOG_ERROR("real dlsym not found");
        }
    }
    if (handle == RTLD_NEXT) {
        void *h = real_dlsym(RTLD_NEXT,symbol);
        int tid;
        pthread_mutex_lock(&dlsym_lock);
        tid = pthread_self();
        if (check_dlmap(tid,h)){
            LOG_WARN("recursive dlsym : %s\n",symbol);
            h = NULL;
        }
        pthread_mutex_unlock(&dlsym_lock);
        return h;
    }
    // 对cu前缀的进行特殊处理：使用 vgpulib 加载的进行替换
    if (symbol[0] == 'c' && symbol[1] == 'u') {
        //Compatible with cuda 12.8+ fix
        if (strcmp(symbol,"cuGetExportTable")!=0)
            pthread_once(&pre_cuinit_flag,(void(*)(void))preInit);
        void *f = real_dlsym(vgpulib,symbol);
        if (f!=NULL)
            return f;
    }
    // 对nvml前缀的进行特殊处理：使用 __dlsym_hook_section_nvml(handle, symbol)进行替换
    #ifdef HOOK_NVML_ENABLE
    if (symbol[0] == 'n' && symbol[1] == 'v' &&
          symbol[2] == 'm' && symbol[3] == 'l' ) {
        void* f = __dlsym_hook_section_nvml(handle, symbol);
        if (f != NULL) {
            return f;
        }
    }
#endif
    // 返回原驱动的API
    return real_dlsym(handle, symbol);
}

2.3 cu 前缀处理

先判断该 symbol 是否在vgpulib 加载的列表中，如果存在则替换为 vgpulib 加载的。

2.4 nvml 前缀处理

__dlsym_hook_section_nvml 定义了对于 nvml 开头的符号，具体处理如下：

void* __dlsym_hook_section_nvml(void* handle, const char* symbol) {
    // 一系列 DLSYM_HOOK_FUNC 宏定义调用（核心）
    DLSYM_HOOK_FUNC(nvmlInit);
    /** nvmlShutdown */
    DLSYM_HOOK_FUNC(nvmlShutdown);
    /** nvmlErrorString */
    DLSYM_HOOK_FUNC(nvmlErrorString);
    // ...
    /** nvmlVgpuTypeGetGpuInstanceProfileId */
    DLSYM_HOOK_FUNC(nvmlVgpuTypeGetGpuInstanceProfileId);
    // 如果没找到返回 NULL
    return NULL;
}

2.4.5 宏定义 DLSYM_HOOK_FUNC

#if defined(DLSYM_HOOK_DEBUG)
#define DLSYM_HOOK_FUNC(f)                                       \
    if (0 == strcmp(symbol, #f)) {                               \
        LOG_DEBUG("Detect dlsym for %s\n", #f);                  \
        return (void*) f; }                                      \

#define DLSYM_HOOK_FUNC_REPLACE(f)                               \
    if (0 == strcmp(symbol, hacked_#f)) {                        \
        return (void*) f; }                                      \

#else 

#define DLSYM_HOOK_FUNC(f)                                       \
    if (0 == strcmp(symbol, #f)) {                               \
        return (void*) f; }                                      \

#define DLSYM_HOOK_FUNC_REPLACE(f)                               \
    if (0 == strcmp(symbol, #f)) {                        \
        return (void*) hacked_##f; }                                      \

#endif

#f：这是宏中预处理器的一个特殊操作符，它将传入的参数 f 转换为字符串文字。例如，#f 将 cuGraphDestroy 转换成字符串 "cuGraphDestroy"。
strcmp(symbol, #f)：strcmp 是一个比较两个字符串的函数。如果 symbol 与 #f 字符串匹配（即 symbol 的值为 "cuGraphDestroy"），则 strcmp 返回 0。
return (void*) f;：如果 strcmp 返回 0，即 symbol 等于 f，则返回 f 对应的函数指针。(void*) f 将函数指针强制转换为 void* 类型，表示返回一个通用的函数指针。

2.5 hook.c 的实现

hook.c 通过 dlopen 和 dlsym 函数加载 CUDA 库并重定向 CUDA 库中的函数调用，以实现拦截、监控或修改 CUDA 函数的行为。

2.5.1 定义待拦截 CUDA 函数列表

通过cuda_library_entry 定义了哪些 CUDA 函数需要进行拦截，具体代码如下：

cuda_entry_t cuda_library_entry[] = {
    /* Init Part    */ 
    {.name = "cuInit"},
    /* Deivce Part */
    {.name = "cuDeviceGetAttribute"},
    {.name = "cuDeviceGet"},
    {.name = "cuDeviceGetCount"},
    {.name = "cuDeviceGetName"},
    // ...
    {.name = "cuGetProcAddress"},
    {.name = "cuGetProcAddress_v2"},
};

2.5.2 加载 CUDA API 地址

使用 load_cuda_libraries 函数获取各个 CUDA 函数的地址：

通过 dlopen 打开共享库 libcuda.so.1；
然后通过 real_dlsym 根据函数名称查询到对应的地址并存储到前面创建的 cuda_library_entry 中。

void load_cuda_libraries() {
    void *table = NULL;
    int i = 0;
    char cuda_filename[FILENAME_MAX];
    char tmpfunc[500];

    LOG_INFO("Start hijacking");

    snprintf(cuda_filename, FILENAME_MAX - 1, "%s","libcuda.so.1");
    cuda_filename[FILENAME_MAX - 1] = '\0';

    table = dlopen(cuda_filename, RTLD_NOW | RTLD_NODELETE);
    if (!table) {
        LOG_WARN("can't find library %s", cuda_filename);
    }

    for (i = 0; i < CUDA_ENTRY_END; i++) {
        LOG_DEBUG("LOADING %s %d",cuda_library_entry[i].name,i);
        cuda_library_entry[i].fn_ptr = real_dlsym(table, cuda_library_entry[i].name);
        if (!cuda_library_entry[i].fn_ptr) {
            cuda_library_entry[i].fn_ptr=real_dlsym(RTLD_NEXT,cuda_library_entry[i].name);
            if (!cuda_library_entry[i].fn_ptr){
                LOG_INFO("can't find function %s in %s", cuda_library_entry[i].name,cuda_filename);
                memset(tmpfunc,0,500);
                strcpy(tmpfunc,cuda_library_entry[i].name);
                while (prior_function(tmpfunc)) {
                    cuda_library_entry[i].fn_ptr=real_dlsym(RTLD_NEXT,tmpfunc);
                    if (cuda_library_entry[i].fn_ptr) {
                        LOG_INFO("found prior function %s",tmpfunc);
                        break;
                    } 
                }
            }
        }
    }
    LOG_INFO("loaded_cuda_libraries");
    if (cuda_library_entry[0].fn_ptr==NULL){
        LOG_WARN("is NULL");
    }
    dlclose(table);
}

2.6 libcuda_hook.h 的实现

libcuda_hook.h 是上一步拿到的 CUDA 函数后的拦截实现。

原始 CUDA 函数符号表 cuda_entry_t 就是 hook.c 中得到的：

typedef struct {
  void *fn_ptr;
  char *name;
} cuda_entry_t;

cuda_entry_t 这个结构体保存了 CUDA 函数的指针 (fn_ptr) 和函数名 (name)。通过将所有要拦截的 CUDA 函数存储在一个数组中，程序可以动态找到并调用这些函数。

函数枚举 cuda_override_enum_t cuda_override_enum_t 枚举了所有要拦截的 CUDA 函数。每个 CUDA 函数都有一个对应的枚举值，通过这个枚举可以索引到函数指针表中的具体函数。

typedef enum {
  CUDA_OVERRIDE_ENUM(cuInit),
  // ...
  CUDA_OVERRIDE_ENUM(cuGraphInstantiate),
  CUDA_ENTRY_END
} cuda_override_enum_t;

符号覆写宏：CUDA_OVERRIDE_ENUM

 // 将 x 前面加上 OVERRIDE_ 前缀，cuInit 会转换成 OVERRIDE_cuInit
#define CUDA_OVERRIDE_ENUM(x) OVERRIDE_##x

函数调用重定向宏：CUDA_OVERRIDE_CALL

#define CUDA_FIND_ENTRY(table, sym) ({ (table)[CUDA_OVERRIDE_ENUM(sym)].fn_ptr; })

#define CUDA_OVERRIDE_CALL(table, sym, ...)                                    \
  ({    \
    LOG_DEBUG("Hijacking %s", #sym);                                           \
    cuda_sym_t _entry = (cuda_sym_t)CUDA_FIND_ENTRY(table, sym);               \
    _entry(__VA_ARGS__);                                                       \
  })

CUDA_OVERRIDE_CALL 宏通过函数表中的函数指针来重定向 CUDA 函数调用：

CUDA_FIND_ENTRY 会根据传入的 sym（函数枚举）从 table 中找到对应的函数指针。 cuda_sym_t 定义为一个函数指针类型，用于调用 CUDA 函数。 _entry(VA_ARGS) 就是调用找到的 CUDA 函数，并传入参数。

这个宏在每次调用时都会输出日志，例如 LOG_DEBUG("Hijacking %s", #sym) 表示拦截了某个函数。就是根据函数枚举值，在 cuda_library_entry 中找到具体的函数地址，然后进行调用。CUDA_FIND_ENTRY 则在 table 根据名称查询对应的函数地址。

原创声明：本文系作者授权腾讯云开发者社区发表，未经许可，不得转载。

如有侵权，请联系 cloudcommunity@tencent.com 删除。

kubernetes

gpu

原创声明：本文系作者授权腾讯云开发者社区发表，未经许可，不得转载。

如有侵权，请联系 cloudcommunity@tencent.com 删除。

kubernetes

gpu

#HAMi

登录后参与评论

0 条评论

热度