[WIP] Adding taint during image pre-pulling to prevent user pod being scheduled

chartpress.yaml

@@ -58,3 +58,7 @@ charts:
       # singleuser-sample, a primitive user container to start with.
       singleuser-sample:
         valuesPath: singleuser.image
+
+      # taint-manager, an initContainer to add/remove taint for a node
+      taint-manager:
+        valuesPath: prePuller.taintmanager.image

images/taint-manager/.gitignore

@@ -0,0 +1,2 @@
+.idea
+taintmanager

images/taint-manager/Dockerfile

@@ -0,0 +1,25 @@
+# syntax=docker/dockerfile:1
+
+## Build
+FROM golang:1.18-bullseye AS build
+
+WORKDIR /app
+
+COPY go.mod ./
+COPY go.sum ./
+RUN go mod download
+
+COPY *.go ./
+
+RUN go build -o /taintmanager
+
+## Deploy
+FROM gcr.io/distroless/base-debian11
+
+WORKDIR /
+
+COPY --from=build /taintmanager /taintmanager
+
+USER nonroot:nonroot
+
+CMD ["/taintmanager"]

images/taint-manager/README.md

@@ -0,0 +1,29 @@
+# In Cluster Taint Manager
+
+To add or remove taint within a pod. The default mode is to run as a pod in a Kubernetes cluster. It will try to
+authenticate to K8S API server with in-cluster config (the service account token mounted inside pod). It also uses
+downward API to retrieve the node name where the pod is running on in order to change the taint.
+
+## Compile
+
+```
+GOOS=linux GOARCH=amd64 go build -o taintmanager taintmanager.go
+```
+
+## Development and Debug
+
+To test outside a cluster, commandline parameter `-kubeconfig` and `-node` can be specified. If `-kubeconfig` is not
+specified, taintmanager will try to load the kubeconfig from default path `HOME/.kube/config`
+
+Using default kubeconfig:
+
+```
+taintmanager -node f6.workers.ctlt.ubc.ca -remove hub.jupyter.org/imagepulling:NoExecute
+```
+
+## Test in Cluster
+
+The `test` directory contains YAML files for deploy a pod with required permissions to run taintmanager.
+Please change `namespace` field in `clusterrolebinding.yaml` before deploying to a cluster.
+
+After deploying yaml files, run `kubectl cp` to copy compiled binary to the target pod and run the binary inside pod.

images/taint-manager/go.mod

@@ -0,0 +1,48 @@
+module taintmanager
+
+go 1.18
+
+require (
+	k8s.io/api v0.25.4
+	k8s.io/apimachinery v0.25.4
+	k8s.io/client-go v0.25.4
+)
+
+require (
+	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/emicklei/go-restful/v3 v3.10.1 // indirect
+	github.com/go-logr/logr v1.2.3 // indirect
+	github.com/go-openapi/jsonpointer v0.19.5 // indirect
+	github.com/go-openapi/jsonreference v0.20.0 // indirect
+	github.com/go-openapi/swag v0.22.3 // indirect
+	github.com/gogo/protobuf v1.3.2 // indirect
+	github.com/golang/protobuf v1.5.2 // indirect
+	github.com/google/gnostic v0.6.9 // indirect
+	github.com/google/go-cmp v0.5.9 // indirect
+	github.com/google/gofuzz v1.2.0 // indirect
+	github.com/imdario/mergo v0.3.6 // 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/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/spf13/pflag v1.0.5 // indirect
+	golang.org/x/net v0.2.0 // indirect
+	golang.org/x/oauth2 v0.2.0 // indirect
+	golang.org/x/sys v0.2.0 // indirect
+	golang.org/x/term v0.2.0 // indirect
+	golang.org/x/text v0.4.0 // indirect
+	golang.org/x/time v0.2.0 // indirect
+	google.golang.org/appengine v1.6.7 // indirect
+	google.golang.org/protobuf v1.28.1 // 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/klog/v2 v2.80.1 // indirect
+	k8s.io/kube-openapi v0.0.0-20221123214604-86e75ddd809a // indirect
+	k8s.io/utils v0.0.0-20221128185143-99ec85e7a448 // indirect
+	sigs.k8s.io/json v0.0.0-20221116044647-bc3834ca7abd // indirect
+	sigs.k8s.io/structured-merge-diff/v4 v4.2.3 // indirect
+	sigs.k8s.io/yaml v1.3.0 // indirect
+)

images/taint-manager/taintmanager.go

