Skip to content

Commit da96644

Browse files
committed
Update Before and Action functions
Signed-off-by: Evan Lezar <[email protected]>
1 parent af81c92 commit da96644

File tree

7 files changed

+36
-31
lines changed

7 files changed

+36
-31
lines changed

api/config/v1/config.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ type Config struct {
4242
// NewConfig builds out a Config struct from a config file (or command line flags).
4343
// The data stored in the config will be populated in order of precedence from
4444
// (1) command line, (2) environment variable, (3) config file.
45-
func NewConfig(c *cli.Context, flags []cli.Flag) (*Config, error) {
45+
func NewConfig(c *cli.Command, flags []cli.Flag) (*Config, error) {
4646
config := &Config{Version: Version}
4747

4848
if configFile := c.String("config-file"); configFile != "" {

api/config/v1/flags.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ func ptr[T any](x T) *T {
2929
}
3030

3131
// updateFromCLIFlag conditionally updates the config flag at 'pflag' to the value of the CLI flag with name 'flagName'
32-
func updateFromCLIFlag[T any](pflag **T, c *cli.Context, flagName string) {
32+
func updateFromCLIFlag[T any](pflag **T, c *cli.Command, flagName string) {
3333
if c.IsSet(flagName) || *pflag == (*T)(nil) {
3434
switch flag := any(pflag).(type) {
3535
case **string:
@@ -112,7 +112,7 @@ type GFDCommandLineFlags struct {
112112
}
113113

114114
// UpdateFromCLIFlags updates Flags from settings in the cli Flags if they are set.
115-
func (f *Flags) UpdateFromCLIFlags(c *cli.Context, flags []cli.Flag) {
115+
func (f *Flags) UpdateFromCLIFlags(c *cli.Command, flags []cli.Flag) {
116116
for _, flag := range flags {
117117
for _, n := range flag.Names() {
118118
// Common flags

cmd/config-manager/main.go

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
package main
1818

1919
import (
20+
"context"
2021
"fmt"
2122
"os"
2223
"path/filepath"
@@ -117,11 +118,11 @@ func main() {
117118
flags := Flags{}
118119

119120
c := cli.Command{}
120-
c.Before = func(c *cli.Context) error {
121-
return validateFlags(c, &flags)
121+
c.Before = func(ctx context.Context, cmd *cli.Command) (context.Context, error) {
122+
return ctx, validateFlags(cmd, &flags)
122123
}
123-
c.Action = func(c *cli.Context) error {
124-
return start(c, &flags)
124+
c.Action = func(_ context.Context, cmd *cli.Command) error {
125+
return start(cmd, &flags)
125126
}
126127

127128
c.Flags = []cli.Flag{
@@ -203,14 +204,14 @@ func main() {
203204
},
204205
}
205206

206-
err := c.Run(os.Args)
207+
err := c.Run(context.Background(), os.Args)
207208
if err != nil {
208209
klog.Error(err)
209210
os.Exit(1)
210211
}
211212
}
212213

213-
func validateFlags(c *cli.Context, f *Flags) error {
214+
func validateFlags(c *cli.Command, f *Flags) error {
214215
if f.NodeName == "" {
215216
return fmt.Errorf("invalid <node-name>: must not be empty string")
216217
}
@@ -226,7 +227,7 @@ func validateFlags(c *cli.Context, f *Flags) error {
226227
return nil
227228
}
228229

229-
func start(c *cli.Context, f *Flags) error {
230+
func start(c *cli.Command, f *Flags) error {
230231
kubeconfig, err := clientcmd.BuildConfigFromFlags("", f.Kubeconfig)
231232
if err != nil {
232233
return fmt.Errorf("error building kubernetes clientcmd config: %s", err)

cmd/gpu-feature-discovery/main.go

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
package main
44

55
import (
6+
"context"
67
"encoding/json"
78
"fmt"
89
"os"
@@ -44,8 +45,8 @@ func main() {
4445
c.Name = "GPU Feature Discovery"
4546
c.Usage = "generate labels for NVIDIA devices"
4647
c.Version = info.GetVersionString()
47-
c.Action = func(ctx *cli.Context) error {
48-
return start(ctx, config)
48+
c.Action = func(_ context.Context, cmd *cli.Command) error {
49+
return start(cmd, config)
4950
}
5051

5152
config.flags = []cli.Flag{
@@ -122,7 +123,7 @@ func main() {
122123

123124
c.Flags = config.flags
124125

125-
if err := c.Run(os.Args); err != nil {
126+
if err := c.Run(context.Background(), os.Args); err != nil {
126127
klog.Error(err)
127128
os.Exit(1)
128129
}
@@ -141,7 +142,7 @@ func validateFlags(config *spec.Config) error {
141142
}
142143

143144
// loadConfig loads the config from the spec file.
144-
func (cfg *Config) loadConfig(c *cli.Context) (*spec.Config, error) {
145+
func (cfg *Config) loadConfig(c *cli.Command) (*spec.Config, error) {
145146
config, err := spec.NewConfig(c, cfg.flags)
146147
if err != nil {
147148
return nil, fmt.Errorf("unable to finalize config: %v", err)
@@ -154,7 +155,7 @@ func (cfg *Config) loadConfig(c *cli.Context) (*spec.Config, error) {
154155
return config, nil
155156
}
156157

157-
func start(c *cli.Context, cfg *Config) error {
158+
func start(c *cli.Command, cfg *Config) error {
158159
defer func() {
159160
klog.Info("Exiting")
160161
}()

cmd/mps-control-daemon/main.go

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
package main
1818

1919
import (
20+
"context"
2021
"encoding/json"
2122
"errors"
2223
"fmt"
@@ -54,8 +55,8 @@ func main() {
5455
c := cli.Command{}
5556
c.Name = "NVIDIA MPS Control Daemon"
5657
c.Version = info.GetVersionString()
57-
c.Action = func(ctx *cli.Context) error {
58-
return start(ctx, config)
58+
c.Action = func(_ context.Context, cmd *cli.Command) error {
59+
return start(cmd, config)
5960
}
6061
c.Commands = []*cli.Command{
6162
mount.NewCommand(),
@@ -78,7 +79,7 @@ func main() {
7879
c.Flags = config.flags
7980

8081
klog.InfoS(c.Name, "version", c.Version)
81-
err := c.Run(os.Args)
82+
err := c.Run(context.Background(), os.Args)
8283
if err != nil {
8384
klog.Error(err)
8485
os.Exit(1)
@@ -91,7 +92,7 @@ func validateFlags(config *spec.Config) error {
9192
}
9293

9394
// loadConfig loads the config from the spec file.
94-
func (cfg *Config) loadConfig(c *cli.Context) (*spec.Config, error) {
95+
func (cfg *Config) loadConfig(c *cli.Command) (*spec.Config, error) {
9596
config, err := spec.NewConfig(c, cfg.flags)
9697
if err != nil {
9798
return nil, fmt.Errorf("unable to finalize config: %w", err)
@@ -105,7 +106,7 @@ func (cfg *Config) loadConfig(c *cli.Context) (*spec.Config, error) {
105106
return config, nil
106107
}
107108

108-
func start(c *cli.Context, cfg *Config) error {
109+
func start(c *cli.Command, cfg *Config) error {
109110
klog.Info("Starting OS watcher.")
110111
sigs := watch.Signals(syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
111112
var started bool
@@ -161,7 +162,7 @@ exit:
161162
return nil
162163
}
163164

164-
func startDaemons(c *cli.Context, cfg *Config) ([]*mps.Daemon, bool, error) {
165+
func startDaemons(c *cli.Command, cfg *Config) ([]*mps.Daemon, bool, error) {
165166
// Load the configuration file
166167
klog.Info("Loading configuration.")
167168
config, err := cfg.loadConfig(c)

cmd/mps-control-daemon/mount/mount-shm.go

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ package mount
1818

1919
import (
2020
"bufio"
21+
"context"
2122
"fmt"
2223
"os"
2324
"os/exec"
@@ -41,7 +42,7 @@ func NewCommand() *cli.Command {
4142
}
4243

4344
// mountShm creates a tmpfs mount at /mps/shm to be used by the mps control daemon.
44-
func mountShm(c *cli.Context) error {
45+
func mountShm(_ context.Context, c *cli.Command) error {
4546
mountExecutable, err := exec.LookPath("mount")
4647
if err != nil {
4748
return fmt.Errorf("error finding 'mount' executable: %w", err)

cmd/nvidia-device-plugin/main.go

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
package main
1818

1919
import (
20+
"context"
2021
"encoding/json"
2122
"errors"
2223
"fmt"
@@ -53,8 +54,8 @@ func main() {
5354
c.Name = "NVIDIA Device Plugin"
5455
c.Usage = "NVIDIA device plugin for Kubernetes"
5556
c.Version = info.GetVersionString()
56-
c.Action = func(ctx *cli.Context) error {
57-
return start(ctx, o)
57+
c.Action = func(ctx context.Context, cmd *cli.Command) error {
58+
return start(ctx, cmd, o)
5859
}
5960

6061
c.Flags = []cli.Flag{
@@ -173,7 +174,7 @@ func main() {
173174
}
174175
o.flags = c.Flags
175176

176-
err := c.Run(os.Args)
177+
err := c.Run(context.Background(), os.Args)
177178
if err != nil {
178179
klog.Error(err)
179180
os.Exit(1)
@@ -227,7 +228,7 @@ func validateFlags(infolib nvinfo.Interface, config *spec.Config) error {
227228
return nil
228229
}
229230

230-
func loadConfig(c *cli.Context, flags []cli.Flag) (*spec.Config, error) {
231+
func loadConfig(c *cli.Command, flags []cli.Flag) (*spec.Config, error) {
231232
config, err := spec.NewConfig(c, flags)
232233
if err != nil {
233234
return nil, fmt.Errorf("unable to finalize config: %v", err)
@@ -236,8 +237,8 @@ func loadConfig(c *cli.Context, flags []cli.Flag) (*spec.Config, error) {
236237
return config, nil
237238
}
238239

239-
func start(c *cli.Context, o *options) error {
240-
klog.InfoS(fmt.Sprintf("Starting %s", c.App.Name), "version", c.App.Version)
240+
func start(ctx context.Context, c *cli.Command, o *options) error {
241+
klog.InfoS(fmt.Sprintf("Starting %s", c.Name), "version", c.Version)
241242

242243
kubeletSocketDir := filepath.Dir(o.kubeletSocket)
243244
klog.Infof("Starting FS watcher for %v", kubeletSocketDir)
@@ -263,7 +264,7 @@ restart:
263264
}
264265

265266
klog.Info("Starting Plugins.")
266-
plugins, restartPlugins, err := startPlugins(c, o)
267+
plugins, restartPlugins, err := startPlugins(ctx, c, o)
267268
if err != nil {
268269
return fmt.Errorf("error starting plugins: %v", err)
269270
}
@@ -317,7 +318,7 @@ exit:
317318
return nil
318319
}
319320

320-
func startPlugins(c *cli.Context, o *options) ([]plugin.Interface, bool, error) {
321+
func startPlugins(ctx context.Context, c *cli.Command, o *options) ([]plugin.Interface, bool, error) {
321322
// Load the configuration file
322323
klog.Info("Loading configuration.")
323324
config, err := loadConfig(c, o.flags)
@@ -359,7 +360,7 @@ func startPlugins(c *cli.Context, o *options) ([]plugin.Interface, bool, error)
359360

360361
// Get the set of plugins.
361362
klog.Info("Retrieving plugins.")
362-
plugins, err := GetPlugins(c.Context, infolib, nvmllib, devicelib, config)
363+
plugins, err := GetPlugins(ctx, infolib, nvmllib, devicelib, config)
363364
if err != nil {
364365
return nil, false, fmt.Errorf("error getting plugins: %v", err)
365366
}

0 commit comments

Comments
 (0)