Add e2e tests for tpu-provisioner

tpu-provisioner/.dockerignore

@@ -2,3 +2,4 @@
 # Ignore build and test binaries.
 bin/
 testbin/
+e2e/

tpu-provisioner/Makefile

@@ -56,6 +56,12 @@ vet: ## Run go vet against code.
 test: manifests fmt vet envtest ## Run tests.
 	KUBEBUILDER_ASSETS="$(shell $(ENVTEST) use $(ENVTEST_K8S_VERSION) --bin-dir $(LOCALBIN) -p path)" go test ./... -v -coverprofile cover.out
 
+.PHONY: test-e2e
+test-e2e:
+	./test/e2e/scripts/up.sh
+	./test/e2e/scripts/run.sh
+	./test/e2e/scripts/down.sh
+
 ##@ Build
 
 .PHONY: build

tpu-provisioner/README.md

@@ -98,3 +98,23 @@ Test that you can apply a TPU Job.
 ```bash
 kubectl apply -f ./examples/v4-2x2x4/
 ```
+
+## Testing
+
+### Unit / Integration Tests
+
+Run unit and integration tests.
+
+```bash
+make test
+```
+
+### End-to-End Tests
+
+Set environment variables (see example here: `./test/e2e/env/example.sh`).
+
+Run end-to-end tests.
+
+```bash
+make test-e2e
+```

tpu-provisioner/internal/cloud/gke.go

@@ -204,6 +204,7 @@ func (g *GKE) nodePoolForPod(name string, p *corev1.Pod) (*containerv1beta1.Node
 				EnableIntegrityMonitoring: true,
 				EnableSecureBoot:          true,
 			},
+			Spot: p.Spec.NodeSelector["cloud.google.com/gke-spot"] == "true",
 			Tags: g.ClusterContext.NodeTags,
 			// NOTE: vendor/ was manually updated to include the field because
 			// it was not currently available at the time of writing:

tpu-provisioner/test/e2e/config/configmap-patch.yaml

@@ -0,0 +1,8 @@
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: manager
+  namespace: tpu-provisioner-system
+data:
+  GCP_ZONE: ${ZONE}
+  GCP_NODE_SERVICE_ACCOUNT: ${NODE_SA_EMAIL}

tpu-provisioner/test/e2e/config/deployment-patch.yaml

@@ -0,0 +1,12 @@
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: controller-manager
+  namespace: tpu-provisioner-system
+spec:
+  template:
+    spec:
+      containers:
+        - name: manager
+          image: ${CONTROLLER_IMAGE}
+          imagePullPolicy: Always

tpu-provisioner/test/e2e/config/kustomization.yaml

@@ -0,0 +1,8 @@
+resources:
+- ../../../config/default
+patches:
+- path: ./deployment-patch.yaml
+- path: ./configmap-patch.yaml
+- path: ./serviceaccount-patch.yaml
+apiVersion: kustomize.config.k8s.io/v1beta1
+kind: Kustomization

tpu-provisioner/test/e2e/config/serviceaccount-patch.yaml

@@ -0,0 +1,7 @@
+apiVersion: v1
+kind: ServiceAccount
+metadata:
+  name: controller-manager
+  namespace: tpu-provisioner-system
+  annotations:
+    iam.gke.io/gcp-service-account: ${CONTROLLER_SA_EMAIL}

tpu-provisioner/test/e2e/env/example.sh

@@ -0,0 +1,16 @@
+export PROJECT_ID="my-project"
+export PROJECT_NUMBER="$(gcloud projects describe ${PROJECT_ID} --format='value(projectNumber)')"
+
+export CONTROLLER_SA_NAME="tpu-provisioner"
+export CONTROLLER_SA_EMAIL="${CONTROLLER_SA_NAME}@${PROJECT_ID}.iam.gserviceaccount.com"
+
+export REGION=us-central2
+export ZONE=us-central2-b
+export CLUSTER_NAME="tpu-provisioner-e2e"
+
+export BUILD_TOOL="podman" # or "docker"
+
+# Assumes you created a regional Google Artifact Registry instance.
+export CONTROLLER_IMAGE="${REGION}-docker.pkg.dev/${PROJECT_ID}/default/tpu-provisioner:$(git describe --tags --dirty --always)"
+
+export NODE_SA_EMAIL="${PROJECT_NUMBER}[email protected]"