@@ -0,0 +1,177 @@
+package main
+
+import (
+	"context"
+	"flag"
+	"fmt"
+	v1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/api/errors"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/util/validation"
+	"k8s.io/client-go/kubernetes"
+	"k8s.io/client-go/rest"
+	"k8s.io/client-go/tools/clientcmd"
+	"k8s.io/client-go/util/homedir"
+	"log"
+	"os"
+	"path/filepath"
+	"strings"
+)
+
+func buildConfig(kubeconfig string) (*rest.Config, error) {
+	// load kubeconfig from command line
+	if kubeconfig != "" {
+		cfg, err := clientcmd.BuildConfigFromFlags("", kubeconfig)
+		if err != nil {
+			return nil, err
+		}
+		log.Printf("Using kubeconfig %v.", kubeconfig)
+		return cfg, nil
+	}
+
+	// try the default location HOME/.kube/config
+	if home := homedir.HomeDir(); home != "" {
+		kubeconfig = filepath.Join(home, ".kube", "config")
+		cfg, err := clientcmd.BuildConfigFromFlags("", kubeconfig)
+		if err == nil {
+			log.Printf("Using kubeconfig %v.", kubeconfig)
+			return cfg, nil
+		}
+	}
+
+	// try in-cluster auth
+	cfg, err := rest.InClusterConfig()
+	if err != nil {
+		return nil, err
+	}
+	log.Printf("Using in-cluster config.")
+	return cfg, nil
+}
+
+func usage() {
+	log.Println("Usage: taintmanager [-add TAINT] [-remove TAINT] [-node NODE_NAME] [-kubeconfig /PATH/TO/KUBECONFIG]")
+	flag.PrintDefaults()
+	os.Exit(1)
+}
+
+func main() {
+	kubeconfig := flag.String("kubeconfig", "", "(optional) absolute path to the kubeconfig file")
+	nodeName := flag.String("node", "", "(optional) The name of the node to add/remove taints. Can be passed through environment variable MY_NODE_NAME")
+	taintAdd := flag.String("add", "", "The taint to add")
+	taintRemove := flag.String("remove", "", "The taint to remove")
+	flag.Parse()
+
+	// load kubeconfig
+	config, err := buildConfig(*kubeconfig)
+	if err != nil {
+		log.Printf(err.Error())
+		log.Printf("Failed to load kubeconfig")
+		usage()
+	}
+
+	// parse taints
+	tAdd, errAdd := parseTaint(*taintAdd)
+	tRemove, errRemove := parseTaint(*taintRemove)
+	if errAdd != nil && errRemove != nil {
+		log.Println("Please specify at least one option -add or -remove")
+		usage()
+	}
+
+	// check node name
+	if *nodeName == "" {
+		// try from env var
+		*nodeName = os.Getenv("MY_NODE_NAME")
+	}
+	if *nodeName == "" {
+		log.Println("Please specify the node name")
+		usage()
+	}
+	log.Printf("Node name %v\n", *nodeName)
+
+	// creates the clientset
+	clientset, err := kubernetes.NewForConfig(config)
+	if err != nil {
+		panic(err.Error())
+	}
+	node, err := clientset.CoreV1().Nodes().Get(context.TODO(), *nodeName, metav1.GetOptions{})
+	if errors.IsNotFound(err) {
+		log.Printf("Node %v not found in default namespace\n", *nodeName)
+	} else if statusError, isStatus := err.(*errors.StatusError); isStatus {
+		log.Printf("Error getting node %v\n", statusError.ErrStatus.Message)
+	} else if err != nil {
+		panic(err.Error())
+	} else {
+		log.Printf("Found node %v\n", node.GetName())
+		for k, v := range node.Spec.Taints {
+			log.Printf("%v: %v\n", k, v)
+		}
+		if errAdd == nil {
+			node.Spec.Taints = append(node.Spec.Taints, tAdd)
+			clientset.CoreV1().Nodes().Update(context.TODO(), node, metav1.UpdateOptions{})
+			log.Printf("Taint %v is added to node %v.", tAdd.ToString(), *nodeName)
+		}
+		if errRemove == nil {
+			for i, taint := range node.Spec.Taints {
+				if taint.Key == tRemove.Key {
+					node.Spec.Taints = append(node.Spec.Taints[:i], node.Spec.Taints[i+1:]...)
+					log.Printf("Taint %v is removed from node %v.", tRemove.ToString(), *nodeName)
+					clientset.CoreV1().Nodes().Update(context.TODO(), node, metav1.UpdateOptions{})
+				}
+			}
+		}
+	}
+}
+
+// copied from https://github.com/kubernetes/kubernetes/blob/v1.25.4/pkg/util/taints/taints.go
+// parseTaint parses a taint from a string, whose form must be either
+// '<key>=<value>:<effect>', '<key>:<effect>', or '<key>'.
+func parseTaint(st string) (v1.Taint, error) {
+	var taint v1.Taint
+
+	var key string
+	var value string
+	var effect v1.TaintEffect
+
+	parts := strings.Split(st, ":")
+	switch len(parts) {
+	case 1:
+		key = parts[0]
+	case 2:
+		effect = v1.TaintEffect(parts[1])
+		if err := validateTaintEffect(effect); err != nil {
+			return taint, err
+		}
+
+		partsKV := strings.Split(parts[0], "=")
+		if len(partsKV) > 2 {
+			return taint, fmt.Errorf("invalid taint spec: %v", st)
+		}
+		key = partsKV[0]
+		if len(partsKV) == 2 {
+			value = partsKV[1]
+			if errs := validation.IsValidLabelValue(value); len(errs) > 0 {
+				return taint, fmt.Errorf("invalid taint spec: %v, %s", st, strings.Join(errs, "; "))
+			}
+		}
+	default:
+		return taint, fmt.Errorf("invalid taint spec: %v", st)
+	}
+
+	if errs := validation.IsQualifiedName(key); len(errs) > 0 {
+		return taint, fmt.Errorf("invalid taint spec: %v, %s", st, strings.Join(errs, "; "))
+	}
+
+	taint.Key = key
+	taint.Value = value
+	taint.Effect = effect
+
+	return taint, nil
+}
+
+func validateTaintEffect(effect v1.TaintEffect) error {
+	if effect != v1.TaintEffectNoSchedule && effect != v1.TaintEffectPreferNoSchedule && effect != v1.TaintEffectNoExecute {
+		return fmt.Errorf("invalid taint effect: %v, unsupported taint effect", effect)
+	}
+
+	return nil
+}

