Skip to content

Commit f312fc9

Browse files
mikkeloscarmikkeloscar
committed
e2e: Check for GPU Driver/ML framework compatibilty
Signed-off-by: Mikkel Oscar Lyderik Larsen <mikkel.larsen@zalando.de>
1 parent 34cffc9 commit f312fc9

2 files changed

Lines changed: 73 additions & 47 deletions

File tree

‎test/e2e/gpu.go‎

Lines changed: 57 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -37,41 +37,63 @@ var _ = describe("GPU job processing", func() {
3737
cs = f.ClientSet
3838
})
3939

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

47-
By("Creating a vector pod which runs on a GPU node")
48-
pod := createVectorPod(nameprefix, ns, labels)
49-
_, err := cs.CoreV1().Pods(ns).Create(ctx, pod, metav1.CreateOptions{})
50-
framework.ExpectNoError(err, "Could not create POD %s", pod.Name)
51-
framework.ExpectNoError(e2epod.WaitForPodSuccessInNamespaceTimeout(ctx, f.ClientSet, pod.Name, pod.Namespace, 15*time.Minute))
52-
for {
53-
p, err := cs.CoreV1().Pods(ns).Get(ctx, pod.Name, metav1.GetOptions{})
54-
if err != nil {
55-
framework.ExpectNoError(err, "Could not get POD %s", pod.Name)
56-
return
57-
}
58-
if p.Status.ContainerStatuses[0].State.Terminated == nil {
59-
time.Sleep(10 * time.Second)
60-
continue
61-
}
62-
n := p.Status.ContainerStatuses[0].State.Terminated.ExitCode
63-
if n != 0 {
64-
framework.ExpectNoError(fmt.Errorf("expected POD %s to terminate with exit code 0", pod.Name))
65-
return
66-
}
67-
logs, err := getPodLogs(cs, ns, pod.Name, "cuda-vector-add", false)
68-
framework.ExpectNoError(err, "Should be able to get logs for pod %v", pod.Name)
69-
regex := regexp.MustCompile("PASSED")
70-
if regex.MatchString(logs) {
71-
return
72-
}
73-
framework.ExpectNoError(err, "Expected vector job to succeed")
74-
return
75-
}
44+
f.It("Should compile and run a CUDA kernel on a gpu node [Zalando] [GPU]", f.WithSlow(), func(ctx context.Context) {
45+
runGPUTest(ctx, f, cs, "gpu-test-", "nvidia/cuda:13.2.1-devel-ubuntu24.04", []string{"bash", "-c", `cat > /tmp/t.cu <<EOF
46+
#include <cstdio>
47+
__global__ void add(int *a){ *a += 41; }
48+
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;}
49+
EOF
50+
nvcc --version | grep release
51+
nvcc -o /tmp/t /tmp/t.cu && /tmp/t`}, "PASS")
52+
})
53+
54+
f.It("Should run a PyTorch CUDA job on a gpu node [Zalando] [GPU]", f.WithSlow(), func(ctx context.Context) {
55+
runGPUTest(ctx, f, cs, "gpu-test-", "pytorch/pytorch:2.12.1-cuda13.2-cudnn9-runtime", []string{"python", "-c",
56+
"import torch; v=torch.version.cuda; assert torch.cuda.is_available(); " +
57+
"assert tuple(map(int,v.split('.')))>=(13,2), f'CUDA {v} < 13.2'; " +
58+
"x=torch.rand(3,device='cuda'); torch.cuda.synchronize(); " +
59+
"print(f'PASS: torch {torch.__version__}, CUDA {v}, {torch.cuda.get_device_name()}')",
60+
}, "PASS")
7661
})
7762
})
63+
64+
func runGPUTest(ctx context.Context, f *framework.Framework, cs kubernetes.Interface, nameprefix, image string, command []string, logPattern string) {
65+
ns := f.Namespace.Name
66+
labels := map[string]string{
67+
"application": "vector-add",
68+
}
69+
70+
By("Creating a vector pod which runs on a GPU node")
71+
pod := createVectorPod(nameprefix, ns, labels, image, command)
72+
_, err := cs.CoreV1().Pods(ns).Create(ctx, pod, metav1.CreateOptions{})
73+
framework.ExpectNoError(err, "Could not create POD %s", pod.Name)
74+
framework.ExpectNoError(e2epod.WaitForPodSuccessInNamespaceTimeout(ctx, f.ClientSet, pod.Name, pod.Namespace, 15*time.Minute))
75+
for {
76+
p, err := cs.CoreV1().Pods(ns).Get(ctx, pod.Name, metav1.GetOptions{})
77+
if err != nil {
78+
framework.ExpectNoError(err, "Could not get POD %s", pod.Name)
79+
return
80+
}
81+
if p.Status.ContainerStatuses[0].State.Terminated == nil {
82+
time.Sleep(10 * time.Second)
83+
continue
84+
}
85+
n := p.Status.ContainerStatuses[0].State.Terminated.ExitCode
86+
if n != 0 {
87+
framework.ExpectNoError(fmt.Errorf("expected POD %s to terminate with exit code 0", pod.Name))
88+
return
89+
}
90+
logs, err := getPodLogs(cs, ns, pod.Name, "main", false)
91+
framework.ExpectNoError(err, "Should be able to get logs for pod %v", pod.Name)
92+
regex := regexp.MustCompile(logPattern)
93+
if regex.MatchString(logs) {
94+
return
95+
}
96+
framework.ExpectNoError(err, "Expected vector job to succeed")
97+
return
98+
}
99+
}

‎test/e2e/util.go‎

Lines changed: 16 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -912,7 +912,21 @@ func createVegetaDeployment(hostPath string, rate int) *appsv1.Deployment {
912912

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

915-
func createVectorPod(nameprefix, namespace string, labels map[string]string) *v1.Pod {
915+
func createVectorPod(nameprefix, namespace string, labels map[string]string, image string, command []string) *v1.Pod {
916+
container := v1.Container{
917+
Name: "main",
918+
Image: image,
919+
Resources: v1.ResourceRequirements{
920+
Limits: v1.ResourceList{
921+
v1.ResourceCPU: resource.MustParse("100m"),
922+
v1.ResourceMemory: resource.MustParse("1Gi"),
923+
NVIDIAGPUResourceName: *resource.NewQuantity(1, resource.DecimalSI),
924+
},
925+
},
926+
}
927+
if len(command) > 0 {
928+
container.Command = command
929+
}
916930
return &v1.Pod{
917931
TypeMeta: metav1.TypeMeta{
918932
Kind: "Pod",
@@ -925,17 +939,7 @@ func createVectorPod(nameprefix, namespace string, labels map[string]string) *v1
925939
},
926940
Spec: v1.PodSpec{
927941
RestartPolicy: v1.RestartPolicyNever,
928-
Containers: []v1.Container{
929-
{
930-
Name: "cuda-vector-add",
931-
Image: "registry.k8s.io/cuda-vector-add:v0.1",
932-
Resources: v1.ResourceRequirements{
933-
Limits: v1.ResourceList{
934-
NVIDIAGPUResourceName: *resource.NewQuantity(1, resource.DecimalSI),
935-
},
936-
},
937-
},
938-
},
942+
Containers: []v1.Container{container},
939943
NodeSelector: map[string]string{
940944
"kubernetes.io/arch": "amd64",
941945
},

0 commit comments

Comments
 (0)