Repository: kubernetes-sigs/lws (go)

Issue:
All LWS groups get updated all together during rolling update when LWS controller pod is restarted or controller re-sync happens
<!-- Please use this template while reporting a bug and provide as much info as possible. Not doing so may result in your bug not being addressed in a timely manner. Thanks!

If the matter is security related, please disclose it privately via https://kubernetes.io/security/
-->

**What happened**:
When deploying changes that update the leader and worker templates, the rolling update will be triggered. During the rolling update, I observed 2 scenarios where LWS controller will reconcile all worker statefulsets and trigger "patch" event to update all worker statefulsets to use the new version at the same time. This causes service downtime since all model replicas are updated at the same time. Here are the 2 scenarios observed:

1. This can happen when the lws controller pod is restarted during rolling update. This happens when I need to update the LWS template and patch the underlying nodes at the same time. The rolling update is triggered for both LWS and the underlying instances/nodes. When the node running the lws controller is terminated, a new lws controller pod will be set up. We observed that the lws controller pod restart will trigger reconciliations. The pod controller will then update all worker statefulsets all at once.

2. Seems like the lws controller is using the default 10 hours controller SyncPeriod (https://github.com/kubernetes-sigs/controller-runtime/blob/v0.17.2/pkg/cache/cache.go#L171). For every certain period, the pod controller will reconcile the worker statefulset. If this happens during the LWS rolling update where worker template is updated, it will update all worker statefulset all at once and cause service downtime.   

**What you expected to happen**:

I would expect the worker statefulset update follows the same order as the leader pods rolling update. For controller pod restart and controller resync, they should not trigger worker statefulset update.

The LWS rolling update is fully controlled by the leader statefulset's UpdateStrategy. For worker statefulset, currently it is not configured explicitly and uses the default RollingUpdate strategy. I switched to use OnDelete strategy for worker statefulset and it bypasses the issue. This works since the LWS is configured to use ```RecreateGroupOnPodRestart``` RestartPolicy, during leader pods rolling update, all worker pods will be restarted which will trigger worker StatefulSet update when using ```OnDelete``` UpdateStrategy. In this case, worker StatefulSet update will follow the leader pod rolling update order and ignore the pod controller reconciliation from the above 2 scenarios. 

**How to reproduce it (as minimally and precisely as possible)**:
* Set up a running LWS
* Update the LWS definition including changes to both leader and worker templates. Say changing the image URL or environment variables passed for both leader and worker templates
* Run ```kubectl delete pod <lws controller pod> -n lws-system``` to kill the lws controller pod to force a controller pod restart (scenario 1)
* After a few minutes, all worker statefulsets will get updated instead of following rolling update.

**Anything else we need to know?**: 
Tried with the following change for pod controller (https://github.com/kubernetes-sigs/lws/blob/main/pkg/controllers/pod_controller.go#L307-L316) to bypass the issue.
```
modified   pkg/controllers/pod_controller.go
@@ -311,6 +311,7 @@ func constructWorkerStatefulSetApplyConfiguration(leaderPod corev1.Pod, lws lead                                                                        
       WithPodManagementPolicy(appsv1.ParallelPodManagement).                                                                                                               
       WithTemplate(&podTemplateApplyConfiguration).                                                                                                                        
       WithOrdinals(appsapplyv1.StatefulSetOrdinals().WithStart(1)).                                                                                                        
+      WithUpdateStrategy(appsapplyv1.StatefulSetUpdateStrategy().WithType(appsv1.OnDeleteStatefulSetStrategyType)).                                                        
       WithSelector(metaapplyv1.LabelSelector().                                                                                                                            
         WithMatchLabels(selectorMap))).                                                                                                                                    
     WithLabels(labelMap)
```

**Environment**: I set up the LWS using AWS EKS

- Kubernetes version (use `kubectl version`): 1.29
- LWS version (use `git describe --tags --dirty --always`): 0.3.0
- Cloud provider or hardware configuration: AWS EKS
- OS (e.g: `cat /etc/os-release`): Amazon Linux
- Kernel (e.g. `uname -a`): Linux ip-21-16-69-244.ec2.internal 5.10.224-212.876.amzn2.x86_64 #1 SMP Thu Aug 22 16:55:24 UTC 2024 x86_64 x86_64 x86_64 GNU/Linux
- Install tools:
- Others:


Pull request description:
Creating new workerStatefulSet instead of patching when updating LWS template
<!--  Thanks for sending a pull request!  Here are some tips for you:

1. If this is your first time, please read our contributor guidelines: https://git.k8s.io/community/contributors/guide/first-contribution.md#your-first-contribution and developer guide https://git.k8s.io/community/contributors/devel/development.md#development-guide
2. Please label this pull request according to what type of issue you are addressing, especially if this is a release targeted pull request. For reference on required PR/issue labels, read here:
https://git.k8s.io/community/contributors/devel/sig-release/release.md#issuepr-kind-label
3. Ensure you have added or ran the appropriate tests for your PR: https://git.k8s.io/community/contributors/devel/sig-testing/testing.md
4. If you want *faster* PR reviews, read how: https://git.k8s.io/community/contributors/guide/pull-requests.md#best-practices-for-faster-reviews
5. If the PR is unfinished, see how to mark it: https://git.k8s.io/community/contributors/guide/pull-requests.md#marking-unfinished-pull-requests
-->

#### What type of PR is this?
/kind bug 
<!--
Add one of the following kinds:
/kind bug
/kind cleanup
/kind documentation
/kind feature

Optionally add one or more of the following kinds if applicable:
/kind api-change
/kind deprecation
/kind failing-test
/kind flake
/kind regression
-->

#### What this PR does / why we need it
Changes from patching the worker sts to update it, to create a new worker sts when it doesn't exist. The need for this change is described in https://github.com/kubernetes-sigs/lws/issues/225#issuecomment-2395473982 

#### Which issue(s) this PR fixes
<!--
*Automatically closes linked issue when PR is merged.
Usage: `Fixes #<issue number>`, or `Fixes (paste link of issue)`.
_If PR is about `failing-tests or flakes`, please post the related issues/tests in a comment and do not use `Fixes`_*
-->
Fixes #225 

#### Special notes for your reviewer

#### Does this PR introduce a user-facing change?
<!--
If no, just write "NONE" in the release-note block below.
If yes, a release note is required:
Enter your extended release note in the block below. If the PR requires additional action from users switching to the new release, include the string "action required".
-->
```release-note

```


Patch:
diff --git a/pkg/controllers/pod_controller.go b/pkg/controllers/pod_controller.go
index 1a7c6ad..e404522 100644
--- a/pkg/controllers/pod_controller.go
+++ b/pkg/controllers/pod_controller.go
@@ -32,7 +32,6 @@ import (
 	coreapplyv1 "k8s.io/client-go/applyconfigurations/core/v1"
 	metaapplyv1 "k8s.io/client-go/applyconfigurations/meta/v1"
 	"k8s.io/klog/v2"
-	"k8s.io/utils/ptr"
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/client/apiutil"
@@ -146,19 +145,18 @@ func (r *PodReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.R
 	if err != nil {
 		return ctrl.Result{}, err
 	}
-	patch := &unstructured.Unstructured{
+	workerStatefulSet := &unstructured.Unstructured{
 		Object: obj,
 	}
-	// Use server side apply and add fieldmanagaer to the lws owned fields
-	// If there are conflicts in the fields owned by the lws controller, lws will obtain the ownership and force override
-	// these fields to the ones desired by the lws controller. These fields are specified in the StatefulSetApplyConfiguration
-	// TODO b/316776287 add E2E test for SSA
-	err = r.Patch(ctx, patch, client.Apply, &client.PatchOptions{
-		FieldManager: fieldManager,
-		Force:        ptr.To[bool](true),
-	})
-	if err != nil {
-		return ctrl.Result{}, err
+
+	var workerSts appsv1.StatefulSet
+	if err := r.Get(ctx, types.NamespacedName{Name: pod.Name, Namespace: leaderWorkerSet.Namespace}, &workerSts); err != nil {
+		if client.IgnoreNotFound(err) != nil {
+			return ctrl.Result{}, err
+		}
+		if err = r.Create(ctx, workerStatefulSet); err != nil {
+			return ctrl.Result{}, err
+		}
 	}
 	log.V(2).Info("Worker Reconcile completed.")
 	return ctrl.Result{}, nil
diff --git a/pkg/utils/accelerators/tpu.go b/pkg/utils/accelerators/tpu.go
index 235b3cd..de4e8f5 100644
--- a/pkg/utils/accelerators/tpu.go
+++ b/pkg/utils/accelerators/tpu.go
@@ -32,6 +32,7 @@ const (
 	TpuResourceName                 corev1.ResourceName = corev1.ResourceName("google.com/tpu")
 	TpuWorkerHostNames              string              = "TPU_WORKER_HOSTNAMES"
 	TpuWorkerId                     string              = "TPU_WORKER_ID"
+	TpuName                         string              = "TPU_NAME"
 	LeaderRequestsTPUsAnnotationKey string              = "leaderworkerset.sigs.k8s.io/leader-requests-tpus"
 )
 
@@ -154,6 +155,10 @@ func addTPUVariablesSubGroup(pod *corev1.Pod) error {
 			Name:  TpuWorkerId,
 			Value: fmt.Sprint(tpuWorkerId),
 		},
+		corev1.EnvVar{
+			Name:  TpuName,
+			Value: fmt.Sprint(leaderName),
+		},
 	)
 	return nil
 
@@ -211,6 +216,10 @@ func AddTPUVariables(pod *corev1.Pod, size int) error {
 			Name:  TpuWorkerId,
 			Value: fmt.Sprint(tpuWorkerId),
 		},
+		corev1.EnvVar{
+			Name:  TpuName,
+			Value: fmt.Sprint(leaderName),
+		},
 	)
 	return nil
 }


Test patch:
diff --git a/pkg/utils/accelerators/tpu_test.go b/pkg/utils/accelerators/tpu_test.go
index e5524cf..c18d501 100644
--- a/pkg/utils/accelerators/tpu_test.go
+++ b/pkg/utils/accelerators/tpu_test.go
@@ -35,6 +35,7 @@ func TestAddTPUVariables(t *testing.T) {
 		hasWorkerIndexLabelKey     bool
 		expectedTpuWorkerHostNames string
 		expectedTpuWorkerId        string
+		expectedTpuName            string
 	}{
 		{
 			name: "Worker Index is 0",
@@ -52,6 +53,7 @@ func TestAddTPUVariables(t *testing.T) {
 			hasWorkerIndexLabelKey:     true,
 			expectedTpuWorkerHostNames: "test-sample-1.default",
 			expectedTpuWorkerId:        "0",
+			expectedTpuName:            "test-sample-1",
 		},
 		{
 			name: "Worker Index is non-zero, size is above 2",
@@ -72,6 +74,7 @@ func TestAddTPUVariables(t *testing.T) {
 			hasWorkerIndexLabelKey:     true,
 			expectedTpuWorkerHostNames: "test-sample-1.default,test-sample-1-1.default,test-sample-1-2.default,test-sample-1-3.default,test-sample-1-4.default",
 			expectedTpuWorkerId:        "3",
+			expectedTpuName:            "test-sample-1",
 		},
 	}
 
@@ -95,6 +98,9 @@ func TestAddTPUVariables(t *testing.T) {
 				if diff := cmp.Diff(tc.pod.Spec.Containers[0].Env[1].Value, tc.expectedTpuWorkerId); diff != "" {
 					t.Errorf("unexpected add TPU worker ID operation: %s", diff)
 				}
+				if diff := cmp.Diff(tc.pod.Spec.Containers[0].Env[2].Value, tc.expectedTpuName); diff != "" {
+					t.Errorf("unexpected add TPU Name operation: %s", diff)
+				}
 			}
 		})
 	}
@@ -106,6 +112,7 @@ func TestAddTPUVariablesSubGroup(t *testing.T) {
 		pod                        *corev1.Pod
 		expectedTpuWorkerHostNames string
 		expectedTpuWorkerId        string
+		expectedTpuName            string
 	}{
 		{
 			name: "Leader requests TPU resources",
@@ -126,6 +133,7 @@ func TestAddTPUVariablesSubGroup(t *testing.T) {
 			},
 			expectedTpuWorkerId:        "3",
 			expectedTpuWorkerHostNames: "test-sample-1.default,test-sample-1-1.default,test-sample-1-2.default,test-sample-1-3.default,test-sample-1-4.default",
+			expectedTpuName:            "test-sample-1",
 		},
 		{
 			name: "Leader requests TPU resources, worker with subgroup index > 0",
@@ -146,6 +154,7 @@ func TestAddTPUVariablesSubGroup(t *testing.T) {
 			},
 			expectedTpuWorkerId:        "3",
 			expectedTpuWorkerHostNames: "test-sample-1-4.default,test-sample-1-5.default,test-sample-1-6.default,test-sample-1-7.default",
+			expectedTpuName:            "test-sample-1",
 		},
 		{
 			name: "Leader does not request TPU resources, worker with subgroup index > 0",
@@ -165,6 +174,7 @@ func TestAddTPUVariablesSubGroup(t *testing.T) {
 			},
 			expectedTpuWorkerId:        "0",
 			expectedTpuWorkerHostNames: "test-sample-1-5.default,test-sample-1-6.default,test-sample-1-7.default,test-sample-1-8.default",
+			expectedTpuName:            "test-sample-1",
 		},
 	}
 	for _, tc := range tests {
@@ -179,6 +189,9 @@ func TestAddTPUVariablesSubGroup(t *testing.T) {
 			if diff := cmp.Diff(tc.pod.Spec.Containers[0].Env[1].Value, tc.expectedTpuWorkerId); diff != "" {
 				t.Errorf("unexpected add TPU worker ID operation: %s", diff)
 			}
+			if diff := cmp.Diff(tc.pod.Spec.Containers[0].Env[2].Value, tc.expectedTpuName); diff != "" {
+				t.Errorf("unexpected add TPU Name operation: %s", diff)
+			}
 		})
 	}
 }
