kserve/pkg/controller/v1alpha2/llmisvc/workload_multi_node.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

import (
	"context"
	"fmt"

	corev1 "k8s.io/api/core/v1"
	rbacv1 "k8s.io/api/rbac/v1"
	"k8s.io/apimachinery/pkg/api/equality"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/api/meta"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/client-go/util/retry"
	"knative.dev/pkg/kmeta"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/log"
	lwsapi "sigs.k8s.io/lws/api/leaderworkerset/v1"

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

func (r *LLMISVCReconciler) reconcileMultiNodeWorkload(ctx context.Context, llmSvc *v1alpha2.LLMInferenceService, config *Config) error {
	log.FromContext(ctx).Info("Reconciling multi-node workload")

	if err := r.reconcileManagedDRA(ctx, llmSvc); err != nil {
		return fmt.Errorf("failed to reconcile managed DRA: %w", err)
	}

	if err := r.reconcileMultiNodeMainServiceAccount(ctx, llmSvc, config); err != nil {
		return fmt.Errorf("failed to reconcile multi-node service account: %w", err)
	}
	if err := r.reconcileMultiNodePrefillServiceAccount(ctx, llmSvc); err != nil {
		return fmt.Errorf("failed to reconcile multi-node service account: %w", err)
	}
	if err := r.reconcileMultiNodeMainWorkload(ctx, llmSvc, config); err != nil {
		return fmt.Errorf("failed to reconcile multi-node main workload: %w", err)
	}
	if err := r.reconcileMultiNodePrefillWorkload(ctx, llmSvc, config); err != nil {
		return fmt.Errorf("failed to reconcile multi-node prefill workload: %w", err)
	}
	return nil
}

func (r *LLMISVCReconciler) reconcileMultiNodeMainWorkload(ctx context.Context, llmSvc *v1alpha2.LLMInferenceService, config *Config) error {
	if isStopped := utils.GetForceStopRuntime(llmSvc); isStopped || llmSvc.Spec.Worker == nil {
		if isStopped {
			llmSvc.MarkWorkerWorkloadNotReady("Stopped", "Service is stopped")
		} else {
			llmSvc.MarkWorkerWorkloadUnset()
		}
		return Delete(ctx, r, llmSvc, &lwsapi.LeaderWorkerSet{
			ObjectMeta: metav1.ObjectMeta{
				Name:      mainLWSName(llmSvc),
				Namespace: llmSvc.GetNamespace(),
			},
		})
	}

	expected, err := r.expectedMainMultiNodeLWS(ctx, llmSvc, config)
	if err != nil {
		return fmt.Errorf("failed to build the expected main LWS: %w", err)
	}
	if err := Reconcile(ctx, r, llmSvc, &lwsapi.LeaderWorkerSet{}, expected, semanticLWSIsEqual, PreserveLWSReplicas()); err != nil {
		return err
	}
	return r.propagateLeaderWorkerSetStatus(ctx, expected, llmSvc.MarkWorkerWorkloadReady, llmSvc.MarkWorkerWorkloadNotReady)
}

func (r *LLMISVCReconciler) reconcileMultiNodePrefillWorkload(ctx context.Context, llmSvc *v1alpha2.LLMInferenceService, config *Config) error {
	if isStopped := utils.GetForceStopRuntime(llmSvc); isStopped || llmSvc.Spec.Prefill == nil || llmSvc.Spec.Prefill.Worker == nil {
		if isStopped {
			llmSvc.MarkPrefillWorkerWorkloadNotReady("Stopped", "Service is stopped")
		} else {
			llmSvc.MarkPrefillWorkerWorkloadUnset()
		}
		return Delete(ctx, r, llmSvc, &lwsapi.LeaderWorkerSet{
			ObjectMeta: metav1.ObjectMeta{
				Name:      prefillLWSName(llmSvc),
				Namespace: llmSvc.GetNamespace(),
			},
		})
	}

	expected, err := r.expectedPrefillMultiNodeLWS(ctx, llmSvc, config)
	if err != nil {
		return fmt.Errorf("failed to build the expected prefill LWS: %w", err)
	}
	if err := Reconcile(ctx, r, llmSvc, &lwsapi.LeaderWorkerSet{}, expected, semanticLWSIsEqual, PreserveLWSReplicas()); err != nil {
		return err
	}
	return r.propagateLeaderWorkerSetStatus(ctx, expected, llmSvc.MarkPrefillWorkerWorkloadReady, llmSvc.MarkPrefillWorkerWorkloadNotReady)
}

func (r *LLMISVCReconciler) propagateLeaderWorkerSetStatus(ctx context.Context, expected *lwsapi.LeaderWorkerSet, ready func(), notReady func(reason string, messageFormat string, messageA ...interface{})) error {
	logger := log.FromContext(ctx)

	curr := &lwsapi.LeaderWorkerSet{}
	err := retry.OnError(retry.DefaultRetry, apierrors.IsNotFound, func() error {
		return r.Get(ctx, client.ObjectKeyFromObject(expected), curr)
	})
	if err != nil {
		return fmt.Errorf("failed to get current leaderworkerset %s/%s: %w", expected.GetNamespace(), expected.GetName(), err)
	}
	for _, cond := range curr.Status.Conditions {
		if cond.Type == string(lwsapi.LeaderWorkerSetAvailable) {
			if cond.Status == metav1.ConditionTrue {
				ready()
			} else {
				notReady(cond.Reason, cond.Message)
			}
			return nil
		}
	}
	logger.Info("LeaderWorkerSet processed")
	notReady(string(lwsapi.LeaderWorkerSetProgressing), "LWS is progressing")
	return nil
}

func (r *LLMISVCReconciler) expectedMainMultiNodeLWS(ctx context.Context, llmSvc *v1alpha2.LLMInferenceService, config *Config) (*lwsapi.LeaderWorkerSet, error) {
	workerLabels := map[string]string{
		constants.KubernetesComponentLabelKey: constants.LLMComponentWorkloadWorker,
		constants.KubernetesAppNameLabelKey:   llmSvc.GetName(),
		constants.KubernetesPartOfLabelKey:    constants.LLMInferenceServicePartOfValue,
	}
	if llmSvc.Spec.Template == nil {
		// When there is no leader template, workers become part of the InferencePool selector.
		workerLabels[constants.KServeComponentLabelKey] = constants.KServeComponentWorkload
		workerLabels[constants.LLMDRoleLabelKey] = constants.LLMDRoleDecode
	}
	role := constants.LLMDRoleDecode
	if llmSvc.Spec.Prefill == nil {
		role = constants.LLMDRoleBoth
	}
	leaderLabels := map[string]string{
		constants.KubernetesComponentLabelKey: constants.LLMComponentWorkloadLeader,
		constants.KubernetesAppNameLabelKey:   llmSvc.GetName(),
		constants.KubernetesPartOfLabelKey:    constants.LLMInferenceServicePartOfValue,
		constants.KServeComponentLabelKey:     constants.KServeComponentWorkload,
		constants.LLMDRoleLabelKey:            role,
	}

	expected := &lwsapi.LeaderWorkerSet{
		ObjectMeta: metav1.ObjectMeta{
			Name:      mainLWSName(llmSvc),
			Namespace: llmSvc.GetNamespace(),
			OwnerReferences: []metav1.OwnerReference{
				*metav1.NewControllerRef(llmSvc, v1alpha2.LLMInferenceServiceGVK),
			},
			Labels: workerLabels,
		},
		Spec: lwsapi.LeaderWorkerSetSpec{
			Replicas: llmSvc.Spec.Replicas,
			LeaderWorkerTemplate: lwsapi.LeaderWorkerTemplate{
				Size: llmSvc.Spec.Parallelism.GetSize(),
				WorkerTemplate: corev1.PodTemplateSpec{
					ObjectMeta: metav1.ObjectMeta{
						Labels: workerLabels,
					},
				},
				RestartPolicy: lwsapi.RecreateGroupOnPodRestart,
			},
			RolloutStrategy: lwsapi.RolloutStrategy{
				Type: lwsapi.RollingUpdateStrategyType,
			},
			StartupPolicy: lwsapi.LeaderCreatedStartupPolicy,
		},
	}

	// Fetch the current LWS once to preserve storage-init images across upgrades.
	currLWS := &lwsapi.LeaderWorkerSet{}
	if err := r.Get(ctx, client.ObjectKeyFromObject(expected), currLWS); err != nil && !apierrors.IsNotFound(err) && !meta.IsNoMatchError(err) {
		return nil, fmt.Errorf("failed to get current leader worker set %s/%s: %w", expected.GetNamespace(), expected.GetName(), err)
	}

	if llmSvc.Spec.Template != nil && !utils.GetForceStopRuntime(llmSvc) {
		expected.Spec.LeaderWorkerTemplate.LeaderTemplate = &corev1.PodTemplateSpec{
			ObjectMeta: metav1.ObjectMeta{
				Labels: leaderLabels,
			},
			Spec: *llmSvc.Spec.Template.DeepCopy(),
		}

		serviceAccount, _, err := r.expectedMultiNodeMainServiceAccount(ctx, llmSvc)
		if err != nil {
			return nil, fmt.Errorf("failed to create expected multi node service account: %w", err)
		}
		expected.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.ServiceAccountName = serviceAccount.GetName()

		var currLeaderSpec corev1.PodSpec
		if currLWS.Spec.LeaderWorkerTemplate.LeaderTemplate != nil {
			currLeaderSpec = currLWS.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec
		}

		if err := r.attachModelArtifacts(ctx, serviceAccount, llmSvc, currLeaderSpec, &expected.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec, config, "main", constants.DefaultModelLocalMountPath, len(config.ResolvedLoRAAdapters) > 0); err != nil {
			return nil, fmt.Errorf("failed to attach model artifacts to leader template: %w", err)
		}

		if hasRoutingSidecar(expected.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec) {
			log.FromContext(ctx).V(2).Info("Main container has a routing sidecar")

			s := routingSidecar(&expected.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec)
			if llmSvc.Spec.Router != nil {
				s.Env = append(s.Env, corev1.EnvVar{
					Name:  "INFERENCE_POOL_NAME",
					Value: llmSvc.Spec.Router.Scheduler.InferencePoolName(llmSvc),
				})
			}
		}
	}
	if llmSvc.Spec.Worker != nil && !utils.GetForceStopRuntime(llmSvc) {
		expected.Spec.LeaderWorkerTemplate.WorkerTemplate.Spec = *llmSvc.Spec.Worker.DeepCopy()

		serviceAccount, _, err := r.expectedMultiNodeMainServiceAccount(ctx, llmSvc)
		if err != nil {
			return nil, fmt.Errorf("failed to create expected multi node service account: %w", err)
		}
		expected.Spec.LeaderWorkerTemplate.WorkerTemplate.Spec.ServiceAccountName = serviceAccount.GetName()

		if err := r.attachModelArtifacts(ctx, serviceAccount, llmSvc, currLWS.Spec.LeaderWorkerTemplate.WorkerTemplate.Spec, &expected.Spec.LeaderWorkerTemplate.WorkerTemplate.Spec, config, "main", constants.DefaultModelLocalMountPath, len(config.ResolvedLoRAAdapters) > 0); err != nil {
			return nil, fmt.Errorf("failed to attach model artifacts to worker template: %w", err)
		}

		if hasRoutingSidecar(expected.Spec.LeaderWorkerTemplate.WorkerTemplate.Spec) {
			log.FromContext(ctx).V(2).Info("Main (worker) container has a routing sidecar")

			s := routingSidecar(&expected.Spec.LeaderWorkerTemplate.WorkerTemplate.Spec)
			if llmSvc.Spec.Router != nil {
				s.Env = append(s.Env, corev1.EnvVar{
					Name:  "INFERENCE_POOL_NAME",
					Value: llmSvc.Spec.Router.Scheduler.InferencePoolName(llmSvc),
				})
			}
		}
	}

	r.propagateTopLevelLeaderWorkerSetMetadata(llmSvc, expected)

	if expected.Spec.LeaderWorkerTemplate.LeaderTemplate != nil {
		utils.PropagateMap(llmSvc.Spec.Labels, &expected.Spec.LeaderWorkerTemplate.LeaderTemplate.Labels)
		utils.PropagateMap(llmSvc.Spec.Annotations, &expected.Spec.LeaderWorkerTemplate.LeaderTemplate.Annotations, AnnotationModelBasedRoutingEnabled)
	}
	utils.PropagateMap(llmSvc.Spec.Labels, &expected.Spec.LeaderWorkerTemplate.WorkerTemplate.Labels)
	utils.PropagateMap(llmSvc.Spec.Annotations, &expected.Spec.LeaderWorkerTemplate.WorkerTemplate.Annotations, AnnotationModelBasedRoutingEnabled)

	// Inject tracing instrumentation when spec.tracing is set
	if llmSvc.Spec.Tracing != nil {
		if expected.Spec.LeaderWorkerTemplate.LeaderTemplate != nil {
			injectServerTracingIntoPodSpec(llmSvc.Spec.Tracing, llmSvc.GetNamespace(), llmSvc.GetName(), "-decode", &expected.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec)
		}
		injectServerTracingIntoPodSpec(llmSvc.Spec.Tracing, llmSvc.GetNamespace(), llmSvc.GetName(), "-decode", &expected.Spec.LeaderWorkerTemplate.WorkerTemplate.Spec)
	}

	log.FromContext(ctx).V(2).Info("Expected main LWS", "leaderworkerset", expected)

	return expected, nil
}

func (r *LLMISVCReconciler) expectedPrefillMultiNodeLWS(ctx context.Context, llmSvc *v1alpha2.LLMInferenceService, config *Config) (*lwsapi.LeaderWorkerSet, error) {
	workerLabels := map[string]string{
		constants.KubernetesComponentLabelKey: constants.LLMComponentWorkloadWorkerPrefill,
		constants.KubernetesAppNameLabelKey:   llmSvc.GetName(),
		constants.KubernetesPartOfLabelKey:    constants.LLMInferenceServicePartOfValue,
	}
	if llmSvc.Spec.Prefill != nil && llmSvc.Spec.Prefill.Template == nil {
		// When there is no leader template, workers become part of the InferencePool selector.
		workerLabels[constants.KServeComponentLabelKey] = constants.KServeComponentWorkload
		workerLabels[constants.LLMDRoleLabelKey] = constants.LLMDRolePrefill
	}
	leaderLabels := map[string]string{
		constants.KubernetesComponentLabelKey: constants.LLMComponentWorkloadLeaderPrefill,
		constants.KubernetesAppNameLabelKey:   llmSvc.GetName(),
		constants.KubernetesPartOfLabelKey:    constants.LLMInferenceServicePartOfValue,
		constants.KServeComponentLabelKey:     constants.KServeComponentWorkload,
		constants.LLMDRoleLabelKey:            constants.LLMDRolePrefill,
	}

	expected := &lwsapi.LeaderWorkerSet{
		ObjectMeta: metav1.ObjectMeta{
			Name:      prefillLWSName(llmSvc),
			Namespace: llmSvc.GetNamespace(),
			OwnerReferences: []metav1.OwnerReference{
				*metav1.NewControllerRef(llmSvc, v1alpha2.LLMInferenceServiceGVK),
			},
			Labels: workerLabels,
		},
		Spec: lwsapi.LeaderWorkerSetSpec{
			LeaderWorkerTemplate: lwsapi.LeaderWorkerTemplate{
				WorkerTemplate: corev1.PodTemplateSpec{
					ObjectMeta: metav1.ObjectMeta{
						Labels: workerLabels,
					},
				},
				RestartPolicy: lwsapi.RecreateGroupOnPodRestart,
			},
			RolloutStrategy: lwsapi.RolloutStrategy{
				Type: lwsapi.RollingUpdateStrategyType,
			},
			StartupPolicy: lwsapi.LeaderCreatedStartupPolicy,
		},
	}

	if llmSvc.Spec.Prefill != nil && !utils.GetForceStopRuntime(llmSvc) {
		expected.Spec.Replicas = llmSvc.Spec.Prefill.Replicas
		expected.Spec.LeaderWorkerTemplate.Size = llmSvc.Spec.Prefill.Parallelism.GetSize()

		serviceAccount, _, err := r.expectedMultiNodePrefillServiceAccount(ctx, llmSvc)
		if err != nil {
			return nil, fmt.Errorf("failed to create exptected multi node service account: %w", err)
		}

		currLWS := &lwsapi.LeaderWorkerSet{}
		if err := r.Get(ctx, client.ObjectKeyFromObject(expected), currLWS); err != nil && !apierrors.IsNotFound(err) && !meta.IsNoMatchError(err) {
			return nil, fmt.Errorf("failed to get current prefill leader worker set %s/%s: %w", expected.GetNamespace(), expected.GetName(), err)
		}

		if llmSvc.Spec.Prefill.Template != nil {
			expected.Spec.LeaderWorkerTemplate.LeaderTemplate = &corev1.PodTemplateSpec{
				ObjectMeta: metav1.ObjectMeta{
					Labels: leaderLabels,
				},
				Spec: *llmSvc.Spec.Prefill.Template.DeepCopy(),
			}

			expected.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.ServiceAccountName = serviceAccount.GetName()

			var currLeaderSpec corev1.PodSpec
			if currLWS.Spec.LeaderWorkerTemplate.LeaderTemplate != nil {
				currLeaderSpec = currLWS.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec
			}

			if err := r.attachModelArtifacts(ctx, serviceAccount, llmSvc, currLeaderSpec, &expected.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec, config, "main", constants.DefaultModelLocalMountPath, len(config.ResolvedLoRAAdapters) > 0); err != nil {
				return nil, fmt.Errorf("failed to attach model artifacts to prefill leader template: %w", err)
			}
		}
		if llmSvc.Spec.Prefill.Worker != nil {
			expected.Spec.LeaderWorkerTemplate.WorkerTemplate.Spec = *llmSvc.Spec.Prefill.Worker.DeepCopy()

			expected.Spec.LeaderWorkerTemplate.WorkerTemplate.Spec.ServiceAccountName = serviceAccount.GetName()

			if err := r.attachModelArtifacts(ctx, serviceAccount, llmSvc, currLWS.Spec.LeaderWorkerTemplate.WorkerTemplate.Spec, &expected.Spec.LeaderWorkerTemplate.WorkerTemplate.Spec, config, "main", constants.DefaultModelLocalMountPath, len(config.ResolvedLoRAAdapters) > 0); err != nil {
				return nil, fmt.Errorf("failed to attach model artifacts to prefill worker template: %w", err)
			}
		}

		if llmSvc.Spec.Prefill.Parallelism.IsDataParallel() && expected.Spec.LeaderWorkerTemplate.Size != nil {
			expected.Spec.LeaderWorkerTemplate.SubGroupPolicy = &lwsapi.SubGroupPolicy{
				SubGroupSize: expected.Spec.LeaderWorkerTemplate.Size,
			}
		}
	}

	r.propagateTopLevelLeaderWorkerSetMetadata(llmSvc, expected)

	// Inject tracing instrumentation when spec.tracing is set
	if llmSvc.Spec.Tracing != nil {
		if expected.Spec.LeaderWorkerTemplate.LeaderTemplate != nil {
			injectServerTracingIntoPodSpec(llmSvc.Spec.Tracing, llmSvc.GetNamespace(), llmSvc.GetName(), "-prefill", &expected.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec)
		}
		injectServerTracingIntoPodSpec(llmSvc.Spec.Tracing, llmSvc.GetNamespace(), llmSvc.GetName(), "-prefill", &expected.Spec.LeaderWorkerTemplate.WorkerTemplate.Spec)
	}

	if llmSvc.Spec.Prefill != nil {
		if expected.Spec.LeaderWorkerTemplate.LeaderTemplate != nil {
			utils.PropagateMap(llmSvc.Spec.Prefill.Labels, &expected.Spec.LeaderWorkerTemplate.LeaderTemplate.Labels)
			utils.PropagateMap(llmSvc.Spec.Prefill.Annotations, &expected.Spec.LeaderWorkerTemplate.LeaderTemplate.Annotations, AnnotationModelBasedRoutingEnabled)
		}
		utils.PropagateMap(llmSvc.Spec.Prefill.Labels, &expected.Spec.LeaderWorkerTemplate.WorkerTemplate.Labels)
		utils.PropagateMap(llmSvc.Spec.Prefill.Annotations, &expected.Spec.LeaderWorkerTemplate.WorkerTemplate.Annotations, AnnotationModelBasedRoutingEnabled)
	}

	log.FromContext(ctx).V(2).Info("Expected prefill LWS", "leaderworkerset", expected)

	return expected, nil
}

func (r *LLMISVCReconciler) reconcileMultiNodeMainServiceAccount(ctx context.Context, llmSvc *v1alpha2.LLMInferenceService, config *Config) error {
	serviceAccount, useExistingServiceAccount, err := r.expectedMultiNodeMainServiceAccount(ctx, llmSvc)
	if err != nil {
		return fmt.Errorf("failed to create expected multi node service account: %w", err)
	}
	if !useExistingServiceAccount {
		if utils.GetForceStopRuntime(llmSvc) || llmSvc.Spec.Worker == nil {
			return Delete(ctx, r, llmSvc, serviceAccount)
		}

		if err := Reconcile(ctx, r, llmSvc, &corev1.ServiceAccount{}, serviceAccount, semanticServiceAccountIsEqual); err != nil {
			return fmt.Errorf("failed to reconcile multi node service account %s/%s: %w", serviceAccount.GetNamespace(), serviceAccount.GetName(), err)
		}
	}
	if err := r.reconcileMultiNodeMainRole(ctx, llmSvc, config); err != nil {
		return err
	}

	return r.reconcileMultiNodeMainRoleBinding(ctx, llmSvc, serviceAccount, config)
}

func (r *LLMISVCReconciler) reconcileMultiNodePrefillServiceAccount(ctx context.Context, llmSvc *v1alpha2.LLMInferenceService) error {
	serviceAccount, useExistingServiceAccount, err := r.expectedMultiNodePrefillServiceAccount(ctx, llmSvc)
	if err != nil {
		return fmt.Errorf("failed to create expected multi node service account: %w", err)
	}
	if !useExistingServiceAccount {
		if utils.GetForceStopRuntime(llmSvc) || llmSvc.Spec.Prefill == nil || llmSvc.Spec.Prefill.Worker == nil {
			return Delete(ctx, r, llmSvc, serviceAccount)
		}

		if err := Reconcile(ctx, r, llmSvc, &corev1.ServiceAccount{}, serviceAccount, semanticServiceAccountIsEqual); err != nil {
			return fmt.Errorf("failed to reconcile multi node service account %s/%s: %w", serviceAccount.GetNamespace(), serviceAccount.GetName(), err)
		}
	}

	return nil
}

func (r *LLMISVCReconciler) reconcileMultiNodeMainRole(ctx context.Context, llmSvc *v1alpha2.LLMInferenceService, config *Config) error {
	lws, err := r.expectedMainMultiNodeLWS(ctx, llmSvc, config)
	if err != nil {
		return fmt.Errorf("failed to build the expected main LWS for building the Role: %w", err)
	}

	role := r.expectedMultiNodeMainRole(llmSvc)
	if !hasRoutingSidecar(lws.Spec.LeaderWorkerTemplate.WorkerTemplate.Spec) && (lws.Spec.LeaderWorkerTemplate.LeaderTemplate == nil || !hasRoutingSidecar(lws.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec)) {
		return Delete(ctx, r, llmSvc, role)
	}

	if err := Reconcile(ctx, r, llmSvc, &rbacv1.Role{}, role, semanticRoleIsEqual); err != nil {
		return fmt.Errorf("failed to reconcile multi node role %s/%s: %w", role.GetNamespace(), role.GetName(), err)
	}

	return nil
}

func (r *LLMISVCReconciler) reconcileMultiNodeMainRoleBinding(ctx context.Context, llmSvc *v1alpha2.LLMInferenceService, sa *corev1.ServiceAccount, config *Config) error {
	lws, err := r.expectedMainMultiNodeLWS(ctx, llmSvc, config)
	if err != nil {
		return fmt.Errorf("failed to build the expected main LWS for building the RoleBinding: %w", err)
	}

	roleBinding := r.expectedMultiNodeRoleBinding(llmSvc, sa)
	if !hasRoutingSidecar(lws.Spec.LeaderWorkerTemplate.WorkerTemplate.Spec) && (lws.Spec.LeaderWorkerTemplate.LeaderTemplate == nil || !hasRoutingSidecar(lws.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec)) {
		return Delete(ctx, r, llmSvc, roleBinding)
	}

	if err := Reconcile(ctx, r, llmSvc, &rbacv1.RoleBinding{}, roleBinding, semanticRoleBindingIsEqual); err != nil {
		return fmt.Errorf("failed to reconcile multi node rolebinding %s/%s: %w", roleBinding.GetNamespace(), roleBinding.GetName(), err)
	}

	return nil
}

func (r *LLMISVCReconciler) expectedMultiNodeMainServiceAccount(ctx context.Context, llmSvc *v1alpha2.LLMInferenceService) (*corev1.ServiceAccount, bool, error) {
	useExistingServiceAccount := false
	expectedServiceAccountName := mainLWSName(llmSvc)

	// An existing service account attached to the main leader template takes precedence over any attached to the prefill worker template.
	var existingServiceAccountName string
	if llmSvc.Spec.Template != nil && llmSvc.Spec.Template.ServiceAccountName != "" {
		existingServiceAccountName = llmSvc.Spec.Template.ServiceAccountName
	} else if llmSvc.Spec.Worker != nil && llmSvc.Spec.Worker.ServiceAccountName != "" {
		existingServiceAccountName = llmSvc.Spec.Worker.ServiceAccountName
	}

	if existingServiceAccountName != "" && existingServiceAccountName != expectedServiceAccountName {
		useExistingServiceAccount = true
		log.FromContext(ctx).V(2).Info("Using existing service account for multi node main workload", "serviceAccountName", existingServiceAccountName)
		existingServiceAccount := &corev1.ServiceAccount{}
		err := r.Get(ctx, types.NamespacedName{Name: existingServiceAccountName, Namespace: llmSvc.Namespace}, existingServiceAccount)
		if err != nil {
			return nil, useExistingServiceAccount, fmt.Errorf("failed to fetch existing multi node main service account %s/%s: %w", llmSvc.Namespace, existingServiceAccountName, err)
		}
		return existingServiceAccount, useExistingServiceAccount, nil
	}

	expectedServiceAccount := &corev1.ServiceAccount{
		ObjectMeta: metav1.ObjectMeta{
			Name:      expectedServiceAccountName,
			Namespace: llmSvc.GetNamespace(),
			OwnerReferences: []metav1.OwnerReference{
				*metav1.NewControllerRef(llmSvc, v1alpha2.LLMInferenceServiceGVK),
			},
		},
	}

	r.injectSecretsFromDefaultServiceAccount(ctx, expectedServiceAccount)

	// Add required labels to the created service account
	if expectedServiceAccount.Labels == nil {
		expectedServiceAccount.Labels = make(map[string]string)
	}
	expectedServiceAccount.Labels[constants.KubernetesAppNameLabelKey] = llmSvc.GetName()
	expectedServiceAccount.Labels[constants.KubernetesPartOfLabelKey] = constants.LLMInferenceServicePartOfValue

	return expectedServiceAccount, useExistingServiceAccount, nil
}

func (r *LLMISVCReconciler) expectedMultiNodePrefillServiceAccount(ctx context.Context, llmSvc *v1alpha2.LLMInferenceService) (*corev1.ServiceAccount, bool, error) {
	useExistingServiceAccount := false
	expectedServiceAccountName := prefillLWSName(llmSvc)

	// An existing service account attached to the prefill leader template takes precedence over any attached to the prefill worker template.
	var existingServiceAccountName string
	if llmSvc.Spec.Prefill != nil && llmSvc.Spec.Prefill.Template != nil && llmSvc.Spec.Prefill.Template.ServiceAccountName != "" {
		existingServiceAccountName = llmSvc.Spec.Prefill.Template.ServiceAccountName
	} else if llmSvc.Spec.Prefill != nil && llmSvc.Spec.Prefill.Worker != nil && llmSvc.Spec.Prefill.Worker.ServiceAccountName != "" {
		existingServiceAccountName = llmSvc.Spec.Prefill.Worker.ServiceAccountName
	}

	if existingServiceAccountName != "" && existingServiceAccountName != expectedServiceAccountName {
		useExistingServiceAccount = true
		log.FromContext(ctx).V(2).Info("Using existing service account for multi node prefill workload", "serviceAccountName", existingServiceAccountName)
		existingServiceAccount := &corev1.ServiceAccount{}
		err := r.Get(ctx, types.NamespacedName{Name: existingServiceAccountName, Namespace: llmSvc.Namespace}, existingServiceAccount)
		if err != nil {
			return nil, useExistingServiceAccount, fmt.Errorf("failed to fetch existing multi node prefill service account %s/%s: %w", llmSvc.Namespace, existingServiceAccountName, err)
		}
		return existingServiceAccount, useExistingServiceAccount, nil
	}

	expectedServiceAccount := &corev1.ServiceAccount{
		ObjectMeta: metav1.ObjectMeta{
			Name:      expectedServiceAccountName,
			Namespace: llmSvc.GetNamespace(),
			OwnerReferences: []metav1.OwnerReference{
				*metav1.NewControllerRef(llmSvc, v1alpha2.LLMInferenceServiceGVK),
			},
		},
	}

	r.injectSecretsFromDefaultServiceAccount(ctx, expectedServiceAccount)

	// Add required labels to the created service account
	if expectedServiceAccount.Labels == nil {
		expectedServiceAccount.Labels = make(map[string]string)
	}
	expectedServiceAccount.Labels[constants.KubernetesAppNameLabelKey] = llmSvc.GetName()
	expectedServiceAccount.Labels[constants.KubernetesPartOfLabelKey] = constants.LLMInferenceServicePartOfValue

	return expectedServiceAccount, useExistingServiceAccount, nil
}

func (r *LLMISVCReconciler) expectedMultiNodeMainRole(llmSvc *v1alpha2.LLMInferenceService) *rbacv1.Role {
	ro := &rbacv1.Role{
		ObjectMeta: metav1.ObjectMeta{
			Name:      kmeta.ChildName(llmSvc.GetName(), "-kserve-mn-role"),
			Namespace: llmSvc.GetNamespace(),
			OwnerReferences: []metav1.OwnerReference{
				*metav1.NewControllerRef(llmSvc, v1alpha2.LLMInferenceServiceGVK),
			},
			Labels: map[string]string{
				constants.KubernetesAppNameLabelKey: llmSvc.GetName(),
				constants.KubernetesPartOfLabelKey:  constants.LLMInferenceServicePartOfValue,
			},
		},
	}
	ro.Rules = append(ro.Rules, sidecarSSRFProtectionRules...)
	return ro
}

func (r *LLMISVCReconciler) expectedMultiNodeRoleBinding(llmSvc *v1alpha2.LLMInferenceService, sa *corev1.ServiceAccount) *rbacv1.RoleBinding {
	return &rbacv1.RoleBinding{
		ObjectMeta: metav1.ObjectMeta{
			Name:      kmeta.ChildName(llmSvc.GetName(), "-kserve-mn-rb"),
			Namespace: llmSvc.GetNamespace(),
			OwnerReferences: []metav1.OwnerReference{
				*metav1.NewControllerRef(llmSvc, v1alpha2.LLMInferenceServiceGVK),
			},
			Labels: map[string]string{
				constants.KubernetesAppNameLabelKey: llmSvc.GetName(),
				constants.KubernetesPartOfLabelKey:  constants.LLMInferenceServicePartOfValue,
			},
		},
		Subjects: []rbacv1.Subject{{
			Kind:      "ServiceAccount",
			Name:      sa.GetName(),
			Namespace: sa.GetNamespace(),
		}},
		RoleRef: rbacv1.RoleRef{
			APIGroup: "rbac.authorization.k8s.io",
			Kind:     "Role",
			Name:     kmeta.ChildName(llmSvc.GetName(), "-kserve-mn-role"),
		},
	}
}

func (r *LLMISVCReconciler) propagateTopLevelLeaderWorkerSetMetadata(llmSvc *v1alpha2.LLMInferenceService, expected *lwsapi.LeaderWorkerSet) {
	// Define the prefixes to approve for annotations and labels
	approvedAnnotationPrefixes := []string{
		"leaderworkerset.sigs.k8s.io",
		"k8s.v1.cni.cncf.io",
		constants.KueueAPIGroupName,
		"prometheus.io",
		constants.LocalModelLabel,
	}
	approvedLabelPrefixes := []string{
		constants.KueueAPIGroupName,
		constants.LocalModelLabel,
	}

	// Propagate approved annotations to the LeaderWorkerSet's top-level metadata
	utils.PropagatePrefixedMap(llmSvc.GetAnnotations(), &expected.Annotations, approvedAnnotationPrefixes...)

	if expected.Spec.LeaderWorkerTemplate.LeaderTemplate != nil {
		utils.PropagatePrefixedMap(
			llmSvc.GetAnnotations(),
			&expected.Spec.LeaderWorkerTemplate.LeaderTemplate.Annotations,
			approvedAnnotationPrefixes...,
		)
	}

	utils.PropagatePrefixedMap(
		llmSvc.GetAnnotations(),
		&expected.Spec.LeaderWorkerTemplate.WorkerTemplate.Annotations,
		approvedAnnotationPrefixes...,
	)

	// Propagate approved labels
	utils.PropagatePrefixedMap(llmSvc.GetLabels(), &expected.Labels, approvedLabelPrefixes...)
	if expected.Spec.LeaderWorkerTemplate.LeaderTemplate != nil {
		utils.PropagatePrefixedMap(
			llmSvc.GetLabels(),
			&expected.Spec.LeaderWorkerTemplate.LeaderTemplate.Labels,
			approvedLabelPrefixes...,
		)
	}
	utils.PropagatePrefixedMap(
		llmSvc.GetLabels(),
		&expected.Spec.LeaderWorkerTemplate.WorkerTemplate.Labels,
		approvedLabelPrefixes...,
	)
}

func mainLWSName(llmSvc *v1alpha2.LLMInferenceService) string {
	return kmeta.ChildName(llmSvc.GetName(), "-kserve-mn")
}

func prefillLWSName(llmSvc *v1alpha2.LLMInferenceService) string {
	return kmeta.ChildName(llmSvc.GetName(), "-kserve-mn-prefill")
}

func semanticLWSIsEqual(expected *lwsapi.LeaderWorkerSet, curr *lwsapi.LeaderWorkerSet) bool {
	isLeaderEqual := (expected.Spec.LeaderWorkerTemplate.LeaderTemplate != nil) == (curr.Spec.LeaderWorkerTemplate.LeaderTemplate != nil)

	if expected.Spec.LeaderWorkerTemplate.LeaderTemplate != nil && curr.Spec.LeaderWorkerTemplate.LeaderTemplate != nil {
		// Use DeepEqual for the Pod Spec so that when fields are removed (like resource requirements, we push them down
		// to the child resource)
		isLeaderEqual = equality.Semantic.DeepEqual(
			expected.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec,
			curr.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec,
		)
	}

	// Use DeepEqual for the Pod Spec so that when fields are removed (like resource requirements, we push them down
	// to the child resource)
	isWorkerEqual := equality.Semantic.DeepEqual(expected.Spec.LeaderWorkerTemplate.WorkerTemplate.Spec, curr.Spec.LeaderWorkerTemplate.WorkerTemplate.Spec)

	return isLeaderEqual &&
		isWorkerEqual &&
		equality.Semantic.DeepDerivative(expected.Spec, curr.Spec) &&
		equality.Semantic.DeepDerivative(expected.Labels, curr.Labels) &&
		equality.Semantic.DeepDerivative(expected.Annotations, curr.Annotations)
}