|
| 1 | +/* |
| 2 | + * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. |
| 3 | + * |
| 4 | + * Licensed under the Apache License, Version 2.0 (the "License"); |
| 5 | + * you may not use this file except in compliance with the License. |
| 6 | + * You may obtain a copy of the License at |
| 7 | + * |
| 8 | + * http://www.apache.org/licenses/LICENSE-2.0 |
| 9 | + * |
| 10 | + * Unless required by applicable law or agreed to in writing, software |
| 11 | + * distributed under the License is distributed on an "AS IS" BASIS, |
| 12 | + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 13 | + * See the License for the specific language governing permissions and |
| 14 | + * limitations under the License. |
| 15 | + */ |
| 16 | + |
| 17 | +package e2e |
| 18 | + |
| 19 | +import ( |
| 20 | + "context" |
| 21 | + "fmt" |
| 22 | + "strings" |
| 23 | + |
| 24 | + . "github.com/onsi/ginkgo/v2" |
| 25 | + . "github.com/onsi/gomega" |
| 26 | +) |
| 27 | + |
| 28 | +const ( |
| 29 | + // libnvidiaContainerCliTestTemplate is an on-the-fly script that prepares a |
| 30 | + // minimal rootfs with GPU support and finally prints the list of visible GPUs |
| 31 | + // using `nvidia-smi -L`. The exit code as well as the standard output of this |
| 32 | + // script are therefore a good indicator of whether the NVIDIA Container |
| 33 | + // Toolkit is functioning correctly inside the container. |
| 34 | + libnvidiaContainerCliTestTemplate = `#!/usr/bin/env bash |
| 35 | +set -euo pipefail |
| 36 | +
|
| 37 | +apt-get update -y && apt-get install -y curl gnupg2 |
| 38 | +
|
| 39 | +WORKDIR="$(mktemp -d)" |
| 40 | +ROOTFS="${WORKDIR}/rootfs" |
| 41 | +mkdir -p "${ROOTFS}" |
| 42 | +
|
| 43 | +export WORKDIR ROOTFS # make them visible in the child shell |
| 44 | +
|
| 45 | +unshare --mount --pid --fork --propagation private -- bash -eux <<'IN_NS' |
| 46 | + : "${ROOTFS:?}" "${WORKDIR:?}" # abort if either is empty |
| 47 | +
|
| 48 | + # 1 Populate minimal Ubuntu base |
| 49 | + curl -L http://cdimage.ubuntu.com/ubuntu-base/releases/22.04/release/ubuntu-base-22.04-base-amd64.tar.gz \ |
| 50 | + | tar -C "$ROOTFS" -xz |
| 51 | +
|
| 52 | + # 2 Add non-root user |
| 53 | + useradd -R "$ROOTFS" -U -u 1000 -s /bin/bash nvidia |
| 54 | +
|
| 55 | + # 3 Bind-mount new root and unshare mounts |
| 56 | + mount --bind "$ROOTFS" "$ROOTFS" |
| 57 | + mount --make-private "$ROOTFS" |
| 58 | + cd "$ROOTFS" |
| 59 | +
|
| 60 | + # 4 Minimal virtual filesystems |
| 61 | + mount -t proc proc proc |
| 62 | + mount -t sysfs sys sys |
| 63 | + mount -t tmpfs tmp tmp |
| 64 | + mount -t tmpfs run run |
| 65 | +
|
| 66 | + # 5 GPU setup |
| 67 | + nvidia-container-cli --load-kmods --debug=container-cli.log \ |
| 68 | + configure --ldconfig=@/sbin/ldconfig.real \ |
| 69 | + --no-cgroups --utility --device=0 "$(pwd)" |
| 70 | +
|
| 71 | + # 6 Switch root |
| 72 | + mkdir -p mnt |
| 73 | + pivot_root . mnt |
| 74 | + umount -l /mnt |
| 75 | +
|
| 76 | + exec nvidia-smi -L |
| 77 | +IN_NS |
| 78 | +` |
| 79 | + |
| 80 | + dockerRunCmdTemplate = `docker run --name %s -d --privileged --runtime=nvidia \ |
| 81 | + -e NVIDIA_VISIBLE_DEVICES=all \ |
| 82 | + -e NVIDIA_DRIVER_CAPABILITIES=all \ |
| 83 | + -v %s:/libnvidia-container-cli.sh \ |
| 84 | + --entrypoint /libnvidia-container-cli.sh %s` |
| 85 | +) |
| 86 | + |
| 87 | +var _ = Describe("nvidia-container-cli", Ordered, ContinueOnFailure, func() { |
| 88 | + var ( |
| 89 | + runner Runner |
| 90 | + testScript = "/tmp/libnvidia-container-cli.sh" |
| 91 | + dockerImage = "ghcr.io/nvidia/container-toolkit:5e8c1411-ubuntu20.04" |
| 92 | + containerName = "nvidia-cli-e2e" |
| 93 | + dockerRunCmd string |
| 94 | + ) |
| 95 | + |
| 96 | + BeforeAll(func(ctx context.Context) { |
| 97 | + runner = NewRunner( |
| 98 | + WithHost(sshHost), |
| 99 | + WithPort(sshPort), |
| 100 | + WithSshKey(sshKey), |
| 101 | + WithSshUser(sshUser), |
| 102 | + ) |
| 103 | + |
| 104 | + if installCTK { |
| 105 | + installer, err := NewToolkitInstaller( |
| 106 | + WithRunner(runner), |
| 107 | + WithImage(imageName+":"+imageTag), |
| 108 | + WithTemplate(dockerInstallTemplate), |
| 109 | + ) |
| 110 | + Expect(err).ToNot(HaveOccurred()) |
| 111 | + |
| 112 | + err = installer.Install() |
| 113 | + Expect(err).ToNot(HaveOccurred()) |
| 114 | + } |
| 115 | + }) |
| 116 | + |
| 117 | + AfterAll(func(ctx context.Context) { |
| 118 | + // Cleanup: remove the container and the temporary script on the host. |
| 119 | + runner.Run(fmt.Sprintf("docker rm -f %s", containerName)) |
| 120 | + runner.Run(fmt.Sprintf("rm -f %s", testScript)) |
| 121 | + }) |
| 122 | + |
| 123 | + It("should report the same GPUs inside the container as on the host", func(ctx context.Context) { |
| 124 | + // Write the script to the remote host and make it executable. |
| 125 | + createScriptCmd := fmt.Sprintf( |
| 126 | + "cat > %s <<'EOF'\n%s\nEOF\nchmod +x %s", |
| 127 | + testScript, libnvidiaContainerCliTestTemplate, testScript, |
| 128 | + ) |
| 129 | + |
| 130 | + _, _, err := runner.Run(createScriptCmd) |
| 131 | + Expect(err).ToNot(HaveOccurred()) |
| 132 | + |
| 133 | + // If a container with the same name exists from a previous test run, remove it first. |
| 134 | + runner.Run(fmt.Sprintf("docker rm -f %s", containerName)) |
| 135 | + |
| 136 | + // Build the docker run command (detached mode) from the template so it |
| 137 | + // stays readable while still resulting in a single-line invocation. |
| 138 | + dockerRunCmd = fmt.Sprintf(dockerRunCmdTemplate, containerName, testScript, dockerImage) |
| 139 | + |
| 140 | + // Launch the container in detached mode. |
| 141 | + _, _, err = runner.Run(dockerRunCmd) |
| 142 | + Expect(err).ToNot(HaveOccurred()) |
| 143 | + |
| 144 | + // Capture the host GPU list. |
| 145 | + hostOutput, _, err := runner.Run("nvidia-smi -L") |
| 146 | + Expect(err).ToNot(HaveOccurred()) |
| 147 | + |
| 148 | + hostOutput = strings.TrimSpace(strings.ReplaceAll(hostOutput, "\r", "")) |
| 149 | + |
| 150 | + // Poll the logs of the already running container until we observe |
| 151 | + // the GPU list matching the host or until a 5-minute timeout elapses. |
| 152 | + Eventually(func() string { |
| 153 | + logs, _, err := runner.Run(fmt.Sprintf("docker logs %s", containerName)) |
| 154 | + if err != nil { |
| 155 | + return "" |
| 156 | + } |
| 157 | + |
| 158 | + lines := strings.Split(strings.TrimSpace(logs), "\n") |
| 159 | + if len(lines) == 0 { |
| 160 | + return "" |
| 161 | + } |
| 162 | + |
| 163 | + last := strings.TrimSpace(strings.ReplaceAll(lines[len(lines)-1], "\r", "")) |
| 164 | + return last |
| 165 | + }, "5m", "5s").Should(Equal(hostOutput)) |
| 166 | + }) |
| 167 | +}) |
0 commit comments