From e63d4e8d68f1998e09ba1693fb08cad89346629d Mon Sep 17 00:00:00 2001
From: jarradg <jarradg@nvidia.com>
Date: Tue, 4 Aug 2026 15:01:46 -0400
Subject: [PATCH] Tolerate reserved ARM TestKube nodes

---
 v1/providers/testkube/instance.go      |  2 ++
 v1/providers/testkube/instance_test.go | 12 +++++++++++-
 v1/providers/testkube/instancetype.go  | 17 +++++++++++++++++
 3 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/v1/providers/testkube/instance.go b/v1/providers/testkube/instance.go
index fd4c754..3681a49 100644
--- a/v1/providers/testkube/instance.go
+++ b/v1/providers/testkube/instance.go
@@ -75,6 +75,7 @@ func (c *TestKubeClient) CreateInstance(ctx context.Context, attrs cloudv1.Creat
 	return instance, nil
 }
 
+//nolint:funlen // Keep service and pod creation together so partial-resource cleanup remains explicit.
 func (c *TestKubeClient) createInstanceAsK8sResources(ctx context.Context, attrs cloudv1.CreateInstanceAttrs, instanceTypeSpec instanceTypeSpec) (*cloudv1.Instance, error) {
 	// Create a "cloud ID" to emulate a provider-provided instance ID.
 	cloudID := makeCloudID(c.refID, attrs.RefID)
@@ -134,6 +135,7 @@ func (c *TestKubeClient) createInstanceAsK8sResources(ctx context.Context, attrs
 				RestartPolicy:                 corev1.RestartPolicyNever,
 				TerminationGracePeriodSeconds: int64Ptr(1),
 				NodeSelector:                  maps.Clone(instanceTypeSpec.nodeSelector),
+				Tolerations:                   slices.Clone(instanceTypeSpec.tolerations),
 				Containers: []corev1.Container{
 					{
 						Name:            "vm",
diff --git a/v1/providers/testkube/instance_test.go b/v1/providers/testkube/instance_test.go
index 906e87c..d44eba5 100644
--- a/v1/providers/testkube/instance_test.go
+++ b/v1/providers/testkube/instance_test.go
@@ -193,13 +193,14 @@ func TestARM64InstanceUsesVersionedMultiarchImage(t *testing.T) {
 	require.Equal(t, "ghcr.io/brevdev/cloud/testkube-ubuntu-vm:multiarch-v2", pod.Spec.Containers[0].Image)
 }
 
-func TestInstanceUsesArchitectureNodeSelector(t *testing.T) {
+func TestInstanceUsesArchitectureScheduling(t *testing.T) {
 	ctx := context.Background()
 
 	tests := []struct {
 		name             string
 		instanceType     string
 		nodeArchitecture string
+		tolerations      []corev1.Toleration
 	}{
 		{
 			name:             "x86_64",
@@ -210,6 +211,14 @@ func TestInstanceUsesArchitectureNodeSelector(t *testing.T) {
 			name:             "arm64",
 			instanceType:     InstanceTypeOKCPUARM64,
 			nodeArchitecture: "arm64",
+			tolerations: []corev1.Toleration{
+				{
+					Key:      testKubeComputeTaintKey,
+					Operator: corev1.TolerationOpEqual,
+					Value:    testKubeComputeTaintValue,
+					Effect:   corev1.TaintEffectNoSchedule,
+				},
+			},
 		},
 	}
 
@@ -232,6 +241,7 @@ func TestInstanceUsesArchitectureNodeSelector(t *testing.T) {
 			require.Equal(t, map[string]string{
 				corev1.LabelArchStable: tt.nodeArchitecture,
 			}, pod.Spec.NodeSelector)
+			require.Equal(t, tt.tolerations, pod.Spec.Tolerations)
 		})
 	}
 }
diff --git a/v1/providers/testkube/instancetype.go b/v1/providers/testkube/instancetype.go
index 73b9c3a..79dd93c 100644
--- a/v1/providers/testkube/instancetype.go
+++ b/v1/providers/testkube/instancetype.go
@@ -23,6 +23,10 @@ const (
 	InstanceTypeFailCapacity = "test.fail.capacity"
 	InstanceTypeFailQuota    = "test.fail.quota"
 	InstanceTypeFailBuild    = "test.fail.build" // TODO: trigger build failure, maybe with a process that monitors build?
+
+	// Must match the dedicated ARM node-group taint in brevdev/control-plane-infra.
+	testKubeComputeTaintKey   = "testkube.brev.dev/compute-layer"
+	testKubeComputeTaintValue = "true"
 )
 
 // instanceTypeSpec is used mainly as a tuple of instance type (from devplane) and service type (from k8s). When a request
@@ -32,6 +36,7 @@ type instanceTypeSpec struct {
 	imageID      string
 	image        string
 	nodeSelector map[string]string
+	tolerations  []corev1.Toleration
 	serviceType  corev1.ServiceType
 }
 
@@ -53,6 +58,17 @@ func makeInstanceTypeSpec(
 	if architecture == cloudv1.ArchitectureX86_64 {
 		nodeArchitecture = "amd64"
 	}
+	var tolerations []corev1.Toleration
+	if architecture == cloudv1.ArchitectureARM64 {
+		tolerations = []corev1.Toleration{
+			{
+				Key:      testKubeComputeTaintKey,
+				Operator: corev1.TolerationOpEqual,
+				Value:    testKubeComputeTaintValue,
+				Effect:   corev1.TaintEffectNoSchedule,
+			},
+		}
+	}
 	return instanceTypeSpec{
 		instanceType: makeCPUInstanceType(instanceType, architecture, true, &estimatedDeployTime),
 		imageID:      imageID,
@@ -60,6 +76,7 @@ func makeInstanceTypeSpec(
 		nodeSelector: map[string]string{
 			corev1.LabelArchStable: nodeArchitecture,
 		},
+		tolerations: tolerations,
 		serviceType: corev1.ServiceTypeLoadBalancer,
 	}
 }