nvidia-k8s-device-plugin源码分析

1.引言

nvidia-k8s-device-plugin代码由go语言编写，在此确实要赞叹一下go语言的简洁和强大，想必以后会有越来越多的人喜欢上这门语言。

当然，如果想了解一个程序的代码，首先梳理一下每个文件的作用：

1.main.go：作为程序入口

2.nvidia.go：放置所有调用了nvml有关的函数代码

3.watcher.go：定义监视器的代码

4.server.go：实现与k8s-device-plugin有关流程的代码

在server.go中定义了NvidiaDevicePlugin 结构体，该结构体成员作用如下：

type NvidiaDevicePlugin struct {
	devs   []*pluginapi.Device   # api.protobuf里定义的一个数组，每个成员包括设备ID和其health信息
	socket string   # nvidia-device-plugin监听端口路径，实际为/var/lib/kubelet/device-plugins/nvidia.sock

	stop   chan interface{}  # 接受启停命令的管道
	health chan *pluginapi.Device  # 接受不健康设备的管道，发来pluginapi.Device的结构

	server *grpc.Server   # grcpserver，用来保存于kubelet的通讯
}

2.执行逻辑

main.go作为程序入口，首次执行代码逻辑如下。

1.首先加载nvml库，如果没有问题进行下一步，有问题则报错

log.Println("Loading NVML")
	if err := nvml.Init(); err != nil {
		log.Printf("Failed to initialize NVML: %s.", err)
		log.Printf("If this is a GPU node, did you set the docker default runtime to `nvidia`?")
		log.Printf("You can check the prerequisites at: https://github.com/NVIDIA/k8s-device-plugin#prerequisites")
		log.Printf("You can learn how to set the runtime at: https://github.com/NVIDIA/k8s-device-plugin#quick-start")

		select {}
	} 
	defer func() { log.Println("Shutdown of NVML returned:", nvml.Shutdown()) }()

2.获得当前宿主机设备数量，若为0则log出等待信息

log.Println("Fetching d  evices.")
	if len(getDevices()) == 0 {                       
		log.Println("No devices found. Waiting indefinitely.")
		select {}
	}

3.创建对于/var/lib/kubelet/device-plugins/文件夹的fsnotify监视器watcher，监视了所有的文件更改操作。

log.Println("Starting FS watcher.")
	watcher, err := newFSWatcher(pluginapi.DevicePluginPath) //constants.go->"/var/lib/kubelet/device-plugins/"，监视了所有的文件更改操作
	if err != nil {
		log.Println("Failed to created FS watcher.")
		os.Exit(1)
	}
	defer watcher.Close()

4.创建系统调用信号监视器sigs，监视系统调用信号

defer watcher.Close()

	log.Println("Starting OS watcher.")
	sigs := newOSWatcher(syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)  //监听信号，将系统的对应信号发送给sigs

5.监视deviceplugin的状态和系统信号，并作出相应反应

for循环L有两个功能块组成：

1）重启模块：

如果是第一次启动则创建新的NvidiaDevicePlugin结构体并填充信息，开启NvidiaDevicePlugin服务，否则停止之前的deviceplugin并重新创建

2）监视器模块：

针对watcher和sigs的传来不同信号的情况针对性处理，直至收到系统发来的停止信号则退出。

	restart := true
	var devicePlugin *NvidiaDevicePlugin

