/* * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. * * 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 e2e import ( "context" "fmt" "strings" "text/template" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" ) const ( libnvidiaContainerCliDockerRunTemplate = ` docker run -d --name test-nvidia-container-cli \ --privileged \ --runtime=nvidia \ -e NVIDIA_VISIBLE_DEVICES=all \ -e NVIDIA_DRIVER_CAPABILITIES=all \ -v $HOME/libnvidia-container-cli.sh:/usr/local/bin/libnvidia-container-cli.sh \ -v /usr/bin/nvidia-container-cli:/usr/bin/nvidia-container-cli \ -v /usr/bin/nvidia-ctk:/usr/bin/nvidia-ctk \ -v /usr/bin/nvidia-container-runtime:/usr/bin/nvidia-container-runtime \ -v /usr/bin/nvidia-container-runtime-hook:/usr/bin/nvidia-container-runtime-hook \ -v /usr/bin/nvidia-container-toolkit:/usr/bin/nvidia-container-toolkit \ -v /usr/bin/nvidia-cdi-hook:/usr/bin/nvidia-cdi-hook \ -v /usr/bin/nvidia-container-runtime.cdi:/usr/bin/nvidia-container-runtime.cdi \ -v /usr/bin/nvidia-container-runtime.legacy:/usr/bin/nvidia-container-runtime.legacy \ -v /usr/local/nvidia/toolkit:/usr/local/nvidia/toolkit \ -v /etc/nvidia-container-runtime:/etc/nvidia-container-runtime \ -v /usr/lib/x86_64-linux-gnu/libnvidia-container.so.1:/usr/lib/x86_64-linux-gnu/libnvidia-container.so.1 \ -v /usr/lib/x86_64-linux-gnu/{{.LibNvidiaContainerTarget}}:/usr/lib/x86_64-linux-gnu/{{.LibNvidiaContainerTarget}} \ -v /usr/lib/x86_64-linux-gnu/libnvidia-container-go.so.1:/usr/lib/x86_64-linux-gnu/libnvidia-container-go.so.1 \ -v /usr/lib/x86_64-linux-gnu/{{.LibNvidiaContainerGoTarget}}:/usr/lib/x86_64-linux-gnu/{{.LibNvidiaContainerGoTarget}} \ -e LD_LIBRARY_PATH=/usr/lib64:/usr/lib/x86_64-linux-gnu:/usr/lib/aarch64-linux-gnu:/lib64:/lib/x86_64-linux-gnu:/lib/aarch64-linux-gnu \ --entrypoint /usr/local/bin/libnvidia-container-cli.sh \ ubuntu ` libnvidiaContainerCliTestTemplate = `#!/usr/bin/env bash set -euo pipefail apt-get update -y && apt-get install -y curl gnupg2 WORKDIR="$(mktemp -d)" ROOTFS="${WORKDIR}/rootfs" mkdir -p "${ROOTFS}" export WORKDIR ROOTFS # make them visible in the child shell unshare --mount --pid --fork --propagation private -- bash -eux <<'IN_NS' : "${ROOTFS:?}" "${WORKDIR:?}" # abort if either is empty # 1 Populate minimal Ubuntu base curl -L http://cdimage.ubuntu.com/ubuntu-base/releases/22.04/release/ubuntu-base-22.04-base-amd64.tar.gz \ | tar -C "$ROOTFS" -xz # 2 Add non-root user useradd -R "$ROOTFS" -U -u 1000 -s /bin/bash nvidia # 3 Bind-mount new root and unshare mounts mount --bind "$ROOTFS" "$ROOTFS" mount --make-private "$ROOTFS" cd "$ROOTFS" # 4 Minimal virtual filesystems mount -t proc proc proc mount -t sysfs sys sys mount -t tmpfs tmp tmp mount -t tmpfs run run # 5 GPU setup nvidia-container-cli --load-kmods --debug=container-cli.log \ configure --ldconfig=@/sbin/ldconfig.real \ --no-cgroups --utility --device=0 "$(pwd)" # 6 Switch root mkdir -p mnt pivot_root . mnt umount -l /mnt exec nvidia-smi -L IN_NS ` ) // Integration tests for Docker runtime var _ = Describe("nvidia-container-cli", Ordered, ContinueOnFailure, func() { var runner Runner // Install the NVIDIA Container Toolkit BeforeAll(func(ctx context.Context) { runner = NewRunner( WithHost(sshHost), WithPort(sshPort), WithSshKey(sshKey), WithSshUser(sshUser), ) if installCTK { installer, err := NewToolkitInstaller( WithRunner(runner), WithImage(imageName+":"+imageTag), WithTemplate(dockerInstallTemplate), ) Expect(err).ToNot(HaveOccurred()) err = installer.Install() Expect(err).ToNot(HaveOccurred()) } }) When("running nvidia-smi -L", Ordered, func() { var hostOutput string var err error BeforeAll(func(ctx context.Context) { hostOutput, _, err = runner.Run("nvidia-smi -L") Expect(err).ToNot(HaveOccurred()) _, _, err := runner.Run("docker pull ubuntu") Expect(err).ToNot(HaveOccurred()) }) AfterAll(func(ctx context.Context) { _, _, err := runner.Run("docker rm -f test-nvidia-container-cli") Expect(err).ToNot(HaveOccurred()) }) It("should support NVIDIA_VISIBLE_DEVICES and NVIDIA_DRIVER_CAPABILITIES", func(ctx context.Context) { // 1. Create the test script on the remote host at $HOME/test.sh using a here-document testScriptPath := "$HOME/libnvidia-container-cli.sh" testScript := libnvidiaContainerCliTestTemplate createScriptCmd := fmt.Sprintf("cat > %s <<'EOF'\n%s\nEOF\nchmod +x %s", testScriptPath, testScript, testScriptPath) _, _, err := runner.Run(createScriptCmd) Expect(err).ToNot(HaveOccurred()) // 2. Discover the symlink targets for the libraries on the remote host getTargetCmd := func(lib string) string { return fmt.Sprintf("readlink -f /usr/lib/x86_64-linux-gnu/%s.1", lib) } libNvidiaContainerTarget, _, err := runner.Run(getTargetCmd("libnvidia-container.so")) Expect(err).ToNot(HaveOccurred()) libNvidiaContainerTarget = strings.TrimSpace(libNvidiaContainerTarget) libNvidiaContainerTarget = strings.TrimPrefix(libNvidiaContainerTarget, "/usr/lib/x86_64-linux-gnu/") libNvidiaContainerGoTarget, _, err := runner.Run(getTargetCmd("libnvidia-container-go.so")) Expect(err).ToNot(HaveOccurred()) libNvidiaContainerGoTarget = strings.TrimSpace(libNvidiaContainerGoTarget) libNvidiaContainerGoTarget = strings.TrimPrefix(libNvidiaContainerGoTarget, "/usr/lib/x86_64-linux-gnu/") // 3. Render the docker run template with the discovered targets tmpl, err := template.New("dockerRun").Parse(libnvidiaContainerCliDockerRunTemplate) Expect(err).ToNot(HaveOccurred()) var dockerRunCmdBuilder strings.Builder err = tmpl.Execute(&dockerRunCmdBuilder, map[string]string{ "LibNvidiaContainerTarget": libNvidiaContainerTarget, "LibNvidiaContainerGoTarget": libNvidiaContainerGoTarget, }) Expect(err).ToNot(HaveOccurred()) dockerRunCmd := dockerRunCmdBuilder.String() // 4. Start the container using the rendered docker run command _, _, err = runner.Run(dockerRunCmd) Expect(err).ToNot(HaveOccurred()) // 5. Use Eventually to check the container logs contain hostOutput Eventually(func() string { logs, _, err := runner.Run("docker logs test-nvidia-container-cli") if err != nil { return "" } return logs }, "5m", "5s").Should(ContainSubstring(hostOutput)) }) }) })