Skip to content

Test AMI with updated nvidia driver #11528

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
mikkeloscar wants to merge 2 commits into from
Closed

Test AMI with updated nvidia driver #11528

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
64a7087
Test AMI with updated nvidia driver
mikkeloscar Jul 1, 2026
937d2bd
e2e: Check for GPU Driver/ML framework compatibilty
mikkeloscar Jul 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions cluster/config-defaults.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -907,8 +907,8 @@ tracing_coredns_local_zone_traces_endpoint: ""
# AMI id given the image name and the Image AWS account owner.
#
# [0]: https://github.com/zalando-incubator/cluster-lifecycle-manager/blob/8a9bd1cb2d094038a9e23e646421f8146b48886a/provisioner/template.go#L116
kuberuntu_image_v1_35_new_amd64: {{ amiID "zalando-ubuntu-jammy-22.04-kubernetes-production-v1.35.5-amd64-master-405" "861068367966" }}
kuberuntu_image_v1_35_new_arm64: {{ amiID "zalando-ubuntu-jammy-22.04-kubernetes-production-v1.35.5-arm64-master-405" "861068367966" }}
kuberuntu_image_v1_35_new_amd64: {{ amiID "zalando-ubuntu-jammy-22.04-kubernetes-test-v1.35.5-amd64-pr-444-7" "861068367966" }}
kuberuntu_image_v1_35_new_arm64: {{ amiID "zalando-ubuntu-jammy-22.04-kubernetes-test-v1.35.5-arm64-pr-444-7" "861068367966" }}

# This is used to determine which AMI to use for the cluster or individual node
# pools. Possible values are 'new' or 'old'
Expand Down
92 changes: 57 additions & 35 deletions test/e2e/gpu.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,41 +37,63 @@ var _ = describe("GPU job processing", func() {
cs = f.ClientSet
})

f.It("Should run a job on a gpu node [Zalando] [GPU]", f.WithSlow(), func(ctx context.Context) {
ns := f.Namespace.Name
nameprefix := "gpu-test-"
labels := map[string]string{
"application": "vector-add",
}
f.It("Should run a vector-add job on a gpu node [Zalando] [GPU]", f.WithSlow(), func(ctx context.Context) {
runGPUTest(ctx, f, cs, "gpu-test-", "nvcr.io/nvidia/k8s/cuda-sample:vectoradd-cuda12.5.0-ubi8", nil, "PASSED")
})

By("Creating a vector pod which runs on a GPU node")
pod := createVectorPod(nameprefix, ns, labels)
_, err := cs.CoreV1().Pods(ns).Create(ctx, pod, metav1.CreateOptions{})
framework.ExpectNoError(err, "Could not create POD %s", pod.Name)
framework.ExpectNoError(e2epod.WaitForPodSuccessInNamespaceTimeout(ctx, f.ClientSet, pod.Name, pod.Namespace, 15*time.Minute))
for {
p, err := cs.CoreV1().Pods(ns).Get(ctx, pod.Name, metav1.GetOptions{})
if err != nil {
framework.ExpectNoError(err, "Could not get POD %s", pod.Name)
return
}
if p.Status.ContainerStatuses[0].State.Terminated == nil {
time.Sleep(10 * time.Second)
continue
}
n := p.Status.ContainerStatuses[0].State.Terminated.ExitCode
if n != 0 {
framework.ExpectNoError(fmt.Errorf("expected POD %s to terminate with exit code 0", pod.Name))
return
}
logs, err := getPodLogs(cs, ns, pod.Name, "cuda-vector-add", false)
framework.ExpectNoError(err, "Should be able to get logs for pod %v", pod.Name)
regex := regexp.MustCompile("PASSED")
if regex.MatchString(logs) {
return
}
framework.ExpectNoError(err, "Expected vector job to succeed")
return
}
f.It("Should compile and run a CUDA kernel on a gpu node [Zalando] [GPU]", f.WithSlow(), func(ctx context.Context) {
runGPUTest(ctx, f, cs, "gpu-test-", "nvidia/cuda:13.2.1-devel-ubuntu24.04", []string{"bash", "-c", `cat > /tmp/t.cu <<EOF
#include <cstdio>
__global__ void add(int *a){ *a += 41; }
int main(){int *d,h=1; cudaMalloc(&d,4); cudaMemcpy(d,&h,4,cudaMemcpyHostToDevice);add<<<1,1>>>(d); cudaDeviceSynchronize();cudaError_t e=cudaGetLastError();if(e){ printf("FAIL: %s\n", cudaGetErrorString(e)); return 1; }cudaMemcpy(&h,d,4,cudaMemcpyDeviceToHost);printf("%s (result=%d)\n", h==42?"PASS":"FAIL", h); return h==42?0:1;}
EOF
nvcc --version | grep release
nvcc -o /tmp/t /tmp/t.cu && /tmp/t`}, "PASS")
})

f.It("Should run a PyTorch CUDA job on a gpu node [Zalando] [GPU]", f.WithSlow(), func(ctx context.Context) {
runGPUTest(ctx, f, cs, "gpu-test-", "pytorch/pytorch:2.12.1-cuda13.2-cudnn9-runtime", []string{"python", "-c",
"import torch; v=torch.version.cuda; assert torch.cuda.is_available(); " +
"assert tuple(map(int,v.split('.')))>=(13,2), f'CUDA {v} < 13.2'; " +
"x=torch.rand(3,device='cuda'); torch.cuda.synchronize(); " +
"print(f'PASS: torch {torch.__version__}, CUDA {v}, {torch.cuda.get_device_name()}')",
}, "PASS")
})
})

