e2e: support GPU models on separate nodes

Author: msanft

e2e/gpu/gpu_test.go

@@ -74,9 +74,12 @@ func TestGPU(t *testing.T) {
 	require.NoError(t, err)
 
 	var resources []any
-	for _, config := range gpuConfigs {
+	for i, config := range gpuConfigs {
 		t.Logf("Using GPU resource %s (%s) with quantity %d", config.resource, config.model, config.quantity)
-		resources = append(resources, kuberesource.GPU(config.deploymentName(), string(config.resource), config.quantity)...)
+		// The hostpath CSI volume is node-local. Exercise the block-device
+		// regression once without preventing other GPU models from running on
+		// their respective nodes.
+		resources = append(resources, kuberesource.GPU(config.deploymentName(), string(config.resource), config.quantity, i == 0)...)
 	}
 
 	coordinator := kuberesource.CoordinatorBundle()

internal/kuberesource/parts.go

@@ -87,6 +87,20 @@ func NodeInstallers(namespace string, platforms []platforms.Platform) ([]*applya
 	return out, nil
 }
 
+func nvidiaGPUToleration() *applycorev1.TolerationApplyConfiguration {
+	return applycorev1.Toleration().
+		WithKey("nvidia.com/gpu").
+		WithOperator(corev1.TolerationOpExists).
+		WithEffect(corev1.TaintEffectNoSchedule)
+}
+
+func nodeInstallerTolerations(platform platforms.Platform) []*applycorev1.TolerationApplyConfiguration {
+	if !platforms.IsGPU(platform) {
+		return nil
+	}
+	return []*applycorev1.TolerationApplyConfiguration{nvidiaGPUToleration()}
+}
+
 // NodeInstaller constructs a node installer daemon set.
 func NodeInstaller(namespace string, platform platforms.Platform) (*applyappsv1.DaemonSetApplyConfiguration, error) {
 	runtimeHandler, err := manifest.RuntimeHandler(platform)
@@ -124,6 +138,7 @@ func NodeInstaller(namespace string, platform platforms.Platform) (*applyappsv1.
 						WithSpec(
 							PodSpec().
 								WithHostPID(true).
+								WithTolerations(nodeInstallerTolerations(platform)...).
 								WithInitContainers(
 									Container().
 										WithName("installer").

internal/kuberesource/parts_test.go

@@ -4,11 +4,69 @@
 package kuberesource
 
 import (
+	"fmt"
 	"testing"
 
+	"github.com/edgelesssys/contrast/internal/platforms"
 	"github.com/stretchr/testify/require"
+	corev1 "k8s.io/api/core/v1"
+	applyappsv1 "k8s.io/client-go/applyconfigurations/apps/v1"
 )
 
+func TestNodeInstallerGPUToleration(t *testing.T) {
+	for _, tc := range []struct {
+		name            string
+		platform        platforms.Platform
+		wantTolerations int
+	}{
+		{
+			name:            "GPU",
+			platform:        platforms.MetalQEMUTDXGPU,
+			wantTolerations: 1,
+		},
+		{
+			name:     "non-GPU",
+			platform: platforms.MetalQEMUTDX,
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			require := require.New(t)
+			tolerations := nodeInstallerTolerations(tc.platform)
+			require.Len(tolerations, tc.wantTolerations)
+			if tc.wantTolerations == 0 {
+				return
+			}
+			require.Equal("nvidia.com/gpu", *tolerations[0].Key)
+			require.Equal(corev1.TolerationOpExists, *tolerations[0].Operator)
+			require.Equal(corev1.TaintEffectNoSchedule, *tolerations[0].Effect)
+		})
+	}
+}
+
+func TestGPUBlockDevice(t *testing.T) {
+	for _, withBlockDevice := range []bool{false, true} {
+		t.Run(fmt.Sprintf("withBlockDevice=%t", withBlockDevice), func(t *testing.T) {
+			require := require.New(t)
+			resources := GPU("gpu-tester", "nvidia.com/test", 1, withBlockDevice)
+			require.Len(resources, 1)
+
+			deployment, ok := resources[0].(*applyappsv1.DeploymentApplyConfiguration)
+			require.True(ok)
+			spec := deployment.Spec.Template.Spec
+			require.Len(spec.Containers, 3)
+			noGPUContainer := spec.Containers[2]
+			require.Equal("no-gpu", *noGPUContainer.Name)
+
+			wantBlockDevices := 0
+			if withBlockDevice {
+				wantBlockDevices = 1
+			}
+			require.Len(spec.Volumes, wantBlockDevices)
+			require.Len(noGPUContainer.VolumeDevices, wantBlockDevices)
+		})
+	}
+}
+
 func TestNewPortForwarder(t *testing.T) {
 	require := require.New(t)
 

internal/kuberesource/resourcegen/main.go

@@ -87,7 +87,7 @@ func main() {
 		case "vault":
 			subResources = kuberesource.PatchRuntimeHandlers(kuberesource.Vault(*namespace), "contrast-cc")
 		case "gpu":
-			subResources = kuberesource.PatchRuntimeHandlers(kuberesource.GPU("gpu-tester", *gpuClass, int64(*gpuCount)), "contrast-cc")
+			subResources = kuberesource.PatchRuntimeHandlers(kuberesource.GPU("gpu-tester", *gpuClass, int64(*gpuCount), true), "contrast-cc")
 		default:
 			log.Fatalf("Error: unknown set: %s\n", set)
 		}

internal/kuberesource/sets.go

@@ -804,7 +804,8 @@ done
 //
 // gpuClass must be a vendor/class pair, like nvidia.com/GB100_B200.
 // gpuQuantity is the number of GPUs.
-func GPU(name string, gpuClass string, gpuQuantity int64) []any {
+// withBlockDevice adds a block device to the non-GPU container as a regression test.
+func GPU(name string, gpuClass string, gpuQuantity int64, withBlockDevice bool) []any {
 	component := "gpu-test"
 	tester := Deployment(name, "").
 		WithSpec(
@@ -816,6 +817,7 @@ func GPU(name string, gpuClass string, gpuQuantity int64) []any {
 						WithLabels(SelectorLabels(name, component)).
 						WithSpec(
 							PodSpec().
+								WithTolerations(nvidiaGPUToleration()).
 								WithContainers(
 									Container().
 										WithName("gpu-tester-direct"). // This container directly requests a GPU.
@@ -876,6 +878,16 @@ func GPU(name string, gpuClass string, gpuQuantity int64) []any {
 						),
 				),
 		)
+	if !withBlockDevice {
+		tester.Spec.Template.Spec.Volumes = nil
+		for i := range tester.Spec.Template.Spec.Containers {
+			container := &tester.Spec.Template.Spec.Containers[i]
+			if container.Name != nil && *container.Name == "no-gpu" {
+				container.VolumeDevices = nil
+				break
+			}
+		}
+	}
 
 	return []any{tester}
 }