Add support for tolerations in operator (#1648)

Summary: Adds support for tolerations to enable users to deploy PEMs to nodes that have taints through the CLI. Relevant Issues: #598 Type of change: /kind feature Test Plan: Used skaffold to deploy an unreleased version of the operator. Verified that tolerations parameter actually gets passed through to the pod spec. Signed-off-by: Benjamin Kilimnik <[email protected]>
pixie-io · Aug 4, 2023 · 6796db6 · 6796db6
1 parent c6a954a
commit 6796db6
Show file tree

Hide file tree

Showing 8 changed files with 701 additions and 74 deletions.
diff --git a/.arclint b/.arclint
@@ -19,6 +19,7 @@
     "(^private\/credentials\/.*\\.sh)",
     "(^private\/credentials\/.*\\.yaml)",
     "(^src/operator/client/versioned/)",
+    "(^src/operator/apis/px.dev/v1alpha1/zz_generated.deepcopy.go)",
     "(^src/stirling/bpf_tools/bcc_bpf/system-headers)",
     "(^src/stirling/mysql/testing/.*\\.json$)",
     "(^src/stirling/obj_tools/testdata/go/test_go_binary.go)",

diff --git a/k8s/operator/crd/base/px.dev_viziers.yaml b/k8s/operator/crd/base/px.dev_viziers.yaml
@@ -232,6 +232,49 @@ spec:
                         format: int64
                         type: integer
                     type: object
+                  tolerations:
+                    description: 'Tolerations allows scheduling pods on nodes with
+                      matching taints. More info: https://kubernetes.io/docs/concepts/scheduling-eviction/taint-and-toleration/:
+                      This field cannot be updated once the cluster is created.'
+                    items:
+                      description: The pod this Toleration is attached to tolerates
+                        any taint that matches the triple <key,value,effect> using
+                        the matching operator <operator>.
+                      properties:
+                        effect:
+                          description: Effect indicates the taint effect to match.
+                            Empty means match all taint effects. When specified, allowed
+                            values are NoSchedule, PreferNoSchedule and NoExecute.
+                          type: string
+                        key:
+                          description: Key is the taint key that the toleration applies
+                            to. Empty means match all taint keys. If the key is empty,
+                            operator must be Exists; this combination means to match
+                            all values and all keys.
+                          type: string
+                        operator:
+                          description: Operator represents a key's relationship to
+                            the value. Valid operators are Exists and Equal. Defaults
+                            to Equal. Exists is equivalent to wildcard for value,
+                            so that a pod can tolerate all taints of a particular
+                            category.
+                          type: string
+                        tolerationSeconds:
+                          description: TolerationSeconds represents the period of
+                            time the toleration (which must be of effect NoExecute,
+                            otherwise this field is ignored) tolerates the taint.
+                            By default, it is not set, which means tolerate the taint
+                            forever (do not evict). Zero and negative values will
+                            be treated as 0 (evict immediately) by the system.
+                          format: int64
+                          type: integer
+                        value:
+                          description: Value is the taint value the toleration matches
+                            to. If the operator is Exists, the value should be empty,
+                            otherwise just a regular string.
+                          type: string
+                      type: object
+                    type: array
                 type: object
               registry:
                 description: 'Registry specifies the image registry to use rather

diff --git a/src/api/proto/vizierconfigpb/vizier_types.pb.go b/src/api/proto/vizierconfigpb/vizier_types.pb.go
diff --git a/src/api/proto/vizierconfigpb/vizier_types.proto b/src/api/proto/vizierconfigpb/vizier_types.proto
@@ -23,6 +23,7 @@ package px.vizierconfigpb;
 option go_package = "vizierconfigpb";
 
 import "github.com/gogo/protobuf/gogoproto/gogo.proto";
+import "google/protobuf/wrappers.proto";
 
 // Note: Any changes to this file should be in sync with operatpro/api/v1alpha1/vizier_types.go.
 // VizierSpec defines the desired state of Vizier
@@ -90,6 +91,9 @@ message PodPolicyReq {
   // NodeSelector is a selector which must be true for the pod to fit on a node.
   // This field cannot be updated once the cluster is created.
   map<string, string> nodeSelector = 4;
+  // NodeSelector allows scheduling pods on nodes with matching taints.
+  // This field cannot be updated once the cluster is created.
+  repeated Toleration tolerations = 5;
 }
 
 // ResourceReqs is copied from the k8s api:
@@ -100,6 +104,16 @@ message ResourceReqs {
   ResourceList requests = 2;
 }
 
+// Toleration is copied from the k8s api:
+// https://pkg.go.dev/k8s.io/api/core/v1#Toleration
+message Toleration {
+  string key = 1;
+  string operator = 2;
+  string value = 3;
+  string effect = 4;
+  google.protobuf.Int64Value toleration_seconds = 5;
+}
+
 // ResourceList is copied from the k8s api: https://pkg.go.dev/k8s.io/api/core/v1#ResourceList
 message ResourceList {
   map<string, ResourceQuantity> resource_list = 1;

diff --git a/src/operator/apis/px.dev/v1alpha1/vizier_types.go b/src/operator/apis/px.dev/v1alpha1/vizier_types.go
@@ -187,6 +187,10 @@ type PodPolicy struct {
 	// More info: https://kubernetes.io/docs/concepts/configuration/assign-pod-node/
 	// This field cannot be updated once the cluster is created.
 	NodeSelector map[string]string `json:"nodeSelector,omitempty"`
+	// Tolerations allows scheduling pods on nodes with matching taints.
+	// More info: https://kubernetes.io/docs/concepts/scheduling-eviction/taint-and-toleration/:
+	// This field cannot be updated once the cluster is created.
+	Tolerations []v1.Toleration `json:"tolerations,omitempty"`
 	// The securityContext which should be set on non-privileged pods. All pods which require privileged permissions
 	// will still require a privileged securityContext.
 	SecurityContext *PodSecurityContext `json:"securityContext,omitempty"`

diff --git a/src/operator/apis/px.dev/v1alpha1/zz_generated.deepcopy.go b/src/operator/apis/px.dev/v1alpha1/zz_generated.deepcopy.go
diff --git a/src/operator/controllers/BUILD.bazel b/src/operator/controllers/BUILD.bazel
@@ -38,6 +38,7 @@ go_library(
         "//src/utils/shared/k8s",
         "@com_github_blang_semver//:semver",
         "@com_github_cenkalti_backoff_v4//:backoff",
+        "@com_github_gogo_protobuf//types",
         "@com_github_sirupsen_logrus//:logrus",
         "@io_k8s_api//apps/v1:apps",
         "@io_k8s_api//core/v1:core",

diff --git a/src/operator/controllers/vizier_controller.go b/src/operator/controllers/vizier_controller.go
@@ -31,6 +31,7 @@ import (
 
 	"github.com/blang/semver"
 	"github.com/cenkalti/backoff/v4"
+	"github.com/gogo/protobuf/types"
 	log "github.com/sirupsen/logrus"
 	"google.golang.org/grpc"
 	appsv1 "k8s.io/api/apps/v1"
@@ -784,7 +785,7 @@ func updateResourceConfiguration(resource *k8s.Resource, vz *v1alpha1.Vizier) er
 	addKeyValueMapToResource("labels", vz.Spec.Pod.Labels, resource.Object.Object)
 	addKeyValueMapToResource("annotations", vz.Spec.Pod.Annotations, resource.Object.Object)
 	updateResourceRequirements(vz.Spec.Pod.Resources, resource.Object.Object)
-	updatePodSpec(vz.Spec.Pod.NodeSelector, vz.Spec.Pod.SecurityContext, resource.Object.Object)
+	updatePodSpec(vz.Spec.Pod.NodeSelector, vz.Spec.Pod.Tolerations, vz.Spec.Pod.SecurityContext, resource.Object.Object)
 	return nil
 }
 
@@ -830,6 +831,7 @@ func generateVizierYAMLsConfig(ctx context.Context, ns string, k8sVersion string
 					Requests: convertResourceType(vz.Spec.Pod.Resources.Requests),
 				},
 				NodeSelector: vz.Spec.Pod.NodeSelector,
+				Tolerations:  convertTolerations(vz.Spec.Pod.Tolerations),
 			},
 			Patches:  vz.Spec.Patches,
 			Registry: vz.Spec.Registry,
@@ -956,7 +958,25 @@ func updateResourceRequirements(requirements v1.ResourceRequirements, res map[st
 		castedContainer["resources"] = resources
 	}
 }
-func updatePodSpec(nodeSelector map[string]string, securityCtx *v1alpha1.PodSecurityContext, res map[string]interface{}) {
+
+func convertTolerations(tolerations []v1.Toleration) []*vizierconfigpb.Toleration {
+	var castedTolerations []*vizierconfigpb.Toleration
+	for _, toleration := range tolerations {
+		castedToleration := &vizierconfigpb.Toleration{
+			Key:      toleration.Key,
+			Operator: string(toleration.Operator),
+			Value:    toleration.Value,
+			Effect:   string(toleration.Effect),
+		}
+		if toleration.TolerationSeconds != nil {
+			castedToleration.TolerationSeconds = &types.Int64Value{Value: *toleration.TolerationSeconds}
+		}
+		castedTolerations = append(castedTolerations, castedToleration)
+	}
+	return castedTolerations
+}
+
+func updatePodSpec(nodeSelector map[string]string, tolerations []v1.Toleration, securityCtx *v1alpha1.PodSecurityContext, res map[string]interface{}) {
 	podSpec := make(map[string]interface{})
 	md, ok, err := unstructured.NestedFieldNoCopy(res, "spec", "template", "spec")
 	if ok && err == nil {
@@ -977,6 +997,7 @@ func updatePodSpec(nodeSelector map[string]string, securityCtx *v1alpha1.PodSecu
 		castedNodeSelector[k] = v
 	}
 	podSpec["nodeSelector"] = castedNodeSelector
+	podSpec["tolerations"] = tolerations
 
 	// Add securityContext only if enabled.
 	if securityCtx == nil || !securityCtx.Enabled {