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）：

廠商自行實現一個管理設備資源的程序，部署到相應的節點上，我們稱之爲插件；
插件需要向kubelet註冊，註冊內容包含自己的endpoint(endpoint就是一個用於通信的地址)、version、resource_name和DevicePluginOptions類型的options，而即這個類型僅有pre_start_required一個bool型成員，註釋解釋的很明白，表示PreStartContainer這個操作是否要在每個container啓動前調用，其實就是啓動容器前先通知插件做一下準備，多一些擴展性
kubelet連接插件的endpoint，就此kubelet和插件就建立了聯繫；
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()

函數中僅有一個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})

	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源碼

nvidia-k8s-device-plugin源碼分析

1.引言

2.執行邏輯

3.總結

4.參考資料

HAProxy實現Harbor的多節點負載均衡

kubernetes device-plugin源碼分析

《企業級大數據平臺構建：架構與實現》閱讀總結

nvidia-k8s-device-plugin源碼分析

ubuntu16.04 apt方式安裝go語言環境

https://yachay.unat.edu.pe/blog/index.php?comment_area=format_blog&comment_component=blog&comment_co

linux以太網驅動總結