kserve/pkg/controller/v1alpha2/llmisvc/controller_int_multi_node_test.go

/*
Copyright 2025 The KServe Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package llmisvc_test

import (
	"context"

	"k8s.io/client-go/util/retry"
	ctrl "sigs.k8s.io/controller-runtime"
	lwsapi "sigs.k8s.io/lws/api/leaderworkerset/v1"

	. "github.com/onsi/ginkgo/v2"
	. "github.com/onsi/gomega"
	corev1 "k8s.io/api/core/v1"
	rbacv1 "k8s.io/api/rbac/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/utils/ptr"
	"knative.dev/pkg/kmeta"

	"github.com/kserve/kserve/pkg/apis/serving/v1alpha2"
	"github.com/kserve/kserve/pkg/constants"

	. "github.com/kserve/kserve/pkg/controller/v1alpha2/llmisvc/fixture"
	. "github.com/kserve/kserve/pkg/testing"
)

var _ = Describe("LLMInferenceService Multi-Node Controller", func() {
	Context("Multi-Node Workload Reconciliation", func() {
		It("should create a basic multi-node deployment with worker spec", func(ctx SpecContext) {
			// given
			svcName := "test-llm-multinode"
			testNs := NewTestNamespace(ctx, envTest, WithIstioShadowService(svcName))

			llmSvc := LLMInferenceService(svcName,
				InNamespace[*v1alpha2.LLMInferenceService](testNs.Name),
				WithModelURI("hf://facebook/opt-125m"),
				WithReplicas(2),
				WithParallelism(ParallelismSpec(
					WithDataParallelism(4),
					WithDataLocalParallelism(1),
					WithTensorParallelism(3),
				)),
				WithTemplate(SimpleWorkerPodSpec()),
				WithWorker(SimpleWorkerPodSpec()),
				WithPrefill(SimpleWorkerPodSpec()),
				WithManagedRoute(),
				WithManagedGateway(),
			)

			// when
			Expect(envTest.Create(ctx, llmSvc)).To(Succeed())
			defer func() {
				testNs.DeleteAndWait(ctx, llmSvc)
			}()

			// then
			expectedLWS := &lwsapi.LeaderWorkerSet{}
			Eventually(func(g Gomega, ctx context.Context) error {
				return envTest.Get(ctx, types.NamespacedName{
					Name:      svcName + "-kserve-mn",
					Namespace: testNs.Name,
				}, expectedLWS)
			}).WithContext(ctx).Should(Succeed())

			Expect(expectedLWS.Spec.Replicas).To(Equal(ptr.To[int32](2)))
			Expect(expectedLWS.Spec.LeaderWorkerTemplate.Size).To(Equal(ptr.To[int32](4)))
			Expect(expectedLWS).To(BeOwnedBy(llmSvc))

			// Verify leader template is set
			Expect(expectedLWS.Spec.LeaderWorkerTemplate.LeaderTemplate).ToNot(BeNil())
			Expect(expectedLWS.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers).To(HaveLen(1))
			Expect(expectedLWS.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers[0].Name).To(Equal("main"))

			// Verify worker template is set
			Expect(expectedLWS.Spec.LeaderWorkerTemplate.WorkerTemplate.Spec.Containers).To(HaveLen(1))
			Expect(expectedLWS.Spec.LeaderWorkerTemplate.WorkerTemplate.Spec.Containers[0].Name).To(Equal("main"))

			// Verify labels
			Expect(expectedLWS.Spec.LeaderWorkerTemplate.LeaderTemplate.Labels).To(HaveKeyWithValue(constants.KServeComponentLabelKey, constants.KServeComponentWorkload))
			Expect(expectedLWS.Spec.LeaderWorkerTemplate.LeaderTemplate.Labels).To(HaveKeyWithValue(constants.LLMDRoleLabelKey, constants.LLMDRoleDecode))
		})

		It("should create multi-node deployment with prefill workload", func(ctx SpecContext) {
			// given
			svcName := "test-llm-multinode-prefill"
			testNs := NewTestNamespace(ctx, envTest, WithIstioShadowService(svcName))

			llmSvc := LLMInferenceService(svcName,
				InNamespace[*v1alpha2.LLMInferenceService](testNs.Name),
				WithModelURI("hf://facebook/opt-125m"),
				WithReplicas(1),
				WithParallelism(ParallelismSpec(
					WithDataParallelism(10),
					WithDataLocalParallelism(2),
					WithTensorParallelism(4),
				)),
				WithWorker(SimpleWorkerPodSpec()),
				WithPrefillParallelism(ParallelismSpec(
					WithDataParallelism(3),
					WithDataLocalParallelism(1),
					WithTensorParallelism(4),
				)),
				WithPrefillWorker(SimpleWorkerPodSpec()),
				WithPrefillReplicas(1),
				WithManagedRoute(),
				WithManagedGateway(),
			)

			// when
			Expect(envTest.Create(ctx, llmSvc)).To(Succeed())
			defer func() {
				testNs.DeleteAndWait(ctx, llmSvc)
			}()

			// then - Check main workload LWS
			expectedMainLWS := &lwsapi.LeaderWorkerSet{}
			Eventually(func(g Gomega, ctx context.Context) error {
				return envTest.Get(ctx, types.NamespacedName{
					Name:      svcName + "-kserve-mn",
					Namespace: testNs.Name,
				}, expectedMainLWS)
			}).WithContext(ctx).Should(Succeed())

			Expect(expectedMainLWS.Spec.Replicas).To(Equal(ptr.To[int32](1)))
			Expect(expectedMainLWS.Spec.LeaderWorkerTemplate.Size).To(Equal(ptr.To[int32](5)))

			// then - Check prefill workload LWS
			expectedPrefillLWS := &lwsapi.LeaderWorkerSet{}
			Eventually(func(g Gomega, ctx context.Context) error {
				return envTest.Get(ctx, types.NamespacedName{
					Name:      svcName + "-kserve-mn-prefill",
					Namespace: testNs.Name,
				}, expectedPrefillLWS)
			}).WithContext(ctx).Should(Succeed())

			Expect(expectedPrefillLWS.Spec.Replicas).To(Equal(ptr.To[int32](1)))
			Expect(expectedPrefillLWS.Spec.LeaderWorkerTemplate.Size).To(Equal(ptr.To[int32](3)))
			Expect(expectedPrefillLWS).To(BeOwnedBy(llmSvc))

			// Verify prefill-specific labels
			Expect(expectedPrefillLWS.Spec.LeaderWorkerTemplate.LeaderTemplate.Labels).To(HaveKeyWithValue(constants.LLMDRoleLabelKey, constants.LLMDRolePrefill))
		})

		It("should create RBAC resources when prefill and decode is used", func(ctx SpecContext) {
			// given
			svcName := "test-llm-multinode-rbac"
			testNs := NewTestNamespace(ctx, envTest, WithIstioShadowService(svcName))

			llmSvc := LLMInferenceService(svcName,
				InNamespace[*v1alpha2.LLMInferenceService](testNs.Name),
				WithModelURI("hf://facebook/opt-125m"),
				WithReplicas(1),
				WithParallelism(ParallelismSpec(
					WithDataParallelism(2),
					WithDataLocalParallelism(1),
					WithTensorParallelism(4),
				)),
				WithWorker(&corev1.PodSpec{}),
				WithManagedRoute(),
				WithManagedScheduler(),
				WithManagedGateway(),
				WithPrefill(SimpleWorkerPodSpec()),
			)

			// when
			Expect(envTest.Create(ctx, llmSvc)).To(Succeed())
			defer func() {
				testNs.DeleteAndWait(ctx, llmSvc)
			}()

			// then - Check ServiceAccount is created
			expectedSA := &corev1.ServiceAccount{}
			Eventually(func(g Gomega, ctx context.Context) error {
				return envTest.Get(ctx, types.NamespacedName{
					Name:      svcName + "-kserve-mn",
					Namespace: testNs.Name,
				}, expectedSA)
			}).WithContext(ctx).Should(Succeed())

			Expect(expectedSA).To(BeOwnedBy(llmSvc))
			Expect(expectedSA.Labels).To(HaveKeyWithValue(constants.KubernetesAppNameLabelKey, svcName))

			// then - Check Role is created
			expectedRole := &rbacv1.Role{}
			Eventually(func(g Gomega, ctx context.Context) error {
				return envTest.Get(ctx, types.NamespacedName{
					Name:      svcName + "-kserve-mn-role",
					Namespace: testNs.Name,
				}, expectedRole)
			}).WithContext(ctx).Should(Succeed())

			Expect(expectedRole).To(BeOwnedBy(llmSvc))
			Expect(expectedRole.Rules).ToNot(BeEmpty())

			// then - Check RoleBinding is created
			expectedRB := &rbacv1.RoleBinding{}
			Eventually(func(g Gomega, ctx context.Context) error {
				return envTest.Get(ctx, types.NamespacedName{
					Name:      svcName + "-kserve-mn-rb",
					Namespace: testNs.Name,
				}, expectedRB)
			}).WithContext(ctx).Should(Succeed())

			Expect(expectedRB).To(BeOwnedBy(llmSvc))
			Expect(expectedRB.Subjects).To(HaveLen(1))
			Expect(expectedRB.Subjects[0].Name).To(Equal(expectedSA.Name))
			Expect(expectedRB.RoleRef.Name).To(Equal(expectedRole.Name))

			// then - Check LWS uses the ServiceAccount
			expectedLWS := &lwsapi.LeaderWorkerSet{}
			Eventually(func(g Gomega, ctx context.Context) error {
				return envTest.Get(ctx, types.NamespacedName{
					Name:      svcName + "-kserve-mn",
					Namespace: testNs.Name,
				}, expectedLWS)
			}).WithContext(ctx).Should(Succeed())

			Expect(expectedLWS).To(BeOwnedBy(llmSvc))
			Expect(expectedLWS.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.ServiceAccountName).To(Equal(expectedSA.Name))
			Expect(expectedLWS.Spec.LeaderWorkerTemplate.WorkerTemplate.Spec.ServiceAccountName).To(Equal(expectedSA.Name))
		})

		It("should propagate imagePullSecrets from default SA to the created multi-node SA", func(ctx SpecContext) {
			// given
			svcName := "test-llm-mn-ips"
			testNs := NewTestNamespace(ctx, envTest,
				WithIstioShadowService(svcName),
				WithDefaultServiceAccountImagePullSecrets(
					corev1.LocalObjectReference{Name: "my-registry-secret"},
					corev1.LocalObjectReference{Name: "other-pull-secret"},
				),
			)

			llmSvc := LLMInferenceService(svcName,
				InNamespace[*v1alpha2.LLMInferenceService](testNs.Name),
				WithModelURI("hf://facebook/opt-125m"),
				WithReplicas(1),
				WithParallelism(ParallelismSpec(
					WithDataParallelism(2),
					WithDataLocalParallelism(1),
					WithTensorParallelism(4),
				)),
				WithWorker(&corev1.PodSpec{}),
				WithManagedRoute(),
				WithManagedScheduler(),
				WithManagedGateway(),
				WithPrefill(SimpleWorkerPodSpec()),
			)

			// when
			Expect(envTest.Create(ctx, llmSvc)).To(Succeed())
			defer func() {
				testNs.DeleteAndWait(ctx, llmSvc)
			}()

			// then - created SA should have the imagePullSecrets from the default SA
			expectedSA := &corev1.ServiceAccount{}
			Eventually(func(g Gomega, ctx context.Context) error {
				return envTest.Get(ctx, types.NamespacedName{
					Name:      svcName + "-kserve-mn",
					Namespace: testNs.Name,
				}, expectedSA)
			}).WithContext(ctx).Should(Succeed())

			Expect(expectedSA).To(BeOwnedBy(llmSvc))
			Expect(expectedSA.ImagePullSecrets).To(ConsistOf(
				corev1.LocalObjectReference{Name: "my-registry-secret"},
				corev1.LocalObjectReference{Name: "other-pull-secret"},
			))
		})

		It("should propagate imagePullSecrets from default SA to the created multi-node prefill SA", func(ctx SpecContext) {
			// given
			svcName := "test-llm-mn-ips-prefill"
			testNs := NewTestNamespace(ctx, envTest,
				WithIstioShadowService(svcName),
				WithDefaultServiceAccountImagePullSecrets(
					corev1.LocalObjectReference{Name: "my-registry-secret"},
				),
			)

			llmSvc := LLMInferenceService(svcName,
				InNamespace[*v1alpha2.LLMInferenceService](testNs.Name),
				WithModelURI("hf://facebook/opt-125m"),
				WithReplicas(1),
				WithParallelism(ParallelismSpec(
					WithDataParallelism(2),
					WithDataLocalParallelism(1),
				)),
				WithWorker(&corev1.PodSpec{}),
				WithPrefillParallelism(ParallelismSpec(
					WithDataParallelism(2),
					WithDataLocalParallelism(1),
				)),
				WithPrefillWorker(&corev1.PodSpec{}),
				WithPrefillReplicas(1),
				WithManagedRoute(),
				WithManagedGateway(),
			)

			// when
			Expect(envTest.Create(ctx, llmSvc)).To(Succeed())
			defer func() {
				testNs.DeleteAndWait(ctx, llmSvc)
			}()

			// then - prefill SA should have the imagePullSecrets from the default SA
			prefillSA := &corev1.ServiceAccount{}
			Eventually(func(g Gomega, ctx context.Context) error {
				return envTest.Get(ctx, types.NamespacedName{
					Name:      kmeta.ChildName(svcName, "-kserve-mn-prefill"),
					Namespace: testNs.Name,
				}, prefillSA)
			}).WithContext(ctx).Should(Succeed())

			Expect(prefillSA).To(BeOwnedBy(llmSvc))
			Expect(prefillSA.ImagePullSecrets).To(ConsistOf(
				corev1.LocalObjectReference{Name: "my-registry-secret"},
			))
		})

		It("should create multi-node SA with empty imagePullSecrets when default SA has none", func(ctx SpecContext) {
			// given
			svcName := "test-llm-mn-ips-empty"
			testNs := NewTestNamespace(ctx, envTest, WithIstioShadowService(svcName))

			llmSvc := LLMInferenceService(svcName,
				InNamespace[*v1alpha2.LLMInferenceService](testNs.Name),
				WithModelURI("hf://facebook/opt-125m"),
				WithReplicas(1),
				WithParallelism(ParallelismSpec(
					WithDataParallelism(2),
					WithDataLocalParallelism(1),
				)),
				WithWorker(&corev1.PodSpec{}),
				WithManagedRoute(),
				WithManagedScheduler(),
				WithManagedGateway(),
				WithPrefill(SimpleWorkerPodSpec()),
			)

			// when
			Expect(envTest.Create(ctx, llmSvc)).To(Succeed())
			defer func() {
				testNs.DeleteAndWait(ctx, llmSvc)
			}()

			// then
			expectedSA := &corev1.ServiceAccount{}
			Eventually(func(g Gomega, ctx context.Context) error {
				return envTest.Get(ctx, types.NamespacedName{
					Name:      svcName + "-kserve-mn",
					Namespace: testNs.Name,
				}, expectedSA)
			}).WithContext(ctx).Should(Succeed())

			Expect(expectedSA).To(BeOwnedBy(llmSvc))
			Expect(expectedSA.ImagePullSecrets).To(BeEmpty())
		})

		It("should delete multi-node resources when worker spec is removed", func(ctx SpecContext) {
			// given
			svcName := "test-llm-multinode-cleanup"
			testNs := NewTestNamespace(ctx, envTest, WithIstioShadowService(svcName))

			parallelismSpec := ParallelismSpec(
				WithDataParallelism(2),
				WithDataLocalParallelism(1),
			)
			parallelismSpec.Expert = true

			llmSvc := LLMInferenceService(svcName,
				InNamespace[*v1alpha2.LLMInferenceService](testNs.Name),
				WithModelURI("hf://facebook/opt-125m"),
				WithParallelism(parallelismSpec),
				WithWorker(&corev1.PodSpec{}),
				WithManagedRoute(),
				WithManagedGateway(),
			)

			Expect(envTest.Create(ctx, llmSvc)).To(Succeed())
			defer func() {
				testNs.DeleteAndWait(ctx, llmSvc)
			}()

			lwsName := svcName + "-kserve-mn"

			// Verify LWS is created
			Eventually(func(g Gomega, ctx context.Context) error {
				lws := &lwsapi.LeaderWorkerSet{}
				return envTest.Get(ctx, types.NamespacedName{
					Name:      lwsName,
					Namespace: testNs.Name,
				}, lws)
			}).WithContext(ctx).Should(Succeed())

			// when - Remove worker spec
			errRetry := retry.RetryOnConflict(retry.DefaultRetry, func() error {
				_, errUpdate := ctrl.CreateOrUpdate(ctx, envTest.Client, llmSvc, func() error {
					llmSvc.Spec.Worker = nil
					return nil
				})
				return errUpdate
			})
			Expect(errRetry).ToNot(HaveOccurred())

			// then - LWS should be deleted
			Eventually(func(g Gomega, ctx context.Context) error {
				lws := &lwsapi.LeaderWorkerSet{}
				err := envTest.Get(ctx, types.NamespacedName{
					Name:      lwsName,
					Namespace: testNs.Name,
				}, lws)
				g.Expect(err).To(HaveOccurred())
				return nil
			}).WithContext(ctx).Should(Succeed())
		})

		It("should delete prefill resources when prefill spec is removed", func(ctx SpecContext) {
			// given
			svcName := "test-llm-prefill-cleanup"
			testNs := NewTestNamespace(ctx, envTest, WithIstioShadowService(svcName))

			llmSvc := LLMInferenceService(svcName,
				InNamespace[*v1alpha2.LLMInferenceService](testNs.Name),
				WithModelURI("hf://facebook/opt-125m"),
				WithPrefillParallelism(ParallelismSpec(
					WithDataParallelism(2),
					WithDataLocalParallelism(1),
				)),
				WithPrefillWorker(&corev1.PodSpec{}),
				WithManagedRoute(),
				WithManagedGateway(),
			)

			Expect(envTest.Create(ctx, llmSvc)).To(Succeed())
			defer func() {
				testNs.DeleteAndWait(ctx, llmSvc)
			}()

			prefillLWSName := kmeta.ChildName(svcName, "-kserve-mn-prefill")

			// Verify prefill LWS is created
			Eventually(func(g Gomega, ctx context.Context) error {
				lws := &lwsapi.LeaderWorkerSet{}
				return envTest.Get(ctx, types.NamespacedName{
					Name:      prefillLWSName,
					Namespace: testNs.Name,
				}, lws)
			}).WithContext(ctx).Should(Succeed())

			// when - Remove prefill spec
			errRetry := retry.RetryOnConflict(retry.DefaultRetry, func() error {
				_, errUpdate := ctrl.CreateOrUpdate(ctx, envTest.Client, llmSvc, func() error {
					llmSvc.Spec.Prefill = nil
					return nil
				})
				return errUpdate
			})
			Expect(errRetry).ToNot(HaveOccurred())

			// then - Prefill LWS should be deleted
			Eventually(func(g Gomega, ctx context.Context) error {
				lws := &lwsapi.LeaderWorkerSet{}
				err := envTest.Get(ctx, types.NamespacedName{
					Name:      prefillLWSName,
					Namespace: testNs.Name,
				}, lws)
				g.Expect(err).To(HaveOccurred())
				return nil
			}).WithContext(ctx).Should(Succeed())
		})
	})

	Context("Multi-Node Label and Annotation Management", func() {
		It("should set correct labels and annotation", func(ctx SpecContext) {
			// given
			svcName := "test-llm-lws-labels"
			testNs := NewTestNamespace(ctx, envTest, WithIstioShadowService(svcName))

			localQueueName := "test-local-q"
			preemptPriority := "0"
			testValue := "test"

			llmSvc := LLMInferenceService(svcName,
				InNamespace[*v1alpha2.LLMInferenceService](testNs.Name),
				WithModelURI("hf://facebook/opt-125m"),
				WithParallelism(ParallelismSpec(
					WithDataParallelism(1),
					WithDataLocalParallelism(1),
				)),
				WithWorker(&corev1.PodSpec{}),
				WithManagedRoute(),
				WithManagedGateway(),
				// Add a kueue label and annotation to ensure value propagation to the LWS
				// the kueue functionality itself will not be tested here
				WithAnnotations(map[string]string{
					PreemptionReclaimAnnotationKey: preemptPriority,
					testValue:                      testValue, // dummy value, should not be propagated
				}),
				WithLabels(map[string]string{
					LocalQueueNameLabelKey: localQueueName,
					testValue:              testValue, // dummy value, should not be propagated
				}),
			)

			// safety check
			Expect(llmSvc.Spec.Parallelism.IsDataParallel()).To(BeTrue())
			Expect(llmSvc.Spec.Worker).To(Not(BeNil()))

			// when
			Expect(envTest.Create(ctx, llmSvc)).To(Succeed())
			defer func() {
				testNs.DeleteAndWait(ctx, llmSvc)
			}()

			// then
			expectedLWS := &lwsapi.LeaderWorkerSet{}
			Eventually(func(g Gomega, ctx context.Context) error {
				return envTest.Get(ctx, types.NamespacedName{
					Name:      svcName + "-kserve-mn",
					Namespace: testNs.Name,
				}, expectedLWS)
			}).WithContext(ctx).Should(Succeed())

			By("checking the LeaderWorkerSet's top-level metadata")
			Expect(expectedLWS).To(BeOwnedBy(llmSvc))
			Expect(expectedLWS.Labels).To(HaveKeyWithValue(LocalQueueNameLabelKey, localQueueName))
			Expect(expectedLWS.Labels).ToNot(HaveKeyWithValue(testValue, testValue))

			Expect(expectedLWS.Annotations).To(HaveKeyWithValue(PreemptionReclaimAnnotationKey, preemptPriority))

			By("checking the leader pod template metadata")
			Expect(expectedLWS.Spec.LeaderWorkerTemplate.Size).To(Equal(ptr.To(int32(1))))
			Expect(expectedLWS.Spec.LeaderWorkerTemplate.LeaderTemplate).To(Not(BeNil()))
			Expect(expectedLWS.Spec.LeaderWorkerTemplate.LeaderTemplate.Labels).To(HaveKeyWithValue(constants.KServeComponentLabelKey, constants.KServeComponentWorkload))
			Expect(expectedLWS.Spec.LeaderWorkerTemplate.LeaderTemplate.Labels).To(HaveKeyWithValue(constants.LLMDRoleLabelKey, constants.LLMDRoleBoth))
			Expect(expectedLWS.Spec.LeaderWorkerTemplate.LeaderTemplate.Labels).To(HaveKeyWithValue(LocalQueueNameLabelKey, localQueueName))
			Expect(expectedLWS.Spec.LeaderWorkerTemplate.LeaderTemplate.Labels).ToNot(HaveKeyWithValue(testValue, testValue))

			Expect(expectedLWS.Spec.LeaderWorkerTemplate.LeaderTemplate.Annotations).To(HaveKeyWithValue(PreemptionReclaimAnnotationKey, preemptPriority))

			By("checking the worker pod template metadata")
			Expect(expectedLWS.Spec.LeaderWorkerTemplate.WorkerTemplate).To(Not(BeNil()))
			Expect(expectedLWS.Spec.LeaderWorkerTemplate.WorkerTemplate.Labels).To(HaveKeyWithValue(LocalQueueNameLabelKey, localQueueName))
			Expect(expectedLWS.Spec.LeaderWorkerTemplate.WorkerTemplate.Labels).ToNot(HaveKeyWithValue(testValue, testValue))

			Expect(expectedLWS.Spec.LeaderWorkerTemplate.WorkerTemplate.Annotations).To(HaveKeyWithValue(PreemptionReclaimAnnotationKey, preemptPriority))
		})

		It("should preserve externally set replicas when owner does not specify replicas", func(ctx SpecContext) {
			// given
			svcName := "test-llm-mn-preserve-replicas"
			testNs := NewTestNamespace(ctx, envTest, WithIstioShadowService(svcName))
			lwsName := types.NamespacedName{Name: svcName + "-kserve-mn", Namespace: testNs.Name}

			// Create LLMInferenceService WITHOUT specifying replicas
			llmSvc := LLMInferenceService(svcName,
				InNamespace[*v1alpha2.LLMInferenceService](testNs.Name),
				WithModelURI("hf://facebook/opt-125m"),
				WithParallelism(ParallelismSpec(
					WithDataParallelism(2),
					WithDataLocalParallelism(1),
					WithExpert(false),
				)),
				WithWorker(&corev1.PodSpec{}),
				WithManagedRoute(),
				WithManagedGateway(),
				// Note: Not using WithReplicas() - replicas should be nil
			)

			// Verify replicas is nil in the spec
			Expect(llmSvc.Spec.Replicas).To(BeNil())

			// when
			Expect(envTest.Create(ctx, llmSvc)).To(Succeed())
			defer func() {
				testNs.DeleteAndWait(ctx, llmSvc)
			}()

			// then - LeaderWorkerSet should be created with default replicas (1)
			expectedLWS := &lwsapi.LeaderWorkerSet{}
			Eventually(func(g Gomega, ctx context.Context) error {
				return envTest.Get(ctx, lwsName, expectedLWS)
			}).WithContext(ctx).Should(Succeed())

			Expect(expectedLWS.Spec.Replicas).To(Equal(ptr.To[int32](1)))

			// Simulate external scaling (e.g., HPA scaling to 3 replicas)
			errRetry := retry.RetryOnConflict(retry.DefaultRetry, func() error {
				lws := &lwsapi.LeaderWorkerSet{}
				if err := envTest.Get(ctx, lwsName, lws); err != nil {
					return err
				}
				lws.Spec.Replicas = ptr.To[int32](3)
				return envTest.Update(ctx, lws)
			})
			Expect(errRetry).ToNot(HaveOccurred())

			// Trigger a reconciliation by updating a spec field in the LLMInferenceService.
			// Using a spec change (Parallelism.Expert) rather than an annotation ensures
			// compatibility with GenerationChangedPredicate if it's added in the future.
			errRetry = retry.RetryOnConflict(retry.DefaultRetry, func() error {
				_, errUpdate := ctrl.CreateOrUpdate(ctx, envTest.Client, llmSvc, func() error {
					llmSvc.Spec.Parallelism.Expert = !llmSvc.Spec.Parallelism.Expert
					return nil
				})
				return errUpdate
			})
			Expect(errRetry).ToNot(HaveOccurred())

			// Verify the externally set replicas are preserved after reconciliation
			Consistently(func(g Gomega, ctx context.Context) {
				lws := &lwsapi.LeaderWorkerSet{}
				g.Expect(envTest.Get(ctx, lwsName, lws)).To(Succeed())

				g.Expect(lws.Spec.Replicas).To(Equal(ptr.To[int32](3)),
					"Externally set replicas should be preserved when owner doesn't specify replicas")
			}).WithContext(ctx).Should(Succeed())
		})

		It("should override externally set replicas when owner specifies replicas", func(ctx SpecContext) {
			// given
			svcName := "test-llm-mn-override-replicas"
			testNs := NewTestNamespace(ctx, envTest, WithIstioShadowService(svcName))
			lwsName := types.NamespacedName{Name: svcName + "-kserve-mn", Namespace: testNs.Name}

			// Create LLMInferenceService WITH explicit replicas
			llmSvc := LLMInferenceService(svcName,
				InNamespace[*v1alpha2.LLMInferenceService](testNs.Name),
				WithModelURI("hf://facebook/opt-125m"),
				WithParallelism(ParallelismSpec(
					WithDataParallelism(2),
					WithDataLocalParallelism(1),
					WithExpert(false),
				)),
				WithWorker(&corev1.PodSpec{}),
				WithManagedRoute(),
				WithManagedGateway(),
				WithReplicas(2), // Owner explicitly sets replicas
			)

			// Verify replicas is set in the spec
			Expect(llmSvc.Spec.Replicas).To(Equal(ptr.To[int32](2)))

			// when
			Expect(envTest.Create(ctx, llmSvc)).To(Succeed())
			defer func() {
				testNs.DeleteAndWait(ctx, llmSvc)
			}()

			// then - LeaderWorkerSet should be created with owner-specified replicas (2)
			expectedLWS := &lwsapi.LeaderWorkerSet{}
			Eventually(func(g Gomega, ctx context.Context) error {
				return envTest.Get(ctx, lwsName, expectedLWS)
			}).WithContext(ctx).Should(Succeed())

			Expect(expectedLWS.Spec.Replicas).To(Equal(ptr.To[int32](2)))

			// Simulate external scaling attempt (e.g., someone tries to manually scale to 5)
			errRetry := retry.RetryOnConflict(retry.DefaultRetry, func() error {
				lws := &lwsapi.LeaderWorkerSet{}
				if err := envTest.Get(ctx, lwsName, lws); err != nil {
					return err
				}
				lws.Spec.Replicas = ptr.To[int32](5)
				return envTest.Update(ctx, lws)
			})
			Expect(errRetry).ToNot(HaveOccurred())

			// Trigger a reconciliation by updating a spec field in the LLMInferenceService.
			// Using a spec change (Parallelism.Expert) rather than an annotation ensures
			// compatibility with GenerationChangedPredicate if it's added in the future.
			errRetry = retry.RetryOnConflict(retry.DefaultRetry, func() error {
				_, errUpdate := ctrl.CreateOrUpdate(ctx, envTest.Client, llmSvc, func() error {
					llmSvc.Spec.Parallelism.Expert = !llmSvc.Spec.Parallelism.Expert
					return nil
				})
				return errUpdate
			})
			Expect(errRetry).ToNot(HaveOccurred())

			// Verify the owner-specified replicas override the external change
			Eventually(func(g Gomega, ctx context.Context) {
				lws := &lwsapi.LeaderWorkerSet{}
				g.Expect(envTest.Get(ctx, lwsName, lws)).To(Succeed())

				g.Expect(lws.Spec.Replicas).To(Equal(ptr.To[int32](2)),
					"Owner-specified replicas should override externally set replicas")
			}).WithContext(ctx).Should(Succeed())
		})
	})
})