images/taint-manager/test/clusterrole.yaml

@@ -0,0 +1,8 @@
+apiVersion: rbac.authorization.k8s.io/v1
+kind: ClusterRole
+metadata:
+  name: taintmanager
+rules:
+  - apiGroups: [""] # "" indicates the core API group
+    resources: ["nodes"]
+    verbs: ["get", "update"]

images/taint-manager/test/clusterrolebinding.yaml

@@ -0,0 +1,16 @@
+apiVersion: rbac.authorization.k8s.io/v1
+# This role binding allows "jane" to read pods in the "default" namespace.
+# You need to already have a Role named "pod-reader" in that namespace.
+kind: ClusterRoleBinding
+metadata:
+  name: taintmanager
+subjects:
+  # You can specify more than one "subject"
+  - kind: ServiceAccount
+    name: taintmanager
+    namespace: default
+roleRef:
+  # "roleRef" specifies the binding to a Role / ClusterRole
+  kind: ClusterRole #this must be Role or ClusterRole
+  name: taintmanager # this must match the name of the Role or ClusterRole you wish to bind to
+  apiGroup: rbac.authorization.k8s.io

images/taint-manager/test/deployment-busybox.yaml

@@ -0,0 +1,36 @@
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: taintmanager
+spec:
+  replicas: 1
+  selector:
+    matchLabels:
+      app: taintmanager
+  template:
+    metadata:
+      labels:
+        app: taintmanager
+    spec:
+      serviceAccountName: taintmanager
+      containers:
+        - image: busybox
+          imagePullPolicy: IfNotPresent
+          name: taintmanager
+          command: ["sleep", "100000"]
+          env:
+            - name: MY_POD_NAME
+              valueFrom:
+                fieldRef:
+                  fieldPath: metadata.name
+            - name: MY_NODE_NAME
+              valueFrom:
+                fieldRef:
+                  fieldPath: spec.nodeName
+          ports:
+            - containerPort: 8080
+      tolerations:
+        - key: "hub.jupyter.org/dedicated"
+          operator: "Exists"
+          effect: "NoSchedule"
+      restartPolicy: Always

images/taint-manager/test/sa.yaml

@@ -0,0 +1,5 @@
+apiVersion: v1
+kind: ServiceAccount
+metadata:
+  name: taintmanager
+  namespace: default