Use struct to store cri-o command line flags

Signed-off-by: Evan Lezar <elezar@nvidia.com>
This commit is contained in:
Evan Lezar 2022-06-24 07:03:04 +02:00
parent 9bbf7dcf96
commit abe8ca71e0

View File

@ -32,11 +32,16 @@ const (
defaultHookFilename = "oci-nvidia-hook.json" defaultHookFilename = "oci-nvidia-hook.json"
) )
var hooksDirFlag string // options stores the configuration from the command line or environment variables
var hookFilenameFlag string type options struct {
var tooklitDirArg string hooksDir string
hookFilename string
runtimeDir string
}
func main() { func main() {
options := options{}
// Create the top-level CLI // Create the top-level CLI
c := cli.NewApp() c := cli.NewApp()
c.Name = "crio" c.Name = "crio"
@ -49,15 +54,20 @@ func main() {
setup.Name = "setup" setup.Name = "setup"
setup.Usage = "Create the cri-o hook required to run NVIDIA GPU containers" setup.Usage = "Create the cri-o hook required to run NVIDIA GPU containers"
setup.ArgsUsage = "<toolkit_dirname>" setup.ArgsUsage = "<toolkit_dirname>"
setup.Action = Setup setup.Action = func(c *cli.Context) error {
setup.Before = ParseArgs return Setup(c, &options)
}
setup.Before = func(c *cli.Context) error {
return ParseArgs(c, &options)
}
// Create the 'cleanup' subcommand // Create the 'cleanup' subcommand
cleanup := cli.Command{} cleanup := cli.Command{}
cleanup.Name = "cleanup" cleanup.Name = "cleanup"
cleanup.Usage = "Remove the NVIDIA cri-o hook" cleanup.Usage = "Remove the NVIDIA cri-o hook"
cleanup.Action = Cleanup cleanup.Action = func(c *cli.Context) error {
return Cleanup(c, &options)
}
// Register the subcommands with the top-level CLI // Register the subcommands with the top-level CLI
c.Commands = []*cli.Command{ c.Commands = []*cli.Command{
&setup, &setup,
@ -74,7 +84,7 @@ func main() {
Aliases: []string{"d"}, Aliases: []string{"d"},
Usage: "path to the cri-o hooks directory", Usage: "path to the cri-o hooks directory",
Value: defaultHooksDir, Value: defaultHooksDir,
Destination: &hooksDirFlag, Destination: &options.hooksDir,
EnvVars: []string{"CRIO_HOOKS_DIR"}, EnvVars: []string{"CRIO_HOOKS_DIR"},
DefaultText: defaultHooksDir, DefaultText: defaultHooksDir,
}, },
@ -83,7 +93,7 @@ func main() {
Aliases: []string{"f"}, Aliases: []string{"f"},
Usage: "filename of the cri-o hook that will be created / removed in the hooks directory", Usage: "filename of the cri-o hook that will be created / removed in the hooks directory",
Value: defaultHookFilename, Value: defaultHookFilename,
Destination: &hookFilenameFlag, Destination: &options.hookFilename,
EnvVars: []string{"CRIO_HOOK_FILENAME"}, EnvVars: []string{"CRIO_HOOK_FILENAME"},
DefaultText: defaultHookFilename, DefaultText: defaultHookFilename,
}, },
@ -100,16 +110,16 @@ func main() {
} }
// Setup installs the prestart hook required to launch GPU-enabled containers // Setup installs the prestart hook required to launch GPU-enabled containers
func Setup(c *cli.Context) error { func Setup(c *cli.Context, o *options) error {
log.Infof("Starting 'setup' for %v", c.App.Name) log.Infof("Starting 'setup' for %v", c.App.Name)
err := os.MkdirAll(hooksDirFlag, 0755) err := os.MkdirAll(o.hooksDir, 0755)
if err != nil { if err != nil {
return fmt.Errorf("error creating hooks directory %v: %v", hooksDirFlag, err) return fmt.Errorf("error creating hooks directory %v: %v", o.hooksDir, err)
} }
hookPath := getHookPath(hooksDirFlag, hookFilenameFlag) hookPath := getHookPath(o.hooksDir, o.hookFilename)
err = createHook(tooklitDirArg, hookPath) err = createHook(o.runtimeDir, hookPath)
if err != nil { if err != nil {
return fmt.Errorf("error creating hook: %v", err) return fmt.Errorf("error creating hook: %v", err)
} }
@ -118,10 +128,10 @@ func Setup(c *cli.Context) error {
} }
// Cleanup removes the specified prestart hook // Cleanup removes the specified prestart hook
func Cleanup(c *cli.Context) error { func Cleanup(c *cli.Context, o *options) error {
log.Infof("Starting 'cleanup' for %v", c.App.Name) log.Infof("Starting 'cleanup' for %v", c.App.Name)
hookPath := getHookPath(hooksDirFlag, hookFilenameFlag) hookPath := getHookPath(o.hooksDir, o.hookFilename)
err := os.Remove(hookPath) err := os.Remove(hookPath)
if err != nil { if err != nil {
return fmt.Errorf("error removing hook '%v': %v", hookPath, err) return fmt.Errorf("error removing hook '%v': %v", hookPath, err)
@ -131,14 +141,14 @@ func Cleanup(c *cli.Context) error {
} }
// ParseArgs parses the command line arguments to the CLI // ParseArgs parses the command line arguments to the CLI
func ParseArgs(c *cli.Context) error { func ParseArgs(c *cli.Context, o *options) error {
args := c.Args() args := c.Args()
log.Infof("Parsing arguments: %v", args.Slice()) log.Infof("Parsing arguments: %v", args.Slice())
if c.NArg() != 1 { if c.NArg() != 1 {
return fmt.Errorf("incorrect number of arguments") return fmt.Errorf("incorrect number of arguments")
} }
tooklitDirArg = args.Get(0) o.runtimeDir = args.Get(0)
log.Infof("Successfully parsed arguments") log.Infof("Successfully parsed arguments")
return nil return nil
@ -152,7 +162,7 @@ func createHook(toolkitDir string, hookPath string) error {
defer hook.Close() defer hook.Close()
encoder := json.NewEncoder(hook) encoder := json.NewEncoder(hook)
err = encoder.Encode(generateOciHook(tooklitDirArg)) err = encoder.Encode(generateOciHook(toolkitDir))
if err != nil { if err != nil {
return fmt.Errorf("error writing hook file '%v': %v", hookPath, err) return fmt.Errorf("error writing hook file '%v': %v", hookPath, err)
} }