mirror of
https://github.com/NVIDIA/nvidia-container-toolkit
synced 2024-11-22 08:18:32 +00:00
Use struct to store cri-o command line flags
Signed-off-by: Evan Lezar <elezar@nvidia.com>
This commit is contained in:
parent
9bbf7dcf96
commit
abe8ca71e0
@ -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)
|
||||||
}
|
}
|
||||||
|
Loading…
Reference in New Issue
Block a user