L:
	for {
		if restart {
			if devicePlugin != nil {
				devicePlugin.Stop()
			}
                    //如果还没有创建deviceplugin则创建，否则就停止原来的
			devicePlugin = NewNvidiaDevicePlugin() 
                    //返回一个结构体，里面包含NvidiaDevicePlugin{ devs，socket，stop，health}
			if err := devicePlugin.Serve(); err != nil {    
                    //开启NvidiaDevicePlugin的服务程序，并检查和kubelet的连通性，并
                    //开启健康监测，并向kubelet注册设备
				log.Println("Could not contact Kubelet, retrying. Did you enable the device plugin feature gate?")
				log.Printf("You can check the prerequisites at: https://github.com/NVIDIA/k8s-device-plugin#prerequisites")
				log.Printf("You can learn how to set the runtime at: https://github.com/NVIDIA/k8s-device-plugin#quick-start")
			} else {
				restart = false
			}
		}

		select {
		case event := <-watcher.Events:
			if event.Name == pluginapi.KubeletSocket && event.Op&fsnotify.Create == fsnotify.Create {
				log.Printf("inotify: %s created, restarting.", pluginapi.KubeletSocket)
				restart = true //若有重新创建的行为则重启
			}

		case err := <-watcher.Errors:  //出错则报错
			log.Printf("inotify: %s", err)

		case s := <-sigs:   //若有系统调用信号传来
			switch s {
			case syscall.SIGHUP:  //重启信号
				log.Println("Received SIGHUP, restarting.")
				restart = true
			default:   //其余信号都停止plugin服务
				log.Printf("Received signal \"%v\", shutting down.", s)
				devicePlugin.Stop()
				break L
			}
		}
	}

下面main.go中每个步骤中关键的函数：

分析之前我们先看一下main.go和nvidia.go同时引入的包pluginapi "k8s.io/kubernetes/pkg/kubelet/apis/deviceplugin/v1beta1"

，该路径下有一个api.pb.go和constants.go两个文件，包名同样为v1beta1，api.pb.go为grpc分析api.proto自动生成，constants.go中定义了很多接下来的要用到的常量，列举在这里

// \vendor\k8s.io\kubernetes\pkg\kubelet\apis\deviceplugin\v1beta1\constants.go
package v1beta1

const (
	// Healthy means that the device is healty
	Healthy = "Healthy"
	// UnHealthy means that the device is unhealthy
	Unhealthy = "Unhealthy"

	// Current version of the API supported by kubelet
	Version = "v1beta1"
	// DevicePluginPath is the folder the Device Plugin is expecting sockets to be on
	// Only privileged pods have access to this path
	// Note: Placeholder until we find a "standard path"
	DevicePluginPath = "/var/lib/kubelet/device-plugins/"
	// KubeletSocket is the path of the Kubelet registry socket
	KubeletSocket = DevicePluginPath + "kubelet.sock"
	// Timeout duration in secs for PreStartContainer RPC
	KubeletPreStartContainerRPCTimeoutInSecs = 30
)

var SupportedVersions = [...]string{"v1beta1"}

步骤1：

只有一个nvml.Init()，从字面意思可以知道是nvml进行了一些初始化操作。

步骤2：

1.getDevices()

// nvidia.go
func getDevices() []*pluginapi.Device {
	n, err := nvml.GetDeviceCount()
	check(err)

	var devs []*pluginapi.Device
	for i := uint(0); i < n; i++ {
		d, err := nvml.NewDeviceLite(i)
		check(err)
		devs = append(devs, &pluginapi.Device{
			ID:     d.UUID,
			Health: pluginapi.Healthy,
		})
	}

该函数定义在nvidia.go中，首先其调用了nvml.GetDeviceCount()获得当前宿主机设备数，将所有设备的信息加入devs数组，该数组每个成员是一个pluginapi.Device结构体，其ID被初始化为每个设备的UUID，Health字段初始化为"Healthy"（在constants.go中的const字段定义的Healthy = "Healthy"）

步骤3：

1.newFSWatcher(pluginapi.DevicePluginPath)

该函数定义在watchers.go中，其主要功能是创建一个监视pluginapi.DevicePluginPath路径下的文件变动的watcher并返回，从constants.go中的定义我们可以看到，其监视的路径为/var/lib/kubelet/device-plugins/，即同时监视了kubelet.sock和nvidia.sock

// watchers.go
func newFSWatcher(files ...string) (*fsnotify.Watcher, error) {
	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		return nil, err
	}

	for _, f := range files {
		err = watcher.Add(f)
		if err != nil {
			watcher.Close()
			return nil, err
		}
	}

	return watcher, nil
}

