Skip to content

Commit

Permalink
Add signal handler to handle termination more gracefully
Browse files Browse the repository at this point in the history
Change-Id: I537f81ab494cef7fe8ecc0b38ae091b8aaf60ac1
  • Loading branch information
y2kenny-amd committed Oct 9, 2024
1 parent 8162c06 commit a01b491
Showing 1 changed file with 10 additions and 0 deletions.
10 changes: 10 additions & 0 deletions internal/pkg/plugin/plugin.go
Original file line number Diff line number Diff line change
Expand Up @@ -21,9 +21,11 @@ import (
"bufio"
"fmt"
"os"
"os/signal"
"path/filepath"
"regexp"
"strconv"
"syscall"

"github.com/ROCm/k8s-device-plugin/internal/pkg/amdgpu"
"github.com/ROCm/k8s-device-plugin/internal/pkg/hwloc"
Expand All @@ -37,6 +39,7 @@ import (
type AMDGPUPlugin struct {
AMDGPUs map[string]map[string]int
Heartbeat chan bool
signal chan os.Signal
}

// Start is an optional interface that could be implemented by plugin.
Expand All @@ -45,6 +48,9 @@ type AMDGPUPlugin struct {
// method could be used to prepare resources before they are offered
// to Kubernetes.
func (p *AMDGPUPlugin) Start() error {
p.signal = make(chan os.Signal, 1)
signal.Notify(p.signal, syscall.SIGINT, syscall.SIGQUIT, syscall.SIGTERM)

return nil
}

Expand Down Expand Up @@ -187,6 +193,9 @@ func (p *AMDGPUPlugin) ListAndWatch(e *pluginapi.Empty, s pluginapi.DevicePlugin
devs[i].Health = health
}
s.Send(&pluginapi.ListAndWatchResponse{Devices: devs})
case <-p.signal:
glog.Infof("Received signal, exiting")
break
}
}
// returning a value with this function will unregister the plugin from k8s
Expand Down Expand Up @@ -245,6 +254,7 @@ func (p *AMDGPUPlugin) Allocate(ctx context.Context, r *pluginapi.AllocateReques
type AMDGPULister struct {
ResUpdateChan chan dpm.PluginNameList
Heartbeat chan bool
Signal chan os.Signal
}

// GetResourceNamespace must return namespace (vendor ID) of implemented Lister. e.g. for
Expand Down

0 comments on commit a01b491

Please sign in to comment.