feature: add nvidia-container 2.0

Signed-off-by: codejuan <[email protected]>

Author: CodeJuan

apis/swagger.yml

@@ -2471,6 +2471,29 @@ definitions:
         x-nullable: false
         minimum: 0
         maximum: 1
+      NvidiaConfig:
+        $ref: "#/definitions/NvidiaConfig"
+
+  NvidiaConfig:
+    type: "object"
+    properties:
+      NvidiaVisibleDevices:
+        description: "NvidiaVisibleDevices controls which GPUs will be made accessible inside the container"
+        type: "string"
+        example: |
+                  Possible values.
+                  0,1,2, GPU-fef8089b …: a comma-separated list of GPU UUID(s) or index(es).
+                  all: all GPUs will be accessible, this is the default value in our container images.
+                  none: no GPU will be accessible, but driver capabilities will be enabled.
+        x-nullable: false
+      NvidiaDriverCapabilities:
+        description: "NvidiaDriverCapabilities controls which driver libraries/binaries will be mounted inside the container"
+        type: "string"
+        example: |
+                  Possible values
+                  compute,video, graphics,utility …: a comma-separated list of driver features the container needs.
+                  all: enable all available driver capabilities.
+        x-nullable: false
 
   ThrottleDevice:
     type: "object"

apis/types/nvidia_config.go

