UniqAI/nvidia-container-toolkit
1
0
Fork 0
mirror of https://github.com/NVIDIA/nvidia-container-toolkit synced 2025-06-26 18:18:24 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add runtimeDir as argument

Browse Source 
Thsi change adds the --nvidia-runtime-dir as a command line
argument when configuring container runtimes in the toolkit container.
This removes the need to set it via the command line.

Signed-off-by: Evan Lezar <elezar@nvidia.com>
This commit is contained in:
Evan Lezar
2023-06-09 16:50:51 +02:00
parent dcbf4b4f2a
commit 23041be511
3 changed files with 86 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files

49
tools/container/docker/docker.go
View File

@@ -54,19 +54,13 @@ const (
socketMessageToGetPID = "GET /info HTTP/1.0\r\n\r\n"
)
// nvidiaRuntimeBinaries defines a map of runtime names to binary names
var nvidiaRuntimeBinaries = map[string]string{
nvidiaRuntimeName: nvidiaRuntimeBinary,
nvidiaExperimentalRuntimeName: nvidiaExperimentalRuntimeBinary,
}
// options stores the configuration from the command line or environment variables
type options struct {
config string
socket string
runtimeName string
setAsDefault bool
runtimeDir string
setAsDefault bool
restartMode string
hostRootMount string
}
@@ -88,6 +82,9 @@ func main() {
setup.Action = func(c *cli.Context) error {
return Setup(c, &options)
}
setup.Before = func(c *cli.Context) error {
return ParseArgs(c, &options)
}
// Create the 'cleanup' subcommand
cleanup := cli.Command{}
@@ -97,6 +94,9 @@ func main() {
cleanup.Action = func(c *cli.Context) error {
return Cleanup(c, &options)
}
cleanup.Before = func(c *cli.Context) error {
return ParseArgs(c, &options)
}
// Register the subcommands with the top-level CLI
c.Commands = []*cli.Command{
@@ -148,6 +148,13 @@ func main() {
Destination: &options.runtimeName,
EnvVars: []string{"DOCKER_RUNTIME_NAME", "NVIDIA_RUNTIME_NAME"},
},
&cli.StringFlag{
Name: "nvidia-runtime-dir",
Aliases: []string{"runtime-dir"},
Usage: "The path where the nvidia-container-runtime binaries are located. If this is not specified, the first argument will be used instead",
Destination: &options.runtimeDir,
EnvVars: []string{"NVIDIA_RUNTIME_DIR"},
},
&cli.BoolFlag{
Name: "set-as-default",
Usage: "Set the `nvidia` runtime as the default runtime. If --runtime-name is specified as `nvidia-experimental` the experimental runtime is set as the default runtime instead",
@@ -173,12 +180,6 @@ func main() {
func Setup(c *cli.Context, o *options) error {
log.Infof("Starting 'setup' for %v", c.App.Name)
runtimeDir, err := ParseArgs(c)
if err != nil {
return fmt.Errorf("unable to parse args: %v", err)
}
o.runtimeDir = runtimeDir
cfg, err := docker.New(
docker.WithPath(o.config),
)
@@ -211,11 +212,6 @@ func Setup(c *cli.Context, o *options) error {
func Cleanup(c *cli.Context, o *options) error {
log.Infof("Starting 'cleanup' for %v", c.App.Name)
_, err := ParseArgs(c)
if err != nil {
return fmt.Errorf("unable to parse args: %v", err)
}
cfg, err := docker.New(
docker.WithPath(o.config),
)
@@ -248,17 +244,24 @@ func Cleanup(c *cli.Context, o *options) error {
}
// ParseArgs parses the command line arguments to the CLI
func ParseArgs(c *cli.Context) (string, error) {
func ParseArgs(c *cli.Context, o *options) error {
if o.runtimeDir != "" {
log.Debug("Runtime directory already set")
return nil
}
args := c.Args()
log.Infof("Parsing arguments: %v", args.Slice())
if args.Len() != 1 {
return "", fmt.Errorf("incorrect number of arguments")
if c.NArg() != 1 {
return fmt.Errorf("incorrect number of arguments")
}
runtimeDir := args.Get(0)
o.runtimeDir = args.Get(0)
log.Infof("Successfully parsed arguments")
return runtimeDir, nil
return nil
}
// UpdateConfig updates the docker config to include the nvidia runtimes