步骤4：

1.newOSWatcher(syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)

该函数同样定义在watchers.go中，其返回一个监视系统发来的SIGHUP、SIGINT、SIGTERM、SIGQUIT信号的watcher，该watcher实际上是一个只有一个缓存且成员为os.Signal的chan。main.go的L循环则监视该chan并做出相应的反应，

// watchers.go
func newOSWatcher(sigs ...os.Signal) chan os.Signal {
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, sigs...)   //sigs:syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT,监听

	return sigChan
}

步骤5：

1. devicePlugin.Stop()

定义在server.go中，停止grcp服务并清理现场。

// server.go
func (m *NvidiaDevicePlugin) Stop() error {
	if m.server == nil {
		return nil
	}

	m.server.Stop()
	m.server = nil
	close(m.stop)

	return m.cleanup()
}

2.NewNvidiaDevicePlugin()

返回一个结构体，里面包含NvidiaDevicePlugin{ devs，socket，stop，health}，devs是getDevices()返回的devs，socket是server.go中定义的常量serverSock = pluginapi.DevicePluginPath + "nvidia.sock"，即/var/lib/kubelet/device-plugins/nvidia.sock，stop是一个可以接受任何类型输入的无缓存chan，health是可以可以接受*pluginapi.Device类型输入的无缓存chan，其主要作用的及时将不健康的device报告给kubelet。 //待确定

//server.go
func NewNvidiaDevicePlugin() *NvidiaDevicePlugin {
	return &NvidiaDevicePlugin{
		devs:   getDevices(),    //改
		socket: serverSock,
		stop:   make(chan interface{}),
		health: make(chan *pluginapi.Device),
	}
}

3.devicePlugin.Serve()

该函数是main.go中最重要的函数，其负责开启NvidiaDevicePlugin的服务程序，并开启健康监测，并向kubelet注册设备。

首先其调用了NvidiaDevicePlugin结构体的Start()方法，该方法定义在nvidia.go中，作用为开启NvidiaDevicePlugin的服务程序，检查和kubelet的连通性，并开启健康监测。

然后再调用NvidiaDevicePlugin结构体的Register(pluginapi.KubeletSocket, resourceName)方法，两个参数pluginapi.KubeletSocket是constants.go中的常量，值为/var/lib/kubelet/device-plugins/kubelet.sock，而resourceName是该函数中定义的一个常量，为"nvidia.com/gpu"，结合k8splugin的流程我们很容易知道这个函数的作用就是将"nvidia.com/gpu"这个资源类型通过kubelet.sock注册到kubelet上。

// nvidia.go
func (m *NvidiaDevicePlugin) Serve() error {
	err := m.Start()    //开启NvidiaDevicePlugin的服务程序，并开启健康监测
	if err != nil {
		log.Printf("Could not start device plugin: %s", err)
		return err
	}
	log.Println("Starting to serve on", m.socket)

	err = m.Register(pluginapi.KubeletSocket, resourceName)
	if err != nil {
		log.Printf("Could not register device plugin: %s", err)
		m.Stop()
		return err
	}
	log.Println("Registered device plugin with Kubelet")

	return nil
}

接下来我们深入分析Start()方法和Register()方法，这两个是k8s-plugin流程的核心。

首先我们回忆一下kubernetes实现plugin要求设备厂商遵从的机制（参考https://www.kubernetes.org.cn/4391.html）：

1. 厂商自行实现一个管理设备资源的程序，部署到相应的节点上，我们称之为插件；
2. 插件需要向kubelet注册，注册内容包含自己的endpoint(endpoint就是一个用于通信的地址)、version、resource_name和DevicePluginOptions类型的options，而即这个类型仅有pre_start_required一个bool型成员，注释解释的很明白，表示PreStartContainer这个操作是否要在每个container启动前调用，其实就是启动容器前先通知插件做一下准备，多一些扩展性
3. kubelet连接插件的endpoint，就此kubelet和插件就建立了联系；
4. kubelet监听/var/lib/kubelet/device-plugins/kubelet.sock(unix sockets)这个地址，插件监听的也是类似的地址，只是地址变成了/var/lib/kubelet/device-plugins/nvidia.sock

