昇腾社区首页
中文
注册

选取昇腾AI处理器

任务说明

volcano框架根据节点优选得到分数后为Pod任务选择最优的节点,并将其绑定。在该过程中可以注册回调函数,依据亲和性策略,实现昇腾AI处理器选取。

例如Pod任务需要1个昇腾AI处理器,此时节点两个HCCL环分别剩余1个和3个昇腾AI处理器,那么最终会选择剩余1个昇腾AI处理器的环。

具体实现

具体代码实现请参考开源代码中UseAnnotation方法,其中selectNPUFromNode方法实现了根据亲和性从node上选取昇腾AI处理器的功能。

代码示例如下:
func (tp *module910x8) selectNPUFromNode(task *api.TaskInfo, node plugin.NPUNode) ([]int, error) {
    // 获取task所需NPU数量
    taskNPUNum, err := tp.GetTaskReqNPUNum(task)
    if err != nil {
        klog.V(util.LogErrorLev).Infof("%s ScoreBestNPUNodes err: %s", tp.GetPluginName(), err.Error())
        return nil, err
    }
 // 获取节点NPU数量
    nodeTop, err := tp.getUsableTopFromNode(node, len(tp.Tasks) > 1)
    if err != nil {
        klog.V(util.LogErrorLev).Infof("%s ScoreBestNPUNodes err: %s", tp.GetPluginName(), err.Error())
        return nil, err
    }
 // 如果是8昇腾AI处理器任务,直接全部选取
    if taskNPUNum == nodeNPUNumber {
        if len(nodeTop) == nodeNPUNumber {
            return nodeTop, nil
        }
        err = fmt.Errorf("node<%s> top<%v> can not meet task req<%d>", node.Name, nodeTop, taskNPUNum)
        klog.V(util.LogErrorLev).Infof("%s ScoreBestNPUNodes err: %s", tp.GetPluginName(), err.Error())
        return nil, err
    }
 // 获取任务的优先级数组
    priorityArray, err := getNPUAllocPriorityArray(taskNPUNum)
    if err != nil {
        klog.V(util.LogErrorLev).Info(err.Error())
        return nil, err
    }
    klog.V(util.LogInfoLev).Infof("%s selectNPUFromNode %s[%d] priority:%v in %v.", tp.GetPluginName(),
        task.Name, taskNPUNum, priorityArray, nodeTop)

 // 根据任务优先级数组选出昇腾AI处理器
    leftHccsArray, rightHccsArray := getNodeHccsArray(nodeTop)
    for _, priority := range priorityArray {
        if priority == len(leftHccsArray) {
            return leftHccsArray[:taskNPUNum], nil
        }
        if priority == len(rightHccsArray) {
            return rightHccsArray[:taskNPUNum], nil
        }
    }
    err = fmt.Errorf("node<%s> top<%v> can not meet task req<%d>", node.Name, len(nodeTop), taskNPUNum)
    klog.V(util.LogErrorLev).Infof("%s ScoreBestNPUNodes err: %s", tp.GetPluginName(), err.Error())
    return nil, err
}