|
| 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-cli |
| 33 | + // is functioning correctly inside the container. |
| 34 | + libnvidiaContainerCliTestTemplate = `#!/busybox/sh |
| 35 | +set -euo pipefail |
| 36 | +
|
| 37 | +# Create a temporary directory and rootfs path |
| 38 | +TMPDIR="$(mktemp -d)" |
| 39 | +ROOTFS="${TMPDIR}/rootfs" |
| 40 | +mkdir -p "${ROOTFS}" |
| 41 | +
|
| 42 | +# Expose ROOTFS for the child namespace |
| 43 | +export ROOTFS |
| 44 | +
|
| 45 | +busybox wget -qO- http://cdimage.ubuntu.com/ubuntu-base/releases/22.04/release/ubuntu-base-22.04-base-amd64.tar.gz \ |
| 46 | + | tar -C "$ROOTFS" -xz |
| 47 | +
|
| 48 | +# Enter a new mount + PID namespace so we can pivot_root without touching the |
| 49 | +# container's original filesystem. |
| 50 | +busybox unshare --mount --pid --fork --propagation private -- busybox sh -eux <<'IN_NS' |
| 51 | + : "${ROOTFS:?}" |
| 52 | +
|
| 53 | + # 1 Bind-mount the new root and make the mount private |
| 54 | + busybox mount --bind "$ROOTFS" "$ROOTFS" |
| 55 | + busybox mount --make-private "$ROOTFS" |
| 56 | + cd "$ROOTFS" |
| 57 | +
|
| 58 | + # 2 Minimal virtual filesystems |
| 59 | + busybox mount -t proc proc proc |
| 60 | + busybox mount -t sysfs sys sys |
| 61 | + busybox mount -t tmpfs tmp tmp |
| 62 | + busybox mount -t tmpfs run run |
| 63 | +
|
| 64 | + # 3 GPU setup via nvidia-container-cli |
| 65 | + nvidia-container-cli --load-kmods \ |
| 66 | + configure \ |
| 67 | + --no-cgroups --utility --device=0 "$(pwd)" |
| 68 | +
|
| 69 | + # 4 Switch root into the prepared filesystem |
| 70 | + mkdir -p mnt |
| 71 | + busybox pivot_root . mnt |
| 72 | + busybox umount -l /mnt |
| 73 | +
|
| 74 | + exec nvidia-smi -L |
| 75 | +IN_NS |
| 76 | +` |
| 77 | + |
| 78 | + dockerRunCmdTemplate = `docker run --name %s -d --privileged --runtime=nvidia \ |
| 79 | + -e NVIDIA_VISIBLE_DEVICES=all \ |
| 80 | + -e NVIDIA_DRIVER_CAPABILITIES=all \ |
| 81 | + -v %s:/libnvidia-container-cli.sh \ |
| 82 | + --entrypoint /libnvidia-container-cli.sh %s` |
| 83 | +) |
| 84 | + |
| 85 | +var _ = Describe("nvidia-container-cli", Ordered, ContinueOnFailure, func() { |
| 86 | + var ( |
| 87 | + runner Runner |
| 88 | + dockerImage = fmt.Sprintf("%s:%s", imageName, imageTag) |
| 89 | + containerName = "nvidia-cli-e2e" |
| 90 | + hostOutput string |
| 91 | + ) |
| 92 | + |
| 93 | + BeforeAll(func(ctx context.Context) { |
| 94 | + runner = NewRunner( |
| 95 | + WithHost(sshHost), |
| 96 | + WithPort(sshPort), |
| 97 | + WithSshKey(sshKey), |
| 98 | + WithSshUser(sshUser), |
| 99 | + ) |
| 100 | + |
| 101 | + if installCTK { |
| 102 | + installer, err := NewToolkitInstaller( |
| 103 | + WithRunner(runner), |
| 104 | + WithImage(imageName+":"+imageTag), |
| 105 | + WithTemplate(dockerInstallTemplate), |
| 106 | + ) |
| 107 | + Expect(err).ToNot(HaveOccurred()) |
| 108 | + |
| 109 | + err = installer.Install() |
| 110 | + Expect(err).ToNot(HaveOccurred()) |
| 111 | + } |
| 112 | + |
| 113 | + // Capture the host GPU list. |
| 114 | + var err error |
| 115 | + hostOutput, _, err = runner.Run("nvidia-smi -L") |
| 116 | + Expect(err).ToNot(HaveOccurred()) |
| 117 | + }) |
| 118 | + |
| 119 | + AfterAll(func(ctx context.Context) { |
| 120 | + // Cleanup: remove the container and the temporary script on the host. |
| 121 | + runner.Run(fmt.Sprintf("docker rm -f %s", containerName)) |
| 122 | + }) |
| 123 | + |
| 124 | + It("should report the same GPUs inside the container as on the host", func(ctx context.Context) { |
| 125 | + testScriptPath := "/tmp/libnvidia-container-cli.sh" |
| 126 | + |
| 127 | + // Write the script to the remote host and make it executable. |
| 128 | + createScriptCmd := fmt.Sprintf( |
| 129 | + "cat > %s <<'EOF'\n%s\nEOF\nchmod +x %s", |
| 130 | + testScriptPath, libnvidiaContainerCliTestTemplate, testScriptPath, |
| 131 | + ) |
| 132 | + |
| 133 | + _, _, err := runner.Run(createScriptCmd) |
| 134 | + Expect(err).ToNot(HaveOccurred()) |
| 135 | + |
| 136 | + // If a container with the same name exists from a previous test run, remove it first. |
| 137 | + _, _, err = runner.Run(fmt.Sprintf("docker rm -f %s", containerName)) |
| 138 | + Expect(err).ToNot(HaveOccurred()) |
| 139 | + |
| 140 | + // Build the docker run command (detached mode) from the template so it |
| 141 | + // stays readable while still resulting in a single-line invocation. |
| 142 | + dockerRunCmd := fmt.Sprintf(dockerRunCmdTemplate, containerName, testScriptPath, dockerImage) |
| 143 | + |
| 144 | + // Launch the container in detached mode. |
| 145 | + _, _, err = runner.Run(dockerRunCmd) |
| 146 | + Expect(err).ToNot(HaveOccurred()) |
| 147 | + |
| 148 | + //check if the container is running |
| 149 | + containerOutput, _, err := runner.Run(fmt.Sprintf("docker ps -a | grep %s", containerName)) |
| 150 | + Expect(err).ToNot(HaveOccurred()) |
| 151 | + fmt.Println("containerOutput: ", containerOutput) |
| 152 | + |
| 153 | + hostOutputNormalized := strings.TrimSpace(strings.ReplaceAll(hostOutput, "\r", "")) |
| 154 | + |
| 155 | + // Poll the logs of the already running container until we observe |
| 156 | + // the GPU list matching the host or until a 5-minute timeout elapses. |
| 157 | + Eventually(func() string { |
| 158 | + logs, logsErr, err := runner.Run(fmt.Sprintf("docker logs --tail 10 %s", containerName)) |
| 159 | + if err != nil { |
| 160 | + return "" |
| 161 | + } |
| 162 | + |
| 163 | + fmt.Println("logs: ", logs) |
| 164 | + fmt.Println("logsErr: ", logsErr) |
| 165 | + |
| 166 | + lines := strings.Split(strings.TrimSpace(logs), "\n") |
| 167 | + if len(lines) == 0 { |
| 168 | + return "" |
| 169 | + } |
| 170 | + |
| 171 | + last := strings.TrimSpace(strings.ReplaceAll(lines[len(lines)-1], "\r", "")) |
| 172 | + return last |
| 173 | + }, "5m", "10s").Should(Equal(hostOutputNormalized)) |
| 174 | + |
| 175 | + // Remove the script from the remote host. |
| 176 | + runner.Run(fmt.Sprintf("rm -f %s", testScriptPath)) |
| 177 | + }) |
| 178 | +}) |
0 commit comments