func runGPUTest(ctx context.Context, f *framework.Framework, cs kubernetes.Interface, nameprefix, image string, command []string, logPattern string) {
ns := f.Namespace.Name
labels := map[string]string{
"application": "vector-add",
}

By("Creating a vector pod which runs on a GPU node")
pod := createVectorPod(nameprefix, ns, labels, image, command)
_, err := cs.CoreV1().Pods(ns).Create(ctx, pod, metav1.CreateOptions{})
framework.ExpectNoError(err, "Could not create POD %s", pod.Name)
framework.ExpectNoError(e2epod.WaitForPodSuccessInNamespaceTimeout(ctx, f.ClientSet, pod.Name, pod.Namespace, 15*time.Minute))
for {
p, err := cs.CoreV1().Pods(ns).Get(ctx, pod.Name, metav1.GetOptions{})
if err != nil {
framework.ExpectNoError(err, "Could not get POD %s", pod.Name)
return
}
if p.Status.ContainerStatuses[0].State.Terminated == nil {
time.Sleep(10 * time.Second)
continue
}
n := p.Status.ContainerStatuses[0].State.Terminated.ExitCode
if n != 0 {
framework.ExpectNoError(fmt.Errorf("expected POD %s to terminate with exit code 0", pod.Name))
return
}
logs, err := getPodLogs(cs, ns, pod.Name, "main", false)
framework.ExpectNoError(err, "Should be able to get logs for pod %v", pod.Name)
regex := regexp.MustCompile(logPattern)
if regex.MatchString(logs) {
return
}
framework.ExpectNoError(err, "Expected vector job to succeed")
return
}
}
28 changes: 16 additions & 12 deletions test/e2e/util.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -912,7 +912,21 @@ func createVegetaDeployment(hostPath string, rate int) *appsv1.Deployment {

const NVIDIAGPUResourceName v1.ResourceName = "nvidia.com/gpu"

func createVectorPod(nameprefix, namespace string, labels map[string]string) *v1.Pod {
func createVectorPod(nameprefix, namespace string, labels map[string]string, image string, command []string) *v1.Pod {
container := v1.Container{
Name: "main",
Image: image,
Resources: v1.ResourceRequirements{
Limits: v1.ResourceList{
v1.ResourceCPU: resource.MustParse("100m"),
v1.ResourceMemory: resource.MustParse("1Gi"),
NVIDIAGPUResourceName: *resource.NewQuantity(1, resource.DecimalSI),
},
},
}
if len(command) > 0 {
container.Command = command
}
return &v1.Pod{
TypeMeta: metav1.TypeMeta{
Kind: "Pod",
Expand All @@ -925,17 +939,7 @@ func createVectorPod(nameprefix, namespace string, labels map[string]string) *v1
},
Spec: v1.PodSpec{
RestartPolicy: v1.RestartPolicyNever,
Containers: []v1.Container{
{
Name: "cuda-vector-add",
Image: "registry.k8s.io/cuda-vector-add:v0.1",
Resources: v1.ResourceRequirements{
Limits: v1.ResourceList{
NVIDIAGPUResourceName: *resource.NewQuantity(1, resource.DecimalSI),
},
},
},
},
Containers: []v1.Container{container},
NodeSelector: map[string]string{
"kubernetes.io/arch": "amd64",
},
Expand Down