这样我们便可以明白Start()方法和Register()方法到底在做些什么

3.1 Start()

首先其调用cleanup()方法删除/var/lib/kubelet/device-plugins/nvidia.sock，接下来进行标准的grpc调用操作，首先绑定服务端程序监听的sock（nvidia.sock），然后注册一个新的grpc_server对象,赋值到NvidiaDevicePlugin的server成员上，然后注册该grpc服务，之后用go关键字起一个独立的监听服务，然后测试一下服务是否正常工作，然后再用go关键字启动独立的健康检测程序。

// server.go
// Start starts the gRPC server of the device plugin
func (m *NvidiaDevicePlugin) Start() error {
	//NvidiaDevicePlugin开启自身的grpc服务端程序
	err := m.cleanup()   //删除文件夹下存在的nvidia.sock
	if err != nil {
		return err
	}

	sock, err := net.Listen("unix", m.socket)   //创建服务端程序监听的sock
	if err != nil {
		return err
	}
	m.server = grpc.NewServer([]grpc.ServerOption{}...)  //注册一个新的grpc_server
	pluginapi.RegisterDevicePluginServer(m.server, m) //将deviceplugin这种类型的grpc服务指定由NvidiaDevicePlugin实现

	go m.server.Serve(sock)   //创建独立服务监听

	// Wait for server to start by launching a blocking connexion
	conn, err := dial(m.socket, 5*time.Second)    //?创建一个 gRPC channel 和服务器交互连接，试一下是否服务器是否创建成功
	if err != nil {
		return err
	}  //出问题则报错
	conn.Close()     //关闭连接

	go m.healthcheck()   //开始健康监测

	return nil
}

3.1.1 healthcheck()

healthcheck()同样定义在server.go中，目前健康检查仅支持xids，其首先定义了xids，是一个成员为pluginapi.Device的chan，然后用go关键字调用watchXIDs(ctx, m.devs, xids)，最后用一个for循环select做检查，若NvidiaDevicePlugin中stop收到信息，则调用cancel()函数并返回，若xids中有信息了，则将调用m.unhealth(dev)将其放入m.health成员管道中。

// server.go
func (m *NvidiaDevicePlugin) healthcheck() {
	disableHealthChecks := strings.ToLower(os.Getenv(envDisableHealthChecks))
	if disableHealthChecks == "all" {
		disableHealthChecks = allHealthChecks   //目前健康检测仅支持xids
	}

	ctx, cancel := context.WithCancel(context.Background())

	var xids chan *pluginapi.Device
	if !strings.Contains(disableHealthChecks, "xids") {
		xids = make(chan *pluginapi.Device)
		go watchXIDs(ctx, m.devs, xids)
	}

	for {
		select {
		case <-m.stop:                    //取消健康检查
			cancel()               
			return
		case dev := <-xids:              
			m.unhealthy(dev)              //如果xids中有内容了，则将其中的设备加入m.healthy中
		}
	}
}

3.1.1.1 watchXIDs(ctx context.Context, devs []*pluginapi.Device, xids chan<- *pluginapi.Device)

首先其调用nvml.RegisterEventForDevice(eventSet, nvml.XidCriticalError, d.ID)为每个设备开启驱动端的健康监测，然后根据驱动返回的设备状态码决定是否要把不健康设备传入NvidiaDevicePlugin的health成员管道中。

