add leader election retry

.github/workflows/tests.yml

@@ -27,8 +27,14 @@ jobs:
         run: make test-integration
 
   e2e:
+    strategy:
+      matrix:
+        replicas: ["1", "3"]
+        test_cases:
+          - { requests: 60, expected_replicas: 1 }
+          - { requests: 300, expected_replicas: 2 }
     runs-on: ubuntu-latest
-
+    name: E2E Lingo.replicas=${{ matrix.replicas }} requests=${{ matrix.test_cases.requests }} expected_replicas=${{ matrix.test_cases.expected_replicas }}
     steps:
       - name: Checkout code
         uses: actions/checkout@v2
@@ -49,4 +55,8 @@ jobs:
           sudo mv skaffold /usr/local/bin
 
       - name: Run e2e tests
+        env:
+          REPLICAS: ${{ matrix.replicas }}
+          REQUESTS: ${{ matrix.test_cases.requests }}
+          EXPECTED_REPLICAS: ${{ matrix.test_cases.expected_replicas }}
         run: make test-e2e

cmd/lingo/main.go

@@ -182,7 +182,11 @@ func run() error {
 	}()
 	go func() {
 		setupLog.Info("Starting leader election")
-		le.Start(ctx)
+		err := le.Start(ctx)
+		if err != nil {
+			setupLog.Error(err, "starting leader election")
+			os.Exit(1)
+		}
 	}()
 	defer func() {
 		setupLog.Info("waiting on manager to stop")

pkg/deployments/scaler.go

@@ -79,7 +79,7 @@ func (s *scaler) compareScales(current, desired int32) {
 		s.desiredScale = desired
 	}
 
-	if s.currentScale == -1 || s.desiredScale == -1 {
+	if s.currentScale == -1 || s.desiredScale == -1 || desired == -1 {
 		// Nothing to compare if we only have partial information
 		return
 	}

pkg/leader/election.go

@@ -10,6 +10,7 @@ import (
 	"k8s.io/client-go/kubernetes"
 	"k8s.io/client-go/tools/leaderelection"
 	"k8s.io/client-go/tools/leaderelection/resourcelock"
+	"k8s.io/client-go/util/flowcontrol"
 )
 
 func NewElection(clientset kubernetes.Interface, id, namespace string) *Election {
@@ -35,11 +36,11 @@ func NewElection(clientset kubernetes.Interface, id, namespace string) *Election
 		RetryPeriod:     2 * time.Second,
 		Callbacks: leaderelection.LeaderCallbacks{
 			OnStartedLeading: func(ctx context.Context) {
-				log.Println("Started leading")
+				log.Printf("%v started leading", id)
 				isLeader.Store(true)
 			},
 			OnStoppedLeading: func() {
-				log.Println("Stopped leading")
+				log.Printf("%v stopped leading", id)
 				isLeader.Store(false)
 			},
 			OnNewLeader: func(identity string) {
@@ -54,14 +55,28 @@ func NewElection(clientset kubernetes.Interface, id, namespace string) *Election
 	return &Election{
 		IsLeader: isLeader,
 		config:   config,
+		ID:       id,
 	}
 }
 
 type Election struct {
 	config   leaderelection.LeaderElectionConfig
 	IsLeader *atomic.Bool
+	ID       string
 }
 
-func (le *Election) Start(ctx context.Context) {
-	leaderelection.RunOrDie(ctx, le.config)
+func (le *Election) Start(ctx context.Context) error {
+	backoff := flowcontrol.NewBackOff(1*time.Second, 15*time.Second)
+	const backoffID = "lingo-leader-election"
+	for {
+		leaderelection.RunOrDie(ctx, le.config)
+		backoff.Next(backoffID, backoff.Clock.Now())
+		delay := backoff.Get(backoffID)
+		log.Printf("Leader election stopped on %v, retrying in %v", le.ID, delay)
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-time.After(delay):
+		}
+	}
 }

tests/e2e/test.sh

@@ -6,6 +6,9 @@ set -xe
 HOST=127.0.0.1
 PORT=30080
 BASE_URL="http://$HOST:$PORT/v1"
+REPLICAS=${REPLICAS:-3}
+REQUESTS=60
+EXPECTED_REPLICAS=1
 
 
 if kind get clusters | grep -q substratus-test; then
@@ -42,6 +45,9 @@ if ! kubectl get deployment lingo; then
   skaffold run
 fi
 
+kubectl patch deployment lingo --patch "{\"spec\": {\"replicas\": $REPLICAS}}"
+
+kubectl logs -f deployment/lingo &
 
 kubectl wait --for=condition=available --timeout=30s deployment/lingo
 
@@ -77,53 +83,42 @@ pip3 install openai==1.2.3
 
 # Send 60 requests in parallel to stapi backend using openai python client and threading
 python3 $SCRIPT_DIR/test_openai_embedding.py \
-  --requests 60 --timeout 300 --base-url "${BASE_URL}" \
+  --requests ${REQUESTS} --timeout 300 --base-url "${BASE_URL}" \
   --model text-embedding-ada-002
 
 # Ensure replicas has been scaled up to 1 after sending 60 requests
 replicas=$(kubectl get deployment stapi-minilm-l6-v2 -o jsonpath='{.spec.replicas}')
-if [ "$replicas" -eq 1 ]; then
-  echo "Test passed: Expected 1 replica after sending requests 60 requests"
+if [ "$replicas" -ge "${EXPECTED_REPLICAS}" ]; then
+  echo "Test passed: Expected ${EXPECTED_REPLICAS} or more replicas and got ${replicas} after sending requests ${REQUESTS} requests"
   else
-  echo "Test failed: Expected 1 replica after sending requests 60 requests, got $replicas"
+  echo "Test failed: Expected ${EXPECTED_REPLICAS} or more replicas after sending requests ${REQUESTS} requests, got ${replicas}"
   exit 1
 fi
 
-echo "Waiting for deployment to scale down back to 0 within 2 minutes"
-for i in {1..15}; do
-  if [ "$i" -eq 15 ]; then
+# Verify that leader election works by forcing a 120 second apiserver outage
+KIND_NODE=$(kind get nodes --name=substratus-test)
+docker exec ${KIND_NODE} iptables -I INPUT -p tcp --dport 6443 -j DROP
+sleep 120
+docker exec ${KIND_NODE} iptables -D INPUT -p tcp --dport 6443 -j DROP
+
+until kubectl get pods; do
+  echo "Waiting for apiserver to be back up, waiting for 1 second and trying again"
+  sleep 1
+done
+
+# rerun kubectl logs because previous one got killed when apiserver was down
+kubectl logs --tail=500 -f deployment/lingo &
+
+echo "Waiting for deployment to scale down back to 0 within ~2 minutes"
+for i in {1..30}; do
+  if [ "$i" -eq 30 ]; then
     echo "Test failed: Expected 0 replica after not having requests for more than 1 minute, got $replicas"
     exit 1
   fi
-  replicas=$(kubectl get deployment stapi-minilm-l6-v2 -o jsonpath='{.spec.replicas}')
+  replicas=$(kubectl get deployment stapi-minilm-l6-v2 -o jsonpath='{.spec.replicas}' || true)
   if [ "$replicas" -eq 0 ]; then
     echo "Test passed: Expected 0 replica after not having requests for more than 1 minute"
     break
   fi
-  sleep 8
+  sleep 6
 done
-
-echo "Patching stapi deployment to sleep on startup"
-cat <<EOF | kubectl patch deployment stapi-minilm-l6-v2 --type merge --patch "$(cat)"
-spec:
-  template:
-    spec:
-      initContainers:
-      - name: sleep
-        image: busybox
-        command: ["sh", "-c", "sleep 10"]
-EOF
-
-requests=300
-echo "Send $requests requests in parallel to stapi backend using openai python client and threading"
-python3 $SCRIPT_DIR/test_openai_embedding.py \
-  --requests $requests --timeout 600 --base-url "${BASE_URL}" \
-  --model text-embedding-ada-002
-
-replicas=$(kubectl get deployment stapi-minilm-l6-v2 -o jsonpath='{.spec.replicas}')
-if [ "$replicas" -ge 2 ]; then
-  echo "Test passed: Expected 2 or more replicas after sending more than $requests requests, got $replicas"
-  else
-  echo "Test failed: Expected 2 or more replicas after sending more than $requests requests, got $replicas"
-  exit 1
-fi