@@ -2,6 +2,7 @@ package mgr
 
 import (
 	"fmt"
+	"os"
 	"strconv"
 	"strings"
 
@@ -11,6 +12,7 @@ import (
 	"github.com/alibaba/pouch/pkg/meta"
 	"github.com/alibaba/pouch/pkg/randomid"
 	"github.com/alibaba/pouch/pkg/system"
+	"github.com/alibaba/pouch/pkg/utils"
 
 	"github.com/opencontainers/selinux/go-selinux/label"
 	"github.com/pkg/errors"
@@ -234,6 +236,10 @@ func validateConfig(config *types.ContainerConfig, hostConfig *types.HostConfig,
 	if err != nil {
 		return nil, err
 	}
+	// validates nvidia config
+	if err := validateNvidiaConfig(hostConfig); err != nil {
+		return warnings, err
+	}
 	warnings = append(warnings, warns...)
 
 	if hostConfig.OomScoreAdj < -1000 || hostConfig.OomScoreAdj > 1000 {
@@ -379,3 +385,76 @@ func amendContainerSettings(config *types.ContainerConfig, hostConfig *types.Hos
 		r.MemorySwap = 2 * r.Memory
 	}
 }
+
+// validateNvidiaConfig
+func validateNvidiaConfig(hostConfig *types.HostConfig) error {
+	r := &hostConfig.Resources
+	if r.NvidiaConfig == nil {
+		return nil
+	}
+
+	if err := vlidateNvidiaDriver(r); err != nil {
+		return err
+	}
+
+	if err := validateNvidiaDevice(hostConfig); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func vlidateNvidiaDriver(r *types.Resources) error {
+	n := r.NvidiaConfig
+	n.NvidiaDriverCapabilities = strings.TrimSpace(n.NvidiaDriverCapabilities)
+	if n.NvidiaDriverCapabilities == "" {
+		// use default driver capability: utility
+		return nil
+	}
+
+	if n.NvidiaDriverCapabilities == "all" {
+		// enable all capabilities
+		return nil
+	}
+
+	supportedDrivers := []string{"compute", "compat32", "graphics", "utility", "video", "display"}
+	drivers := strings.Split(n.NvidiaDriverCapabilities, ",")
+	for _, d := range drivers {
+		d = strings.TrimSpace(d)
+		found := utils.StringInSlice(supportedDrivers, d)
+		if !found {
+			return fmt.Errorf("invalid nvidia driver capability (%s)", d)
+		}
+	}
+	return nil
+}
+
+func validateNvidiaDevice(hostConfig *types.HostConfig) error {
+	n := hostConfig.Resources.NvidiaConfig
+	n.NvidiaVisibleDevices = strings.TrimSpace(n.NvidiaVisibleDevices)
+
+	// none: no GPU will be accessible, but driver capabilities will be enabled.
+	// void or empty: no GPU will be accessible, and driver capabilities will be disabled.
+	// all: all GPUs will be accessible
+	if n.NvidiaDriverCapabilities == "" {
+		return nil
+	}
+	supportedDevices := []string{"all", "none", "void"}
+	found := utils.StringInSlice(supportedDevices, n.NvidiaVisibleDevices)
+	if found {
+		return nil
+	}
+	devs := strings.Split(n.NvidiaVisibleDevices, ",")
+	for _, dev := range devs {
+		dev = strings.TrimSpace(dev)
+		if _, err := strconv.Atoi(dev); err == nil {
+			//dev is numeric, the realDev should be /dev/nvidiaN
+			realDev := fmt.Sprintf("/dev/nvidia%s", dev)
+			if _, err := os.Stat(realDev); err != nil {
+				return fmt.Errorf("invalid nvidia device %s", realDev)
+			}
+		}
+		// TODO: how to validate GPU UUID
+	}
+	return nil
+}

apis/types/resources.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"io/ioutil"
 	"os"
+	"os/exec"
 	"path/filepath"
 	"sort"
 	"strconv"
@@ -64,6 +65,11 @@ func setupHook(ctx context.Context, c *Container, specWrapper *SpecWrapper) erro
 		return errors.Wrap(err, "failed to set volume mount tab prestart hook")
 	}
 
+	// set nvidia config
+	if err := setNvidiaHook(ctx, c, specWrapper); err != nil {
+		return errors.Wrap(err, "failed to set nvidia prestart hook")
+	}
+
 	return nil
 }
 
@@ -152,6 +158,24 @@ func setMountTab(ctx context.Context, c *Container, spec *SpecWrapper) error {
 	return nil
 }
 
+func setNvidiaHook(ctx context.Context, c *Container, spec *SpecWrapper) error {
+	n := c.HostConfig.NvidiaConfig
+	if n == nil {
+		return nil
+	}
+	path, err := exec.LookPath("nvidia-container-runtime-hook")
+	if err != nil {
+		return err
+	}
+	args := []string{path}
+	nvidiaPrestart := specs.Hook{
+		Path: path,
+		Args: append(args, "prestart"),
+	}
+	spec.s.Hooks.Prestart = append(spec.s.Hooks.Prestart, nvidiaPrestart)
+	return nil
+}
+
 type hookArray []*wrapperEmbedPrestart
 
 // Len is defined in order to support sort

daemon/mgr/container_utils.go

@@ -2,6 +2,7 @@ package mgr
 
 import (
 	"context"
+	"fmt"
 	"io/ioutil"
 	"os"
 	"strings"
@@ -61,6 +62,9 @@ func setupProcess(ctx context.Context, c *Container, s *specs.Spec) error {
 		return err
 	}
 
+	if err := setupNvidiaEnv(ctx, c, s); err != nil {
+		return err
+	}
 	return nil
 }
 
@@ -163,3 +167,13 @@ func setupRlimits(ctx context.Context, hostConfig *types.HostConfig, s *specs.Sp
 	s.Process.Rlimits = rlimits
 	return nil
 }
+
+func setupNvidiaEnv(ctx context.Context, c *Container, s *specs.Spec) error {
+	n := c.HostConfig.NvidiaConfig
+	if n == nil {
+		return nil
+	}
+	s.Process.Env = append(s.Process.Env, fmt.Sprintf("NVIDIA_DRIVER_CAPABILITIES=%s", n.NvidiaDriverCapabilities))
+	s.Process.Env = append(s.Process.Env, fmt.Sprintf("NVIDIA_VISIBLE_DEVICES=%s", n.NvidiaVisibleDevices))
+	return nil
+}

daemon/mgr/spec_hook.go

@@ -16,6 +16,8 @@ POUCHDIR=$TMP/source
 [ -d "$POUCHDIR" ] || mkdir -p "$POUCHDIR"
 BINDIR=$POUCHDIR/bin
 [ -d "$BINDIR" ] || mkdir -p "$BINDIR"
+LIBDIR=$POUCHDIR/lib
+[ -d "$LIBDIR" ] || mkdir -p "$LIBDIR"
 LXC_DIR=$TMP/lxc
 [ -d "$LXC_DIR" ] || mkdir -p "$LXC_DIR"
 
@@ -34,6 +36,9 @@ CATEGORY='Tools/Pouch'
 MAINTAINER='Pouch [email protected]'
 VENDOR='Pouch'
 
+LIB_NVIDIA_VERSION="1.0.0-rc.2"
+NVIDIA_RUNTIME_VERSION="1.4.0-1"
+
 # build lxcfs
 function build_lxcfs ()
 {
@@ -69,6 +74,24 @@ function build_pouch()
     popd
 }
 
+# install nvidia-container-runtime
+function build_nvidia_runtime(){
+    echo "Downloading libnvidia-container."
+    wget --quiet "https://github.com/NVIDIA/libnvidia-container/releases/download/v${LIB_NVIDIA_VERSION}/libnvidia-container_${LIB_NVIDIA_VERSION}_x86_64.tar.xz" -P "${TMP}"
+    tar -xf "${TMP}/libnvidia-container_${LIB_NVIDIA_VERSION}_x86_64.tar.xz" -C "${TMP}"
+    cp "${TMP}/libnvidia-container_${LIB_NVIDIA_VERSION}/usr/local/bin/nvidia-container-cli" "${BINDIR}/"
+    cp "${TMP}/libnvidia-container_${LIB_NVIDIA_VERSION}/usr/local/lib/libnvidia-container.so" "${LIBDIR}/"
+    cp "${TMP}/libnvidia-container_${LIB_NVIDIA_VERSION}/usr/local/lib/libnvidia-container.so.1" "${LIBDIR}/"
+    cp "${TMP}/libnvidia-container_${LIB_NVIDIA_VERSION}/usr/local/lib/libnvidia-container.so.1.0.0" "${LIBDIR}/"
+
+    echo "Downloading nvidia-container-runtime."
+    wget --quiet "https://github.com/NVIDIA/nvidia-container-runtime/archive/v${NVIDIA_RUNTIME_VERSION}.tar.gz" -P "${TMP}"
+    mkdir -p "${GOPATH}/src/github.com/NVIDIA"
+    tar -xzf "${TMP}/v${NVIDIA_RUNTIME_VERSION}.tar.gz" -C "${GOPATH}/src/github.com/NVIDIA"
+    mv "${GOPATH}/src/github.com/NVIDIA/nvidia-container-runtime-${NVIDIA_RUNTIME_VERSION}" "${GOPATH}/src/github.com/NVIDIA/nvidia-container-runtime"
+    go build -o "${BINDIR}/nvidia-container-runtime-hook" "github.com/NVIDIA/nvidia-container-runtime/hook/nvidia-container-runtime-hook"
+}
+
 function build_rpm ()
 {
     pushd $MOUNTDIR
@@ -110,6 +133,7 @@ function build_rpm ()
           -d fuse-libs \
           -d fuse \
           "$BINDIR/"=/usr/local/bin/ \
+          "$LIBDIR/"=/usr/lib64/ \
           "$SERVICEDIR/"=/usr/lib/systemd/system/ \
           "$LXC_DIR/usr/local/bin/lxcfs"=/usr/bin/pouch-lxcfs \
           "$LXC_DIR/usr/local/lib/lxcfs/libpouchlxcfs.so"=/usr/lib64/libpouchlxcfs.so \
@@ -121,6 +145,7 @@ function main()
 	echo "Building rpm package."
 	build_pouch
 	build_lxcfs
+	build_nvidia_runtime
 	build_rpm
 }