// nvidia.go
func watchXIDs(ctx context.Context, devs []*pluginapi.Device, xids chan<- *pluginapi.Device) {
	eventSet := nvml.NewEventSet()
	defer nvml.DeleteEventSet(eventSet)

	for _, d := range devs {    //为每一个devs开启驱动端健康检测
		err := nvml.RegisterEventForDevice(eventSet, nvml.XidCriticalError, d.ID)
		if err != nil && strings.HasSuffix(err.Error(), "Not Supported") {
			log.Printf("Warning: %s is too old to support healthchecking: %s. Marking it unhealthy.", d.ID, err)

			xids <- d
			continue
		}

		if err != nil {
			log.Panicln("Fatal:", err)
		}
	}

	for {
		select {
		case <-ctx.Done():
			return
		default:
		}  //如果工作完成了就退出健康检查

		e, err := nvml.WaitForEvent(eventSet, 5000)
		if err != nil && e.Etype != nvml.XidCriticalError {
			continue
		}  //错误不是致命错误进行新一轮

		// FIXME: formalize the full list and document it.
		// http://docs.nvidia.com/deploy/xid-errors/index.html#topic_4
		// Application errors: the GPU should still be healthy
		if e.Edata == 31 || e.Edata == 43 || e.Edata == 45 {
			continue
		}//健康的
        
		if e.UUID == nil || len(*e.UUID) == 0 {
			// All devices are unhealthy，将所有的设备号都放入xid——channel中并进行下一轮
			for _, d := range devs {
				xids <- d
			}
			continue
		}
        //有错误将所有有错误的设备都放进去
		for _, d := range devs {
			if d.ID == *e.UUID {
				xids <- d
			}
		}
	}
}

3.2 Register(pluginapi.KubeletSocket, resourceName)

Register(kubeletEndpoint, resourceName string)函数首先通过kubelet.sock建立与kubelet的连接，然后调用kubelet的服务端预预先定义的GRPC方法Register(context.Background(), reqt)方法，将pluginapi.RegisterRequest类型的设备信息reqt注册至kubelet，包括：

reqt := &pluginapi.RegisterRequest{
		Version:      pluginapi.Version,  // v1beta1
		Endpoint:     path.Base(m.socket), // /var/lib/kubelet/device-plugins/nvidia.sock
		ResourceName: resourceName,  // nvidia.com/gpu
	}

// server.go
// Register registers the device plugin for the given resourceName with Kubelet.
func (m *NvidiaDevicePlugin) Register(kubeletEndpoint, resourceName string) error {
	conn, err := dial(kubeletEndpoint, 5*time.Second)  //与kubelet建立连接
	if err != nil {
		return err
	}
	defer conn.Close()

	client := pluginapi.NewRegistrationClient(conn)   //获得并注册远程调用的注册方法
	reqt := &pluginapi.RegisterRequest{
		Version:      pluginapi.Version,
		Endpoint:     path.Base(m.socket),
		ResourceName: resourceName,
	} 

	_, err = client.Register(context.Background(), reqt)   //将设备信息注册到kubelet中
	if err != nil {
		return err
	}
	return nil
}

以上便是Nvidia-Device-Plugin中向kubelet注册插件操作、健康检查程序和插件端服务监听的建立过程，那么我们可以发现，上述代码实现了如何让kubelet发现自己，而根据k8s-deviceplugin的流程和proto的定义，k8s要求插件端实现ListAndWatch、Allocate、GetDevicePluginOptions和PreStartContainer四个操作用来被kubelet操作，下面我们看一下这四个方法是如何实现的，当然，他们都定义在server.go中。

1.ListAndWatch()

函数中先将getDeviceCount()函数发现的本机GPU设备发送给kubelet，然后一个for循环，里面的select关键字表明其在监控两个chan，一个是stop信号，另一个是health信号，如果health通道被填入内容，则说明有设备处于不健康的状态，那么将调用Send函数将设备号报告给kubelet。此处有个注释是FIXME，内容是现阶段无法让失效设备恢复，在未来版本应该会改进这个问题

