From 90518e0ce5ee7681c7f677f8345f4f86f2269f10 Mon Sep 17 00:00:00 2001 From: Evan Lezar Date: Mon, 25 Jul 2022 10:31:31 +0200 Subject: [PATCH] Allow accept-visible-devices config options to be set This change allows the * accept-nvidia-visible-devices-envvar-when-unprivileged * accept-nvidia-visible-devices-as-volume-mounts options to be set in the toolkit-container. These are controlled by command line flags or the following environment variables: * ACCEPT_NVIDIA_VISIBLE_DEVICES_ENVVAR_WHEN_UNPRIVILEGED * ACCEPT_NVIDIA_VISIBLE_DEVICES_AS_VOLUME_MOUNTS Signed-off-by: Evan Lezar --- tools/container/toolkit/toolkit.go | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tools/container/toolkit/toolkit.go b/tools/container/toolkit/toolkit.go index e8f7fde8..b119eb42 100644 --- a/tools/container/toolkit/toolkit.go +++ b/tools/container/toolkit/toolkit.go @@ -45,6 +45,9 @@ type options struct { ContainerRuntimeLogLevel string ContainerCLIDebug string toolkitRoot string + + acceptNVIDIAVisibleDevicesWhenUnprivileged bool + acceptNVIDIAVisibleDevicesAsVolumeMounts bool } func main() { @@ -111,6 +114,18 @@ func main() { Destination: &opts.ContainerCLIDebug, EnvVars: []string{"NVIDIA_CONTAINER_CLI_DEBUG"}, }, + &cli.BoolFlag{ + Name: "accept-nvidia-visible-devices-envvar-when-unprivileged", + Usage: "Set the accept-nvidia-visible-devices-envvar-when-unprivileged config option", + Destination: &opts.acceptNVIDIAVisibleDevicesWhenUnprivileged, + EnvVars: []string{"ACCEPT_NVIDIA_VISIBLE_DEVICES_ENVVAR_WHEN_UNPRIVILEGED"}, + }, + &cli.BoolFlag{ + Name: "accept-nvidia-visible-devices-as-volume-mounts", + Usage: "Set the accept-nvidia-visible-devices-as-volume-mounts config option", + Destination: &opts.acceptNVIDIAVisibleDevicesWhenUnprivileged, + EnvVars: []string{"ACCEPT_NVIDIA_VISIBLE_DEVICES_AS_VOLUME_MOUNTS"}, + }, &cli.StringFlag{ Name: "toolkit-root", Usage: "The directory where the NVIDIA Container toolkit is to be installed", @@ -260,6 +275,10 @@ func installToolkitConfig(toolkitConfigPath string, nvidiaContainerCliExecutable } defer targetConfig.Close() + // Set the options in the root toml table + config.Set("accept-nvidia-visible-devices-envvar-when-unprivileged", opts.acceptNVIDIAVisibleDevicesWhenUnprivileged) + config.Set("accept-nvidia-visible-devices-as-volume-mounts", opts.acceptNVIDIAVisibleDevicesAsVolumeMounts) + nvidiaContainerCliKey := func(p string) []string { return []string{"nvidia-container-cli", p} }