tpu-provisioner/test/e2e/go.mod

@@ -0,0 +1,65 @@
+module e2e
+
+go 1.22.0
+
+require (
+	github.com/beorn7/perks v1.0.1 // indirect
+	github.com/cespare/xxhash/v2 v2.2.0 // indirect
+	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/emicklei/go-restful/v3 v3.11.0 // indirect
+	github.com/evanphx/json-patch/v5 v5.8.0 // indirect
+	github.com/fsnotify/fsnotify v1.7.0 // indirect
+	github.com/go-logr/logr v1.4.1 // indirect
+	github.com/go-openapi/jsonpointer v0.19.6 // indirect
+	github.com/go-openapi/jsonreference v0.20.2 // indirect
+	github.com/go-openapi/swag v0.22.3 // indirect
+	github.com/gogo/protobuf v1.3.2 // indirect
+	github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
+	github.com/golang/protobuf v1.5.3 // indirect
+	github.com/google/gnostic-models v0.6.8 // indirect
+	github.com/google/go-cmp v0.6.0 // indirect
+	github.com/google/gofuzz v1.2.0 // indirect
+	github.com/google/uuid v1.3.0 // indirect
+	github.com/imdario/mergo v0.3.16 // indirect
+	github.com/josharian/intern v1.0.0 // indirect
+	github.com/json-iterator/go v1.1.12 // indirect
+	github.com/mailru/easyjson v0.7.7 // indirect
+	github.com/matttproud/golang_protobuf_extensions/v2 v2.0.0 // indirect
+	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
+	github.com/modern-go/reflect2 v1.0.2 // indirect
+	github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
+	github.com/pkg/errors v0.9.1 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+	github.com/prometheus/client_golang v1.18.0 // indirect
+	github.com/prometheus/client_model v0.5.0 // indirect
+	github.com/prometheus/common v0.45.0 // indirect
+	github.com/prometheus/procfs v0.12.0 // indirect
+	github.com/spf13/pflag v1.0.5 // indirect
+	github.com/stretchr/testify v1.8.4 // indirect
+	golang.org/x/exp v0.0.0-20220722155223-a9213eeb770e // indirect
+	golang.org/x/net v0.19.0 // indirect
+	golang.org/x/oauth2 v0.12.0 // indirect
+	golang.org/x/sys v0.16.0 // indirect
+	golang.org/x/term v0.15.0 // indirect
+	golang.org/x/text v0.14.0 // indirect
+	golang.org/x/time v0.3.0 // indirect
+	gomodules.xyz/jsonpatch/v2 v2.4.0 // indirect
+	google.golang.org/appengine v1.6.7 // indirect
+	google.golang.org/protobuf v1.31.0 // indirect
+	gopkg.in/inf.v0 v0.9.1 // indirect
+	gopkg.in/yaml.v2 v2.4.0 // indirect
+	gopkg.in/yaml.v3 v3.0.1 // indirect
+	k8s.io/api v0.29.2 // indirect
+	k8s.io/apiextensions-apiserver v0.29.0 // indirect
+	k8s.io/apimachinery v0.29.2 // indirect
+	k8s.io/client-go v0.29.2 // indirect
+	k8s.io/component-base v0.29.0 // indirect
+	k8s.io/klog/v2 v2.110.1 // indirect
+	k8s.io/kube-openapi v0.0.0-20231010175941-2dd684a91f00 // indirect
+	k8s.io/utils v0.0.0-20240102154912-e7106e64919e // indirect
+	sigs.k8s.io/controller-runtime v0.17.2 // indirect
+	sigs.k8s.io/jobset v0.3.2 // indirect
+	sigs.k8s.io/json v0.0.0-20221116044647-bc3834ca7abd // indirect
+	sigs.k8s.io/structured-merge-diff/v4 v4.4.1 // indirect
+	sigs.k8s.io/yaml v1.4.0 // indirect
+)