Skip to content

fix(virt-handler): re-register device plugins after kubelet restart #16

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
rohan-togetherai wants to merge 2 commits into
base: release-1.4.0/together
Choose a base branch
from
Open

fix(virt-handler): re-register device plugins after kubelet restart #16

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d4d4497
fix(virt-handler): re-register device plugins after kubelet restart
rohan-togetherai Aug 31, 2026
bb5f66d
test(virt-handler): cover device plugin restart lifecycle
rohan-togetherai Aug 31, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions pkg/virt-handler/device-manager/BUILD.bazel
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ go_test(
"device_controller_test.go",
"device_manager_suite_test.go",
"generic_device_test.go",
"kubelet_restart_test.go",
"mediated_device_test.go",
"mediated_devices_types_test.go",
"pci_device_test.go",
Expand All @@ -58,6 +59,7 @@ go_test(
"//staging/src/kubevirt.io/api/core/v1:go_default_library",
"//staging/src/kubevirt.io/client-go/log:go_default_library",
"//staging/src/kubevirt.io/client-go/testutils:go_default_library",
"//vendor/github.com/fsnotify/fsnotify:go_default_library",
"//vendor/github.com/onsi/ginkgo/v2:go_default_library",
"//vendor/github.com/onsi/gomega:go_default_library",
"//vendor/k8s.io/api/core/v1:go_default_library",
Expand All @@ -68,5 +70,6 @@ go_test(
"//vendor/k8s.io/client-go/testing:go_default_library",
"//vendor/k8s.io/client-go/tools/cache:go_default_library",
"//vendor/k8s.io/client-go/tools/cache/testing:go_default_library",
"@org_golang_google_grpc//:go_default_library",
],
)
19 changes: 19 additions & 0 deletions pkg/virt-handler/device-manager/common.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ package device_manager
import (
"bufio"
"bytes"
"context"
"fmt"
"net"
"os"
Expand Down Expand Up @@ -211,6 +212,24 @@ func waitForGRPCServer(socketPath string, timeout time.Duration) error {
return nil
}

func newDevicePluginGRPCServer() *grpc.Server {
return grpc.NewServer(grpc.WaitForHandlers(true))
}

func registrationContext(stop <-chan struct{}) (context.Context, context.CancelFunc) {
ctx, cancel := context.WithTimeout(context.Background(), connectionTimeout)
if stop != nil {
go func() {
select {
case <-stop:
cancel()
case <-ctx.Done():
}
}()
}
return ctx, cancel
}

// dial establishes the gRPC communication with the registered device plugin.
func gRPCConnect(socketPath string, timeout time.Duration) (*grpc.ClientConn, error) {
c, err := grpc.Dial(socketPath,
Expand Down
15 changes: 8 additions & 7 deletions pkg/virt-handler/device-manager/device_plugin_base.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -177,11 +177,9 @@ func (dpi *DevicePluginBase) Allocate(ctx context.Context, r *pluginapi.Allocate
}

func (dpi *DevicePluginBase) stopDevicePlugin() error {
defer func() {
if !IsChanClosed(dpi.done) {
close(dpi.done)
}
}()
if !IsChanClosed(dpi.done) {
close(dpi.done)
}
ticker := time.NewTicker(1 * time.Second)
defer ticker.Stop()
select {
Expand Down Expand Up @@ -214,7 +212,8 @@ func (dpi *DevicePluginBase) setInitialized(initialized bool) {
}

func (dpi *DevicePluginBase) register() error {
conn, err := gRPCConnect(pluginapi.KubeletSocket, connectionTimeout)
kubeletSocket := filepath.Join(filepath.Dir(dpi.socketPath), filepath.Base(pluginapi.KubeletSocket))
conn, err := gRPCConnect(kubeletSocket, connectionTimeout)
if err != nil {
return err
}
Expand All @@ -227,7 +226,9 @@ func (dpi *DevicePluginBase) register() error {
ResourceName: dpi.resourceName,
}

_, err = client.Register(context.Background(), reqt)
ctx, cancel := registrationContext(dpi.stop)
defer cancel()
_, err = client.Register(ctx, reqt)
if err != nil {
return err
}
Expand Down
62 changes: 36 additions & 26 deletions pkg/virt-handler/device-manager/generic_device.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -130,7 +130,7 @@ func (dpi *GenericDevicePlugin) Start(stop <-chan struct{}) (err error) {
return fmt.Errorf("error creating GRPC server socket: %v", err)
}

dpi.server = grpc.NewServer([]grpc.ServerOption{}...)
dpi.server = newDevicePluginGRPCServer()
defer dpi.stopDevicePlugin()

pluginapi.RegisterDevicePluginServer(dpi.server, dpi)
Expand All @@ -146,13 +146,27 @@ func (dpi *GenericDevicePlugin) Start(stop <-chan struct{}) (err error) {
return fmt.Errorf("error starting the GRPC server: %v", err)
}

socketDir := filepath.Dir(dpi.socketPath)
watcher, err := fsnotify.NewWatcher()
if err != nil {
return fmt.Errorf("failed to creating a fsnotify watcher: %v", err)
}
if err = watcher.Add(socketDir); err != nil {
_ = watcher.Close()
return fmt.Errorf("failed to add the device-plugin kubelet path to the watcher: %v", err)
}
if _, err = os.Stat(dpi.socketPath); err != nil {
_ = watcher.Close()
return fmt.Errorf("failed to stat the device-plugin socket: %v", err)
}
err = dpi.register()
if err != nil {
_ = watcher.Close()
return fmt.Errorf("error registering with device plugin manager: %v", err)
}

go func() {
errChan <- dpi.healthCheck()
errChan <- dpi.healthCheck(watcher)
}()

dpi.setInitialized(true)
Expand All @@ -164,11 +178,9 @@ func (dpi *GenericDevicePlugin) Start(stop <-chan struct{}) (err error) {

// Stop stops the gRPC server
func (dpi *GenericDevicePlugin) stopDevicePlugin() error {
defer func() {
if !IsChanClosed(dpi.done) {
close(dpi.done)
}
}()
if !IsChanClosed(dpi.done) {
close(dpi.done)
}

// Give the device plugin one second to properly deregister
ticker := time.NewTicker(1 * time.Second)
Expand All @@ -184,7 +196,8 @@ func (dpi *GenericDevicePlugin) stopDevicePlugin() error {

// Register registers the device plugin for the given resourceName with Kubelet.
func (dpi *GenericDevicePlugin) register() error {
conn, err := gRPCConnect(pluginapi.KubeletSocket, connectionTimeout)
kubeletSocket := filepath.Join(filepath.Dir(dpi.socketPath), filepath.Base(pluginapi.KubeletSocket))
conn, err := gRPCConnect(kubeletSocket, connectionTimeout)
if err != nil {
return err
}
Expand All @@ -197,7 +210,9 @@ func (dpi *GenericDevicePlugin) register() error {
ResourceName: dpi.resourceName,
}

_, err = client.Register(context.Background(), reqt)
ctx, cancel := registrationContext(dpi.stop)
defer cancel()
_, err = client.Register(ctx, reqt)
if err != nil {
return err
}
Expand Down Expand Up @@ -273,20 +288,16 @@ func (dpi *GenericDevicePlugin) PreStartContainer(_ context.Context, _ *pluginap
return res, nil
}

func (dpi *GenericDevicePlugin) healthCheck() error {
func (dpi *GenericDevicePlugin) healthCheck(watcher *fsnotify.Watcher) error {
logger := log.DefaultLogger()
watcher, err := fsnotify.NewWatcher()
if err != nil {
return fmt.Errorf("failed to creating a fsnotify watcher: %v", err)
}
defer watcher.Close()

// This way we don't have to mount /dev from the node
devicePath := filepath.Join(dpi.deviceRoot, dpi.devicePath)

// Start watching the files before we check for their existence to avoid races
dirName := filepath.Dir(devicePath)
err = watcher.Add(dirName)
err := watcher.Add(dirName)

if err != nil {
return fmt.Errorf("failed to add the device root path to the watcher: %v", err)
Expand All @@ -302,16 +313,8 @@ func (dpi *GenericDevicePlugin) healthCheck() error {
}
logger.Infof("device '%s' is present.", dpi.devicePath)

dirName = filepath.Dir(dpi.socketPath)
err = watcher.Add(dirName)

if err != nil {
return fmt.Errorf("failed to add the device-plugin kubelet path to the watcher: %v", err)
}
_, err = os.Stat(dpi.socketPath)
if err != nil {
return fmt.Errorf("failed to stat the device-plugin socket: %v", err)
}
socketDir := filepath.Dir(dpi.socketPath)
kubeletSocketPath := filepath.Join(socketDir, filepath.Base(pluginapi.KubeletSocket))

for {
select {
Expand All @@ -330,9 +333,16 @@ func (dpi *GenericDevicePlugin) healthCheck() error {
logger.Infof("monitored device %s disappeared", dpi.deviceName)
dpi.health <- deviceHealth{Health: pluginapi.Unhealthy}
}
} else if event.Name == dpi.socketPath && event.Op == fsnotify.Remove {
} else if event.Name == dpi.socketPath && event.Op.Has(fsnotify.Remove) {
logger.Infof("device socket file for device %s was removed, kubelet probably restarted.", dpi.deviceName)
return nil
} else if event.Name == kubeletSocketPath && event.Op.Has(fsnotify.Create) {
logger.Infof("kubelet socket %s was recreated, kubelet probably restarted. Restarting %s device plugin to re-register.", kubeletSocketPath, dpi.deviceName)
return nil
} else if event.Name == socketDir &&
(event.Op.Has(fsnotify.Remove) || event.Op.Has(fsnotify.Rename)) {
logger.Infof("device plugin socket directory %s was removed or replaced, kubelet probably restarted. Restarting %s device plugin to re-register.", socketDir, dpi.deviceName)
return nil
}
}
}
Expand Down
4 changes: 2 additions & 2 deletions pkg/virt-handler/device-manager/generic_device_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ var _ = Describe("Generic Device", func() {

errChan := make(chan error, 1)
go func(errChan chan error) {
errChan <- dpi.healthCheck()
errChan <- dpi.healthCheck(newTestHealthWatcher(dpi.socketPath))
}(errChan)
Consistently(func() string {
return dpi.devs[0].Health
Expand All @@ -63,7 +63,7 @@ var _ = Describe("Generic Device", func() {

os.OpenFile(dpi.socketPath, os.O_RDONLY|os.O_CREATE, 0666)

go dpi.healthCheck()
go dpi.healthCheck(newTestHealthWatcher(dpi.socketPath))
Expect(dpi.devs[0].Health).To(Equal(pluginapi.Healthy))

time.Sleep(1 * time.Second)
Expand Down
Loading