From 94337b74278a7549287fb15d7c42b740c0b8f8df Mon Sep 17 00:00:00 2001
From: Evan Lezar <elezar@nvidia.com>
Date: Fri, 27 Sep 2024 11:09:24 +0200
Subject: [PATCH] Add runtime package for runtime setup

Signed-off-by: Evan Lezar <elezar@nvidia.com>
---
 tools/container/nvidia-toolkit/run.go |   7 +-
 tools/container/runtime/runtime.go    | 108 ++++++++++++++++++++++++++
 2 files changed, 114 insertions(+), 1 deletion(-)
 create mode 100644 tools/container/runtime/runtime.go

diff --git a/tools/container/nvidia-toolkit/run.go b/tools/container/nvidia-toolkit/run.go
index fbb5c400..d17421b3 100644
--- a/tools/container/nvidia-toolkit/run.go
+++ b/tools/container/nvidia-toolkit/run.go
@@ -13,6 +13,7 @@ import (
 	cli "github.com/urfave/cli/v2"
 	unix "golang.org/x/sys/unix"
 
+	"github.com/NVIDIA/nvidia-container-toolkit/tools/container/runtime"
 	"github.com/NVIDIA/nvidia-container-toolkit/tools/container/toolkit"
 )
 
@@ -40,6 +41,7 @@ type options struct {
 	pidFile     string
 
 	toolkitOptions toolkit.Options
+	runtimeOptions runtime.Options
 }
 
 func (o options) toolkitRoot() string {
@@ -115,6 +117,7 @@ func main() {
 	}
 
 	c.Flags = append(c.Flags, toolkit.Flags(&options.toolkitOptions)...)
+	c.Flags = append(c.Flags, runtime.Flags(&options.runtimeOptions)...)
 
 	// Run the CLI
 	log.Infof("Starting %v", c.Name)
@@ -133,7 +136,9 @@ func validateFlags(_ *cli.Context, o *options) error {
 	if err := toolkit.ValidateOptions(&o.toolkitOptions, o.toolkitRoot()); err != nil {
 		return err
 	}
-
+	if err := runtime.ValidateOptions(&o.runtimeOptions, o.toolkitRoot()); err != nil {
+		return err
+	}
 	return nil
 }
 
diff --git a/tools/container/runtime/runtime.go b/tools/container/runtime/runtime.go
new file mode 100644
index 00000000..9f35e885
--- /dev/null
+++ b/tools/container/runtime/runtime.go
@@ -0,0 +1,108 @@
+/**
+# Copyright 2024 NVIDIA CORPORATION
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+**/
+
+package runtime
+
+import (
+	"fmt"
+
+	"github.com/urfave/cli/v2"
+
+	"github.com/NVIDIA/nvidia-container-toolkit/tools/container"
+)
+
+const (
+	defaultSetAsDefault = true
+	// defaultRuntimeName specifies the NVIDIA runtime to be use as the default runtime if setting the default runtime is enabled
+	defaultRuntimeName   = "nvidia"
+	defaultHostRootMount = "/host"
+
+	runtimeSpecificDefault = "RUNTIME_SPECIFIC_DEFAULT"
+)
+
+type Options struct {
+	container.Options
+}
+
+func Flags(opts *Options) []cli.Flag {
+	flags := []cli.Flag{
+		&cli.StringFlag{
+			Name:        "config",
+			Usage:       "Path to the runtime config file",
+			Value:       runtimeSpecificDefault,
+			Destination: &opts.Config,
+			EnvVars:     []string{"RUNTIME_CONFIG"},
+		},
+		&cli.StringFlag{
+			Name:        "socket",
+			Usage:       "Path to the runtime socket file",
+			Value:       runtimeSpecificDefault,
+			Destination: &opts.Socket,
+			EnvVars:     []string{"RUNTIME_SOCKET"},
+		},
+		&cli.StringFlag{
+			Name:        "restart-mode",
+			Usage:       "Specify how the runtime should be restarted; If 'none' is selected it will not be restarted [signal | systemd | none ]",
+			Value:       runtimeSpecificDefault,
+			Destination: &opts.RestartMode,
+			EnvVars:     []string{"RUNTIME_RESTART_MODE"},
+		},
+		&cli.StringFlag{
+			Name:        "host-root",
+			Usage:       "Specify the path to the host root to be used when restarting the runtime using systemd",
+			Value:       defaultHostRootMount,
+			Destination: &opts.HostRootMount,
+			EnvVars:     []string{"HOST_ROOT_MOUNT"},
+		},
+		&cli.StringFlag{
+			Name:        "runtime-name",
+			Aliases:     []string{"nvidia-runtime-name", "runtime-class"},
+			Usage:       "Specify the name of the `nvidia` runtime. If set-as-default is selected, the runtime is used as the default runtime.",
+			Value:       defaultRuntimeName,
+			Destination: &opts.RuntimeName,
+			EnvVars:     []string{"NVIDIA_RUNTIME_NAME"},
+		},
+		&cli.BoolFlag{
+			Name:        "set-as-default",
+			Usage:       "Set the `nvidia` runtime as the default runtime.",
+			Value:       defaultSetAsDefault,
+			Destination: &opts.SetAsDefault,
+			EnvVars:     []string{"NVIDIA_RUNTIME_SET_AS_DEFAULT"},
+			Hidden:      true,
+		},
+	}
+
+	return flags
+}
+
+// ValidateOptions checks whether the specified options are valid
+func ValidateOptions(opts *Options, toolkitRoot string) error {
+	// We set this option here to ensure that it is avalable in future calls.
+	opts.RuntimeDir = toolkitRoot
+	return nil
+}
+
+func Setup(c *cli.Context, opts *Options, runtime string) error {
+	switch runtime {
+	}
+	return fmt.Errorf("undefined runtime %v", runtime)
+}
+
+func Cleanup(c *cli.Context, opts *Options, runtime string) error {
+	switch runtime {
+	}
+	return fmt.Errorf("undefined runtime %v", runtime)
+}