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"
	"sigs.k8s.io/controller-runtime/pkg/client"
	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"
	resourcev1 "k8s.io/api/resource/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/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"

	. "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))

			// Verify status.workloads references
			Eventually(func(g Gomega, ctx context.Context) {
				current := &v1alpha2.LLMInferenceService{}
				g.Expect(envTest.Get(ctx, types.NamespacedName{Name: svcName, Namespace: testNs.Name}, current)).To(Succeed())
				g.Expect(current.Status.Workloads).NotTo(BeNil())
				g.Expect(current.Status.Workloads.Primary).To(Equal(&corev1.TypedLocalObjectReference{
					APIGroup: ptr.To("leaderworkerset.x-k8s.io"),
					Kind:     "LeaderWorkerSet",
					Name:     kmeta.ChildName(svcName, "-kserve-mn"),
				}))
				g.Expect(current.Status.Workloads.Service).To(Equal(&corev1.TypedLocalObjectReference{
					Kind: "Service",
					Name: kmeta.ChildName(svcName, "-kserve-workload-svc"),
				}))
				g.Expect(current.Status.Workloads.Prefill).To(Equal(&corev1.TypedLocalObjectReference{
					APIGroup: ptr.To("apps"),
					Kind:     "Deployment",
					Name:     kmeta.ChildName(svcName, "-kserve-prefill"),
				}))
				g.Expect(current.Status.Workloads.Scheduler).To(BeNil())
			}).WithContext(ctx).Should(Succeed())
		})

		DescribeTable("should create a multi-node deployment with managed DRA",
			func(ctx SpecContext, testName string, podSpec *corev1.PodSpec, extraAnnotations map[string]string, targetContainer string) {
				// given
				svcName := "test-llm-multinode-dra-" + testName
				testNs := NewTestNamespace(ctx, envTest, WithIstioShadowService(svcName))

				annotations := map[string]string{
					constants.ManagedDRADeviceClassAnnotationKey: "gpu.nvidia.com",
					constants.ManagedDRADeviceCountAnnotationKey: "8",
				}
				for k, v := range extraAnnotations {
					annotations[k] = v
				}

				llmSvc := LLMInferenceService(svcName,
					InNamespace[*v1alpha2.LLMInferenceService](testNs.Name),
					WithModelURI("hf://facebook/opt-125m"),
					WithReplicas(2),
					WithParallelism(ParallelismSpec(
						WithDataParallelism(4),
						WithDataLocalParallelism(1),
						WithTensorParallelism(3),
					)),
					WithTemplate(podSpec.DeepCopy()),
					WithWorker(podSpec.DeepCopy()),
					WithAnnotations(annotations),
					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())

				// Verify the ResourceClaimTemplate was created with correct values
				expectedTemplate := &resourcev1.ResourceClaimTemplate{}
				Eventually(func(g Gomega, ctx context.Context) error {
					return envTest.Get(ctx, types.NamespacedName{
						Name:      svcName + "-managed-dra",
						Namespace: testNs.Name,
					}, expectedTemplate)
				}).WithContext(ctx).Should(Succeed())

				// Check that it parsed the annotations correctly
				Expect(expectedTemplate.Spec.Spec.Devices.Requests).To(HaveLen(1))
				Expect(expectedTemplate.Spec.Spec.Devices.Requests[0].Name).To(Equal("device"))
				Expect(expectedTemplate.Spec.Spec.Devices.Requests[0].Exactly.DeviceClassName).To(Equal("gpu.nvidia.com"))
				Expect(expectedTemplate.Spec.Spec.Devices.Requests[0].Exactly.Count).To(Equal(int64(8)))

				verifyPod := func(pod corev1.PodSpec, role string) {
					Expect(pod.ResourceClaims).To(HaveLen(1), "%s pod-level claim", role)
					Expect(pod.ResourceClaims[0].Name).To(Equal("managed-device"))
					Expect(pod.ResourceClaims[0].ResourceClaimTemplateName).To(Not(BeNil()))
					Expect(*pod.ResourceClaims[0].ResourceClaimTemplateName).To(Equal(svcName + "-managed-dra"))

					var targetSeen bool
					for _, c := range pod.Containers {
						if c.Name == targetContainer {
							targetSeen = true
							Expect(c.Resources.Claims).To(HaveLen(1), "%s %s should receive the device claim", role, c.Name)
							Expect(c.Resources.Claims[0].Name).To(Equal("managed-device"))
						} else {
							Expect(c.Resources.Claims).To(BeEmpty(), "%s %s must not receive the device claim", role, c.Name)
						}
					}
					Expect(targetSeen).To(BeTrue(), "%s target container %q should be present", role, targetContainer)
				}

				// Verify leader template has DRA
				verifyPod(expectedLWS.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec, "leader")
				// Verify worker template has DRA
				verifyPod(expectedLWS.Spec.LeaderWorkerTemplate.WorkerTemplate.Spec, "worker")
			},
			Entry("main-named workload, no container annotation -> claim on main",
				"main",
				SimpleWorkerPodSpec(),
				nil,
				"main",
			),
			Entry("non-main first container, no container annotation -> claim on first container",
				"fallback",
				&corev1.PodSpec{Containers: []corev1.Container{
					{Name: "vllm", Image: "test-worker:latest"},
					{Name: "sidecar", Image: "busybox:latest"},
				}},
				nil,
				"vllm",
			),
			Entry("container annotation set -> claim on annotated container",
				"annotated",
				&corev1.PodSpec{Containers: []corev1.Container{
					{Name: "init", Image: "busybox:latest"},
					{Name: "worker", Image: "test-worker:latest"},
					{Name: "sidecar", Image: "busybox:latest"},
				}},
				map[string]string{constants.ManagedDRAContainerNameAnnotationKey: "worker"},
				"worker",
			),
		)

		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 multi-node deployment with prefill workload and managed DRA", func(ctx SpecContext) {
			// given
			svcName := "test-llm-mn-pd-dra"
			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(2),
				)),
				WithTemplate(SimpleWorkerPodSpec()),
				WithWorker(SimpleWorkerPodSpec()),
				WithPrefill(SimpleWorkerPodSpec()),
				WithPrefillWorker(SimpleWorkerPodSpec()),
				WithPrefillParallelism(ParallelismSpec(
					WithDataParallelism(2),
					WithDataLocalParallelism(1),
					WithTensorParallelism(2),
				)),
				WithPrefillReplicas(1),
				WithAnnotations(map[string]string{
					constants.ManagedDRADeviceClassAnnotationKey: "gpu.nvidia.com",
					constants.ManagedDRADeviceCountAnnotationKey: "4",
				}),
				WithManagedRoute(),
				WithManagedGateway(),
			)

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

			// then
			expectedTemplate := &resourcev1.ResourceClaimTemplate{}
			Eventually(func(g Gomega, ctx context.Context) error {
				return envTest.Get(ctx, types.NamespacedName{
					Name:      svcName + "-managed-dra",
					Namespace: testNs.Name,
				}, expectedTemplate)
			}).WithContext(ctx).Should(Succeed())

			Expect(expectedTemplate.Spec.Spec.Devices.Requests).To(HaveLen(1))
			Expect(expectedTemplate.Spec.Spec.Devices.Requests[0].Exactly.DeviceClassName).To(Equal("gpu.nvidia.com"))
			Expect(expectedTemplate.Spec.Spec.Devices.Requests[0].Exactly.Count).To(Equal(int64(4)))

			rctName := svcName + "-managed-dra"

			expectDRAInjected := func(spec corev1.PodSpec, podLabel string) {
				Expect(spec.ResourceClaims).To(HaveLen(1), "%s: pod-level resourceClaim missing", podLabel)
				Expect(spec.ResourceClaims[0].Name).To(Equal("managed-device"), "%s", podLabel)
				Expect(spec.ResourceClaims[0].ResourceClaimTemplateName).To(Not(BeNil()), "%s", podLabel)
				Expect(*spec.ResourceClaims[0].ResourceClaimTemplateName).To(Equal(rctName), "%s", podLabel)

				Expect(spec.Containers).ToNot(BeEmpty(), "%s", podLabel)
				Expect(spec.Containers[0].Name).To(Equal("main"), "%s", podLabel)
				Expect(spec.Containers[0].Resources.Claims).To(HaveLen(1), "%s: main-container claim missing", podLabel)
				Expect(spec.Containers[0].Resources.Claims[0].Name).To(Equal("managed-device"), "%s", podLabel)
			}

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

			Expect(decodeLWS.Spec.LeaderWorkerTemplate.LeaderTemplate).ToNot(BeNil())
			expectDRAInjected(decodeLWS.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec, "decode-leader")
			expectDRAInjected(decodeLWS.Spec.LeaderWorkerTemplate.WorkerTemplate.Spec, "decode-worker")

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

			Expect(prefillLWS.Spec.LeaderWorkerTemplate.LeaderTemplate).ToNot(BeNil())
			expectDRAInjected(prefillLWS.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec, "prefill-leader")
			expectDRAInjected(prefillLWS.Spec.LeaderWorkerTemplate.WorkerTemplate.Spec, "prefill-worker")
		})

		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())
		})

		// NOTE: Topology transition test (LWS -> Deployment when Worker
		// is removed) is covered by the unit test in
		// workload_status_test.go. An envtest variant cannot be written
		// because CRD structural schema validation rejects nulling out
		// a PodSpec field (containers: null). Same limitation affects
		// the existing "should delete multi-node resources" test above.

		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 not propagate model-based-routing annotation to LWS pod templates", func(ctx SpecContext) {
			svcName := "test-llm-lws-no-mbr"
			testNs := NewTestNamespace(ctx, envTest, WithIstioShadowService(svcName))

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

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

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

			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.LeaderWorkerTemplate.LeaderTemplate).To(Not(BeNil()))
			Expect(expectedLWS.Spec.LeaderWorkerTemplate.LeaderTemplate.Annotations).
				NotTo(HaveKey(llmisvc.AnnotationModelBasedRoutingEnabled))
			Expect(expectedLWS.Spec.LeaderWorkerTemplate.WorkerTemplate.Annotations).
				NotTo(HaveKey(llmisvc.AnnotationModelBasedRoutingEnabled))
		})

		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())
		})
	})

	Context("Readiness Event Emission", func() {
		It("should emit a Warning LLMInferenceServiceNotReady Event with WorkerWorkloadReady when transitioning Ready to NotReady", func(ctx SpecContext) {
			svcName := "test-llm-mn-event-notready"
			testNs := NewTestNamespace(ctx, envTest, WithIstioShadowService(svcName))

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

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

			lwsName := types.NamespacedName{
				Name:      kmeta.ChildName(svcName, "-kserve-mn"),
				Namespace: testNs.Name,
			}

			Eventually(func(g Gomega, ctx context.Context) error {
				return envTest.Get(ctx, lwsName, &lwsapi.LeaderWorkerSet{})
			}).WithContext(ctx).Should(Succeed())

			ensureRouterManagedResourcesAreReady(ctx, envTest.Client, llmSvc)
			ensureLWSReady(ctx, envTest.Client, lwsName)

			Eventually(func(g Gomega, ctx context.Context) {
				current := &v1alpha2.LLMInferenceService{}
				g.Expect(envTest.Get(ctx, client.ObjectKeyFromObject(llmSvc), current)).To(Succeed())
				g.Expect(current.Status).To(HaveCondition("Ready", "True"))
			}).WithContext(ctx).Should(Succeed())

			setLWSNotReady(ctx, envTest.Client, lwsName)

			Eventually(func(g Gomega, ctx context.Context) {
				current := &v1alpha2.LLMInferenceService{}
				g.Expect(envTest.Get(ctx, client.ObjectKeyFromObject(llmSvc), current)).To(Succeed())
				g.Expect(current.Status).To(HaveCondition("Ready", "False"))
			}).WithContext(ctx).Should(Succeed())

			Eventually(func(g Gomega, ctx context.Context) {
				ev := findEvent(ctx, envTest.Client, llmSvc, string(llmisvc.LLMInferenceServiceNotReadyState))
				g.Expect(ev).NotTo(BeNil(), "Expected a Warning LLMInferenceServiceNotReady Event")
				g.Expect(ev.Message).To(ContainSubstring(string(v1alpha2.WorkerWorkloadReady)))
			}).WithContext(ctx).Should(Succeed())
		})
	})
})