diff --git a/test/integration/controllers/leaderworkerset_test.go b/test/integration/controllers/leaderworkerset_test.go
index 186b6c0..d063db5 100644
--- a/test/integration/controllers/leaderworkerset_test.go
+++ b/test/integration/controllers/leaderworkerset_test.go
@@ -584,7 +584,7 @@ var _ = ginkgo.Describe("LeaderWorkerSet controller", func() {
 				},
 			},
 		}),
-		ginkgo.Entry("leaderTemplate changed with maxUnavailable=2", &testCase{
+		ginkgo.Entry("workerTemplate changed with maxUnavailable=2", &testCase{
 			makeLeaderWorkerSet: func(nsName string) *testing.LeaderWorkerSetWrapper {
 				return testing.BuildLeaderWorkerSet(nsName).Replica(4).MaxUnavailable(2)
 			},
@@ -662,7 +662,7 @@ var _ = ginkgo.Describe("LeaderWorkerSet controller", func() {
 				},
 			},
 		}),
-		ginkgo.Entry("leaderTemplate changed with maxUnavailable greater than replicas", &testCase{
+		ginkgo.Entry("workerTemplate changed with maxUnavailable greater than replicas", &testCase{
 			makeLeaderWorkerSet: func(nsName string) *testing.LeaderWorkerSetWrapper {
 				return testing.BuildLeaderWorkerSet(nsName).Replica(4).MaxUnavailable(10)
 			},
diff --git a/test/testutils/util.go b/test/testutils/util.go
index 45f72f7..c2ec3b4 100644
--- a/test/testutils/util.go
+++ b/test/testutils/util.go
@@ -267,6 +267,7 @@ func SetLeaderPodToReady(ctx context.Context, k8sClient client.Client, podName s
 			Status: corev1.ConditionTrue,
 		}
 		leaderPod.Status.Conditions = append(leaderPod.Status.Conditions, condition)
+		deleteWorkerStatefulSetIfExists(ctx, k8sClient, podName, lws)
 		return k8sClient.Status().Update(ctx, &leaderPod)
 	}, Timeout, Interval).Should(gomega.Succeed())
 }
@@ -367,7 +368,7 @@ func IsContainerFirstEnvVarLWSLeaderAddress(pod corev1.Pod) error {
 }
 
 func HasTPUEnvVarsPopulated(pod corev1.Pod) bool {
-	return hasAllEnvVarPopulated(pod, []string{acceleratorutils.TpuWorkerHostNames, acceleratorutils.TpuWorkerId})
+	return hasAllEnvVarPopulated(pod, []string{acceleratorutils.TpuWorkerHostNames, acceleratorutils.TpuWorkerId, acceleratorutils.TpuName})
 }
 
 func CheckTPUContainerHasCorrectEnvVars(pod corev1.Pod, envVal string) error {
@@ -538,3 +539,17 @@ func SetLeaderPodsToReady(ctx context.Context, k8sClient client.Client, lws *lea
 		return k8sClient.Status().Update(ctx, &sts)
 	}, Timeout, Interval).Should(gomega.Succeed())
 }
+
+func deleteWorkerStatefulSetIfExists(ctx context.Context, k8sClient client.Client, statefulsetName string, lws *leaderworkerset.LeaderWorkerSet) {
+	// in cases where size = 1, the workerstatefulset does not exist
+	gomega.Eventually(func() error {
+		var sts appsv1.StatefulSet
+		if err := k8sClient.Get(ctx, types.NamespacedName{Name: statefulsetName, Namespace: lws.Namespace}, &sts); err != nil {
+			if client.IgnoreNotFound(err) != nil {
+				return err
+			}
+			return nil
+		}
+		return k8sClient.Delete(ctx, &sts)
+	}, Timeout, Interval).Should(gomega.Succeed())
+}