Invoke PreStart RPC call before container start, if desired by plugin

pkg/kubelet/apis/deviceplugin/v1beta1/BUILD

@@ -33,9 +33,3 @@ filegroup(
     srcs = [":package-srcs"],
     tags = ["automanaged"],
 )
-
-filegroup(
-    name = "go_default_library_protos",
-    srcs = ["api.proto"],
-    visibility = ["//visibility:public"],
-)

pkg/kubelet/apis/deviceplugin/v1beta1/api.proto

 // To regenerate api.pb.go run hack/update-device-plugin.sh
 syntax = 'proto3';
 
-package v1beta1;
+package deviceplugin;
 
 import "github.com/gogo/protobuf/gogoproto/gogo.proto";
 
@@ -24,6 +24,11 @@ service Registration {
 	rpc Register(RegisterRequest) returns (Empty) {}
 }
 
+message DevicePluginOptions {
+        // Indicates if PreStartContainer call is required before each container start
+	bool pre_start_required = 1;
+}
+
 message RegisterRequest {
 	// Version of the API the Device Plugin was built against
 	string version = 1;
@@ -32,6 +37,8 @@ message RegisterRequest {
 	string endpoint = 2;
 	// Schedulable resource name. As of now it's expected to be a DNS Label
 	string resource_name = 3;
+        // Options to be communicated with Device Manager
+        DevicePluginOptions options = 4;
 }
 
 message Empty {
@@ -48,6 +55,11 @@ service DevicePlugin {
 	// Plugin can run device specific operations and instruct Kubelet
 	// of the steps to make the Device available in the container
 	rpc Allocate(AllocateRequest) returns (AllocateResponse) {}
+
+        // PreStartContainer is called, if indicated by Device Plugin during registeration phase,
+        // before each container start. Device plugin can run device specific operations
+        // such as reseting the device before making devices available to the container
+	rpc PreStartContainer(PreStartContainerRequest) returns (PreStartContainerResponse) {}
 }
 
 // ListAndWatch returns a stream of List of Devices
@@ -71,6 +83,18 @@ message Device {
 	string health = 2;
 }
 
+// - PreStartContainer is expected to be called before each container start if indicated by plugin during registration phase.
+// - PreStartContainer allows kubelet to pass reinitialized devices to containers.
+// - PreStartContainer allows Device Plugin to run device specific operations on
+//   the Devices requested
+message PreStartContainerRequest {
+        repeated string devicesIDs = 1;
+}
+
+// PreStartContainerResponse will be send by plugin in response to PreStartContainerRequest
+message PreStartContainerResponse {
+}
+
 // - Allocate is expected to be called during pod creation since allocation
 //   failures for any container would result in pod startup failure.
 // - Allocate allows kubelet to exposes additional artifacts in a pod's

pkg/kubelet/apis/deviceplugin/v1beta1/constants.go

@@ -30,4 +30,6 @@ const (
 	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
 )

pkg/kubelet/cm/container_manager_linux.go

@@ -601,8 +601,10 @@ func (cm *containerManagerImpl) GetResources(pod *v1.Pod, container *v1.Containe
 	opts := &kubecontainer.RunContainerOptions{}
 	// Allocate should already be called during predicateAdmitHandler.Admit(),
 	// just try to fetch device runtime information from cached state here
-	devOpts := cm.deviceManager.GetDeviceRunContainerOptions(pod, container)
-	if devOpts == nil {
+	devOpts, err := cm.deviceManager.GetDeviceRunContainerOptions(pod, container)
+	if err != nil {
+		return nil, err
+	} else if devOpts == nil {
 		return opts, nil
 	}
 	opts.Devices = append(opts.Devices, devOpts.Devices...)

pkg/kubelet/cm/devicemanager/device_plugin_stub.go

@@ -105,7 +105,7 @@ func (m *Stub) Stop() error {
 }
 
 // Register registers the device plugin for the given resourceName with Kubelet.
-func (m *Stub) Register(kubeletEndpoint, resourceName string) error {
+func (m *Stub) Register(kubeletEndpoint, resourceName string, preStartContainerFlag bool) error {
 	conn, err := grpc.Dial(kubeletEndpoint, grpc.WithInsecure(), grpc.WithBlock(),
 		grpc.WithTimeout(10*time.Second),
 		grpc.WithDialer(func(addr string, timeout time.Duration) (net.Conn, error) {
@@ -120,6 +120,7 @@ func (m *Stub) Register(kubeletEndpoint, resourceName string) error {
 		Version:      pluginapi.Version,
 		Endpoint:     path.Base(m.socket),
 		ResourceName: resourceName,
+		Options:      &pluginapi.DevicePluginOptions{PreStartRequired: preStartContainerFlag},
 	}
 
 	_, err = client.Register(context.Background(), reqt)
@@ -129,6 +130,12 @@ func (m *Stub) Register(kubeletEndpoint, resourceName string) error {
 	return nil
 }
 
+// PreStartContainer resets the devices received
+func (m *Stub) PreStartContainer(ctx context.Context, r *pluginapi.PreStartContainerRequest) (*pluginapi.PreStartContainerResponse, error) {
+	log.Printf("PreStartContainer, %+v", r)
+	return &pluginapi.PreStartContainerResponse{}, nil
+}
+
 // ListAndWatch lists devices and update that list according to the Update call
 func (m *Stub) ListAndWatch(e *pluginapi.Empty, s pluginapi.DevicePlugin_ListAndWatchServer) error {
 	log.Println("ListAndWatch")

pkg/kubelet/cm/devicemanager/endpoint.go

@@ -36,6 +36,7 @@ type endpoint interface {
 	run()
 	stop()
 	allocate(devs []string) (*pluginapi.AllocateResponse, error)
+	preStartContainer(devs []string) (*pluginapi.PreStartContainerResponse, error)
 	getDevices() []pluginapi.Device
 	callback(resourceName string, added, updated, deleted []pluginapi.Device)
 }
@@ -65,8 +66,9 @@ func newEndpointImpl(socketPath, resourceName string, devices map[string]plugina
 		client:     client,
 		clientConn: c,
 
-		socketPath:   socketPath,
-		resourceName: resourceName,
+		socketPath:                      socketPath,
+		resourceName:                    resourceName,
+		invokePreStartContainerBoolFlag: false,
 
 		devices: devices,
 		cb:      callback,
@@ -182,6 +184,15 @@ func (e *endpointImpl) allocate(devs []string) (*pluginapi.AllocateResponse, err
 	})
 }
 
+// preStartContainer issues PreStartContainer gRPC call to the device plugin.
+func (e *endpointImpl) preStartContainer(devs []string) (*pluginapi.PreStartContainerResponse, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), pluginapi.KubeletPreStartContainerRPCTimeoutInSecs*time.Second)
+	defer cancel()
+	return e.client.PreStartContainer(ctx, &pluginapi.PreStartContainerRequest{
+		DevicesIDs: devs,
+	})
+}
+
 func (e *endpointImpl) stop() {
 	e.clientConn.Close()
 }

pkg/kubelet/cm/devicemanager/manager.go

@@ -85,6 +85,7 @@ type ManagerImpl struct {
 	// podDevices contains pod to allocated device mapping.
 	podDevices podDevices
 	store      utilstore.Store
+	pluginOpts map[string]*pluginapi.DevicePluginOptions
 }
 
 type sourcesReadyStub struct{}
@@ -112,6 +113,7 @@ func newManagerImpl(socketPath string) (*ManagerImpl, error) {
 		healthyDevices:   make(map[string]sets.String),
 		unhealthyDevices: make(map[string]sets.String),
 		allocatedDevices: make(map[string]sets.String),
+		pluginOpts:       make(map[string]*pluginapi.DevicePluginOptions),
 		podDevices:       make(podDevices),
 	}
 	manager.callback = manager.genericDeviceUpdateCallback
@@ -201,6 +203,7 @@ func (m *ManagerImpl) checkpointFile() string {
 // starts device plugin registration service.
 func (m *ManagerImpl) Start(activePods ActivePodsFunc, sourcesReady config.SourcesReady) error {
 	glog.V(2).Infof("Starting Device Plugin manager")
+	fmt.Println("Starting Device Plugin manager")
 
 	m.activePods = activePods
 	m.sourcesReady = sourcesReady
@@ -340,8 +343,10 @@ func (m *ManagerImpl) addEndpoint(r *pluginapi.RegisterRequest) {
 		glog.Errorf("Failed to dial device plugin with request %v: %v", r, err)
 		return
 	}
-
 	m.mutex.Lock()
+	if r.Options != nil {
+		m.pluginOpts[r.ResourceName] = r.Options
+	}
 	// Check for potential re-registration during the initialization of new endpoint,
 	// and skip updating if re-registration happens.
 	// TODO: simplify the part once we have a better way to handle registered devices
@@ -590,11 +595,7 @@ func (m *ManagerImpl) allocateContainerResources(pod *v1.Pod, container *v1.Cont
 		resource := string(k)
 		needed := int(v.Value())
 		glog.V(3).Infof("needs %d %s", needed, resource)
-		_, registeredResource := m.healthyDevices[resource]
-		_, allocatedResource := m.allocatedDevices[resource]
-		// Continues if this is neither an active device plugin resource nor
-		// a resource we have previously allocated.
-		if !registeredResource && !allocatedResource {
+		if !m.isDevicePluginResource(resource) {
 			continue
 		}
 		// Updates allocatedDevices to garbage collect any stranded resources
@@ -610,6 +611,7 @@ func (m *ManagerImpl) allocateContainerResources(pod *v1.Pod, container *v1.Cont
 		if allocDevices == nil || len(allocDevices) <= 0 {
 			continue
 		}
+
 		startRPCTime := time.Now()
 		// Manager.Allocate involves RPC calls to device plugin, which
 		// could be heavy-weight. Therefore we want to perform this operation outside
@@ -659,10 +661,60 @@ func (m *ManagerImpl) allocateContainerResources(pod *v1.Pod, container *v1.Cont
 // GetDeviceRunContainerOptions checks whether we have cached containerDevices
 // for the passed-in <pod, container> and returns its DeviceRunContainerOptions
 // for the found one. An empty struct is returned in case no cached state is found.
-func (m *ManagerImpl) GetDeviceRunContainerOptions(pod *v1.Pod, container *v1.Container) *DeviceRunContainerOptions {
+func (m *ManagerImpl) GetDeviceRunContainerOptions(pod *v1.Pod, container *v1.Container) (*DeviceRunContainerOptions, error) {
+	podUID := string(pod.UID)
+	contName := container.Name
+	for k := range container.Resources.Limits {
+		resource := string(k)
+		if !m.isDevicePluginResource(resource) {
+			continue
+		}
+		err := m.callPreStartContainerIfNeeded(podUID, contName, resource)
+		if err != nil {
+			return nil, err
+		}
+	}
 	m.mutex.Lock()
 	defer m.mutex.Unlock()
-	return m.podDevices.deviceRunContainerOptions(string(pod.UID), container.Name)
+	return m.podDevices.deviceRunContainerOptions(string(pod.UID), container.Name), nil
+}
+
+func (m *ManagerImpl) callPreStartContainerIfNeeded(podUID, contName, resource string) error {
+	m.mutex.Lock()
+	opts, ok := m.pluginOpts[resource]
+	if !ok {
+		m.mutex.Unlock()
+		glog.V(4).Infof("Plugin options not found in cache for resource: %s. Skip PreStartContainer", resource)
+		return nil
+	}
+
+	if !opts.PreStartRequired {
+		m.mutex.Unlock()
+		glog.V(4).Infof("Plugin options indicate to skip PreStartContainer for resource, %v", resource)
+		return nil
+	}
+
+	devices := m.podDevices.containerDevices(podUID, contName, resource)
+	if devices == nil {
+		m.mutex.Unlock()
+		return fmt.Errorf("no devices found allocated in local cache for pod %s, container %s, resource %s", podUID, contName, resource)
+	}
+
+	e, ok := m.endpoints[resource]
+	if !ok {
+		m.mutex.Unlock()
+		return fmt.Errorf("endpoint not found in cache for a registered resource: %s", resource)
+	}
+
+	m.mutex.Unlock()
+	devs := devices.UnsortedList()
+	glog.V(4).Infof("Issuing an PreStartContainer call for container, %s, of pod %s", contName, podUID)
+	_, err := e.preStartContainer(devs)
+	if err != nil {
+		return fmt.Errorf("device plugin PreStartContainer rpc failed with err: %v", err)
+	}
+	// TODO: Add metrics support for init RPC
+	return nil
 }
 
 // sanitizeNodeAllocatable scans through allocatedDevices in the device manager
@@ -692,3 +744,14 @@ func (m *ManagerImpl) sanitizeNodeAllocatable(node *schedulercache.NodeInfo) {
 		node.SetAllocatableResource(newAllocatableResource)
 	}
 }
+
+func (m *ManagerImpl) isDevicePluginResource(resource string) bool {
+	_, registeredResource := m.healthyDevices[resource]
+	_, allocatedResource := m.allocatedDevices[resource]
+	// Return true if this is either an active device plugin resource or
+	// a resource we have previously allocated.
+	if registeredResource || allocatedResource {
+		return true
+	}
+	return false
+}

pkg/kubelet/cm/devicemanager/manager_stub.go

@@ -53,8 +53,8 @@ func (h *ManagerStub) Allocate(node *schedulercache.NodeInfo, attrs *lifecycle.P
 }
 
 // GetDeviceRunContainerOptions simply returns nil.
-func (h *ManagerStub) GetDeviceRunContainerOptions(pod *v1.Pod, container *v1.Container) *DeviceRunContainerOptions {
-	return nil
+func (h *ManagerStub) GetDeviceRunContainerOptions(pod *v1.Pod, container *v1.Container) (*DeviceRunContainerOptions, error) {
+	return nil, nil
 }
 
 // GetCapacity simply returns nil capacity and empty removed resource list.

pkg/kubelet/cm/devicemanager/pod_devices.go

@@ -117,7 +117,9 @@ func (pdev podDevices) devices() map[string]sets.String {
 				if _, exists := ret[resource]; !exists {
 					ret[resource] = sets.NewString()
 				}
-				ret[resource] = ret[resource].Union(devices.deviceIds)
+				if devices.allocResp != nil {
+					ret[resource] = ret[resource].Union(devices.deviceIds)
+				}
 			}
 		}
 	}

pkg/kubelet/cm/devicemanager/types.go

@@ -51,7 +51,7 @@ type Manager interface {
 	// GetDeviceRunContainerOptions checks whether we have cached containerDevices
 	// for the passed-in <pod, container> and returns its DeviceRunContainerOptions
 	// for the found one. An empty struct is returned in case no cached state is found.
-	GetDeviceRunContainerOptions(pod *v1.Pod, container *v1.Container) *DeviceRunContainerOptions
+	GetDeviceRunContainerOptions(pod *v1.Pod, container *v1.Container) (*DeviceRunContainerOptions, error)
 
 	// GetCapacity returns the amount of available device plugin resource capacity, resource allocatable
 	// and inactive device plugin resources previously registered on the node.

test/e2e_node/device_plugin.go

@@ -75,7 +75,7 @@ var _ = framework.KubeDescribe("Device Plugin [Feature:DevicePlugin] [Serial] [D
 			framework.ExpectNoError(err)
 
 			By("Register resources")
-			err = dp1.Register(pluginapi.KubeletSocket, resourceName)
+			err = dp1.Register(pluginapi.KubeletSocket, resourceName, false)
 			framework.ExpectNoError(err)
 
 			By("Waiting for the resource exported by the stub device plugin to become available on the local node")
@@ -112,7 +112,7 @@ var _ = framework.KubeDescribe("Device Plugin [Feature:DevicePlugin] [Serial] [D
 			err = dp1.Start()
 			framework.ExpectNoError(err)
 
-			err = dp1.Register(pluginapi.KubeletSocket, resourceName)
+			err = dp1.Register(pluginapi.KubeletSocket, resourceName, false)
 			framework.ExpectNoError(err)
 
 			By("Waiting for resource to become available on the local node after re-registration")