func ensureLWSReady(ctx context.Context, c client.Client, lwsName types.NamespacedName) {
	if envTest.UsingExistingCluster() {
		return
	}

	Eventually(func(g Gomega, ctx context.Context) {
		lws := &lwsapi.LeaderWorkerSet{}
		g.Expect(c.Get(ctx, lwsName, lws)).To(Succeed())

		lws.Status.Conditions = []metav1.Condition{
			{
				Type:               string(lwsapi.LeaderWorkerSetAvailable),
				Status:             metav1.ConditionTrue,
				Reason:             "AllReplicasAvailable",
				Message:            "All replicas are available",
				LastTransitionTime: metav1.Now(),
			},
		}
		g.Expect(c.Status().Update(ctx, lws)).To(Succeed())
	}).WithContext(ctx).Should(Succeed())
}

func setLWSNotReady(ctx context.Context, c client.Client, lwsName types.NamespacedName) {
	Eventually(func(g Gomega, ctx context.Context) {
		lws := &lwsapi.LeaderWorkerSet{}
		g.Expect(c.Get(ctx, lwsName, lws)).To(Succeed())

		lws.Status.Conditions = []metav1.Condition{
			{
				Type:               string(lwsapi.LeaderWorkerSetAvailable),
				Status:             metav1.ConditionFalse,
				Reason:             "ReplicasUnavailable",
				Message:            "Not all replicas are available",
				LastTransitionTime: metav1.Now(),
			},
		}
		g.Expect(c.Status().Update(ctx, lws)).To(Succeed())
	}).WithContext(ctx).Should(Succeed())
}