// server.go
func (m *NvidiaDevicePlugin) ListAndWatch(e *pluginapi.Empty, s pluginapi.DevicePlugin_ListAndWatchServer) error {
	s.Send(&pluginapi.ListAndWatchResponse{Devices: m.devs})
    //先将本机初始化检测到的所有健康设备发送给kublet
	for {
		select {
		case <-m.stop:
			return nil
		case d := <-m.health:
			// FIXME: there is no way to recover from the Unhealthy state.
			d.Health = pluginapi.Unhealthy
			s.Send(&pluginapi.ListAndWatchResponse{Devices: m.devs})
		}
	}
}

2.Allocate()

该函数接受kubelet传来的设备分配请求reqs，并返回要在容器中设置的设备的环境变量NVIDIA_VISIBLE_DEVICES的值，其值是由多个设备ID由 ‘,' 连接而成的，结合nvidia-docker的代码我们可以知道，k8s-plugin和nvidia-docker之间的交互是通过环境变量发生的，nvidia-docker中的libnvidia-container的prestarthook在运行时通过容器的环境变量设置来决定mount哪一个设备进入容器，这是一个系统的解决方案问题，所以mount这个动作发生在nvidia-docker而不是k8s这里也情有可原。

// server.go
// Allocate which return list of devices.传回环境变量与nvidia进行交互
func (m *NvidiaDevicePlugin) Allocate(ctx context.Context, reqs *pluginapi.AllocateRequest) (*pluginapi.AllocateResponse, error) {
	devs := m.devs
	responses := pluginapi.AllocateResponse{}
	for _, req := range reqs.ContainerRequests {
		response := pluginapi.ContainerAllocateResponse{
			Envs: map[string]string{
				"NVIDIA_VISIBLE_DEVICES": strings.Join(req.DevicesIDs, ","),
			},
		}

		for _, id := range req.DevicesIDs {
			if !deviceExists(devs, id) {
				return nil, fmt.Errorf("invalid allocation request: unknown device: %s", id)
			}
		}

		responses.ContainerResponses = append(responses.ContainerResponses, &response)
	}

	return &responses, nil
}

3.PreStartContainer()

什么都没有返回，现在应该是还没有实现这个函数的需求

// server.go
func (m *NvidiaDevicePlugin) PreStartContainer(context.Context, *pluginapi.PreStartContainerRequest) (*pluginapi.PreStartContainerResponse, error) {
	return &pluginapi.PreStartContainerResponse{}, nil
}

4.GetDevicePluginOptions()

同样什么都没有返回

// server.go
func (m *NvidiaDevicePlugin) GetDevicePluginOptions(context.Context, *pluginapi.Empty) (*pluginapi.DevicePluginOptions, error) {

	return &pluginapi.DevicePluginOptions{}, nil
}

3.总结

梳理一下nvidia-device-plugin的整体逻辑：

1.getCount()函数通过调用NVML接口，列出当前主机的所有设备信息，从而调用事先定义grpc函数Register()，来向kubelet注册自己。

2.kubelet通过grpc调用nvidia-device-plugin实现的Allocate()函数，向容器注入待分配设备的环境变量信息，容器创建时nvidia-docker通过调用libcontainer的prestartHook获取环境变量信息，在容器中挂载对应的设备。

3.通过调用nvml实现了一个健康检查程序，该健康检查程序负责监控本机GPU设备的健康情况，如果某个设备出了问题，该健康检查程序通过向NvidiaDevicePlugin的health通道发送不健康设备的信息，此时触发ListAndWatch函数的报告程序，将不健康设备健康状态更新为不健康，通过grpc更新设备健康列表并从kubelet.sock发送给kubelet，kubelet收到不健康设备信息后将处理所有受到该设备影响的pod。

4.目前健康检查仅支持xids，且没有支持重新恢复健康状态的设备的重新注册。

 

如果有错误之处还望指出。

4.参考资料

1.https://www.kubernetes.org.cn/4391.html

2.nvidia-k8s-device-plugin源码：Latest commit 2d56964 on 23 Aug

3.nvidia-docker源码