Repository: kubernetes-sigs/kueue

Issue:
MVP support for serving workloads running as LeaderWorkerSet
**What would you like to be added**:

MVP support for LeaderWorkerSet in Kueue. It does not need to be ideal, but we want to have some support to unblock users and collect users' feedback. 

The idea is to base the support on StatefulSets, so the integration would also use Pod Groups, similarly as for regular StatefulSets. Each LeaderWorkerGroup creates a new Pod Group. I a single pod group we will have:

- Leader pod, controller by Leader’s STS
- Worker pods, controller by unique, dedicated STS

The size of the group will be taken from LeaderWorkerSet.Spec.LeaderWorkerTemplate.Size and increased by 1 (to include the leader).

This is a follow up to https://github.com/kubernetes-sigs/kueue/issues/2717.

**Why is this needed**:

We want to support serving primitives in Kueue as there is an increasing demand among users to run clusters mixing AI training and inference who want to manage the expensive GPU resources.

[LeaderWorkerSet](https://github.com/kubernetes-sigs/lws) is a new serving API which is gaining popularity as a primitive to host AI/ML inference. 


Hints from issue comments:
@mimowo: /assign @vladikkuzn @mimowo: /cc @mwielgus @tenzen-y @mimowo: I synced with @mbobrovskyi and @vladikkuzn on the feature and it seems complex, so I propose to have a KEP for it, and go via the Alpha phase so that we can update the implementation in the future easily.@mimowo: The already identified follow ups needed after https://github.com/kubernetes-sigs/kueue/pull/3515:
- support for resizing the LWS group by group  (potentially the same PR, but might be follow up) [already done in the main PR]
- support for the default StartupPolicy=LeaderReady
- support for TAS with rank-based ordering (with e2e tests)
- https://github.com/kubernetes-sigs/kueue/issues/4000@mimowo: /reopen
to finalize the integration with TAS @k8s-ci-robot: @mimowo: Reopened this issue.

<details>

In response to [this](https://github.com/kubernetes-sigs/kueue/issues/3232#issuecomment-2619313252):

>/reopen
>to finalize the integration with TAS 


Instructions for interacting with me using PR comments are available [here](https://git.k8s.io/community/contributors/guide/pull-requests.md).  If you have questions or suggestions related to my behavior, please file an issue against the [kubernetes-sigs/prow](https://github.com/kubernetes-sigs/prow/issues/new?title=Prow%20issue:) repository.
</details>@mimowo: Also, I would like to have e2e test for scaling the LWS when Startup policy is Leader ready (default). @mbobrovskyi: /unassign @vladikkuzn 
/assign@tenzen-y: @mbobrovskyi @mimowo Will you finalize MultiKueue and TAS support within this MVP issue?
@kannon92: What outstanding items do we have for this ticket?@mimowo: I think we can close it, AFAIK the only outstanding work for now is the support of rank-based ordering in TAS for groups.
However, I would  not call it MVP any longer: https://github.com/kubernetes-sigs/kueue/issues/4531@mimowo: /close
Thank you @mbobrovskyi for the implementation!@k8s-ci-robot: @mimowo: Closing this issue.

<details>

In response to [this](https://github.com/kubernetes-sigs/kueue/issues/3232#issuecomment-2709742137):

>/close
>Thank you @mbobrovskyi for the implementation!


Instructions for interacting with me using PR comments are available [here](https://git.k8s.io/community/contributors/guide/pull-requests.md).  If you have questions or suggestions related to my behavior, please file an issue against the [kubernetes-sigs/prow](https://github.com/kubernetes-sigs/prow/issues/new?title=Prow%20issue:) repository.
</details>

Solution patch:
diff --git a/Makefile-test.mk b/Makefile-test.mk
index f51cb33bd..b8f092d4a 100644
--- a/Makefile-test.mk
+++ b/Makefile-test.mk
@@ -155,6 +155,7 @@ run-test-tas-e2e-%: FORCE
 		ARTIFACTS="$(ARTIFACTS)/$@" IMAGE_TAG=$(IMAGE_TAG) GINKGO_ARGS="$(GINKGO_ARGS)" \
 		JOBSET_VERSION=$(JOBSET_VERSION) KUBEFLOW_VERSION=$(KUBEFLOW_VERSION) KUBEFLOW_MPI_VERSION=$(KUBEFLOW_MPI_VERSION) \
 		APPWRAPPER_VERSION=$(APPWRAPPER_VERSION) \
+		LEADERWORKERSET_VERSION=$(LEADERWORKERSET_VERSION) \
 		KIND_CLUSTER_FILE="tas-kind-cluster.yaml" E2E_TARGET_FOLDER="tas" \
 		./hack/e2e-test.sh
 	$(PROJECT_DIR)/bin/ginkgo-top -i $(ARTIFACTS)/$@/e2e.json > $(ARTIFACTS)/$@/e2e-top.yaml
diff --git a/pkg/controller/jobs/leaderworkerset/leaderworkerset_reconciler.go b/pkg/controller/jobs/leaderworkerset/leaderworkerset_reconciler.go
index fa8121cd4..58620a760 100644
--- a/pkg/controller/jobs/leaderworkerset/leaderworkerset_reconciler.go
+++ b/pkg/controller/jobs/leaderworkerset/leaderworkerset_reconciler.go
@@ -166,7 +166,7 @@ func (r *Reconciler) podSets(lws *leaderworkersetv1.LeaderWorkerSet) []kueue.Pod
 		},
 		TopologyRequest: jobframework.PodSetTopologyRequest(
 			&lws.Spec.LeaderWorkerTemplate.WorkerTemplate.ObjectMeta,
-			nil,
+			ptr.To(leaderworkersetv1.WorkerIndexLabelKey),
 			nil,
 			nil,
 		),
diff --git a/pkg/controller/jobs/leaderworkerset/leaderworkerset_webhook.go b/pkg/controller/jobs/leaderworkerset/leaderworkerset_webhook.go
index 51b68d8d1..cdcf77d84 100644
--- a/pkg/controller/jobs/leaderworkerset/leaderworkerset_webhook.go
+++ b/pkg/controller/jobs/leaderworkerset/leaderworkerset_webhook.go
@@ -99,7 +99,9 @@ var (
 	specPath                 = field.NewPath("spec")
 	leaderWorkerTemplatePath = specPath.Child("leaderWorkerTemplate")
 	leaderTemplatePath       = leaderWorkerTemplatePath.Child("leaderTemplate")
+	leaderTemplateMetaPath   = leaderTemplatePath.Child("metadata")
 	workerTemplatePath       = leaderWorkerTemplatePath.Child("workerTemplate")
+	workerTemplateMetaPath   = workerTemplatePath.Child("metadata")
 )
 
 func (wh *Webhook) ValidateCreate(ctx context.Context, obj runtime.Object) (warnings admission.Warnings, err error) {
@@ -108,7 +110,7 @@ func (wh *Webhook) ValidateCreate(ctx context.Context, obj runtime.Object) (warn
 	log := ctrl.LoggerFrom(ctx).WithName("leaderworkerset-webhook")
 	log.V(5).Info("Validating create")
 
-	allErrs := jobframework.ValidateQueueName(lws.Object())
+	allErrs := validateCreate(lws)
 
 	return nil, allErrs.ToAggregate()
 }
@@ -120,11 +122,13 @@ func (wh *Webhook) ValidateUpdate(ctx context.Context, oldObj, newObj runtime.Ob
 	log := ctrl.LoggerFrom(ctx).WithName("leaderworkerset-webhook")
 	log.V(5).Info("Validating update")
 
-	allErrs := apivalidation.ValidateImmutableField(
+	allErrs := validateCreate(newLeaderWorkerSet)
+
+	allErrs = append(allErrs, apivalidation.ValidateImmutableField(
 		jobframework.QueueNameForObject(newLeaderWorkerSet.Object()),
 		jobframework.QueueNameForObject(oldLeaderWorkerSet.Object()),
 		queueNameLabelPath,
-	)
+	)...)
 	allErrs = append(allErrs, jobframework.ValidateUpdateForWorkloadPriorityClassName(
 		newLeaderWorkerSet.Object(),
 		oldLeaderWorkerSet.Object(),
@@ -155,6 +159,22 @@ func GetWorkloadName(uid types.UID, name string, groupIndex string) string {
 	return jobframework.GetWorkloadNameForOwnerWithGVK(fmt.Sprintf("%s-%s", name, groupIndex), uid, gvk)
 }
 
+func validateCreate(lws *LeaderWorkerSet) field.ErrorList {
+	var allErrs field.ErrorList
+	allErrs = append(allErrs, jobframework.ValidateQueueName(lws.Object())...)
+	allErrs = append(allErrs, validateTopologyRequest(lws)...)
+	return allErrs
+}
+
+func validateTopologyRequest(lws *LeaderWorkerSet) field.ErrorList {
+	var allErrs field.ErrorList
+	if lws.Spec.LeaderWorkerTemplate.LeaderTemplate != nil {
+		allErrs = append(allErrs, jobframework.ValidateTASPodSetRequest(leaderTemplateMetaPath, &lws.Spec.LeaderWorkerTemplate.LeaderTemplate.ObjectMeta)...)
+	}
+	allErrs = append(allErrs, jobframework.ValidateTASPodSetRequest(workerTemplateMetaPath, &lws.Spec.LeaderWorkerTemplate.WorkerTemplate.ObjectMeta)...)
+	return allErrs
+}
+
 func validateImmutablePodTemplateSpec(newPodTemplateSpec *corev1.PodTemplateSpec, oldPodTemplateSpec *corev1.PodTemplateSpec, fieldPath *field.Path) field.ErrorList {
 	allErrors := field.ErrorList{}
 	if newPodTemplateSpec == nil || oldPodTemplateSpec == nil {
diff --git a/site/content/en/docs/overview/_index.md b/site/content/en/docs/overview/_index.md
index 920839802..84ce4b83f 100644
--- a/site/content/en/docs/overview/_index.md
+++ b/site/content/en/docs/overview/_index.md
@@ -38,22 +38,22 @@ A core design principle for Kueue is to avoid duplicating mature functionality i
 
 ## Job-integrated features
 
-| Feature                                                                                                         | Batch&nbsp;Job | JobSet | PaddleJob | PytorchJob | TFJob | XGBoostJob | MPIJob | Pod | RayCluster | RayJob | AppWrapper | Deployment | StatefulSet |
-|-----------------------------------------------------------------------------------------------------------------|----------------|--------|-----------|------------|-------|------------|--------|-----|------------|--------|------------|------------|-------------|
-| [Dynamic Reclaim](/docs/concepts/workload/#dynamic-reclaim)                                                     | +              | +      |           |            |       |            |        | +   |            |        |            |            |             |
-| [MultiKueue](/docs/concepts/multikueue/)                                                                        | +              | +      | +         | +          | +     | +          | +      |     | +          | +      | +          |            |             |
-| [MultiKueueBatchJobWithManagedBy](/docs/concepts/multikueue/#multikueuebatchjobwithmanagedby-enabled)           | +              |        |           |            |       |            |        |     |            |        |            |            |             |
-| [PartialAdmission](/docs/tasks/run/jobs/#partial-admission)                                                     | +              |        |           |            |       |            |        |     |            |        |            |            |             |
-| [Workload Priority Class](/docs/concepts/workload_priority_class/)                                              | +              | +      | +         | +          | +     | +          | +      | +   | +          | +      | +          | +          | +           |
-| [FlavorFungibility](/docs/concepts/cluster_queue/#flavorfungibility)                                            | +              | +      | +         | +          | +     | +          | +      | +   | +          | +      | +          | +          | +           |
-| [ProvisioningACC](/docs/admission-check-controllers/provisioning/)                                              | +              | +      | +         | +          | +     | +          | +      | +   | +          | +      | +          | +          | +           |
-| [QueueVisibility](/docs/tasks/manage/monitor_pending_workloads/pending_workloads_in_status/)                    | +              | +      | +         | +          | +     | +          | +      | +   | +          | +      | +          | +          | +           |
-| [VisibilityOnDemand](/docs/tasks/manage/monitor_pending_workloads/pending_workloads_on_demand/)                 | +              | +      | +         | +          | +     | +          | +      | +   | +          | +      | +          | +          | +           |
-| [PrioritySortingWithinCohort](/docs/concepts/cluster_queue/#flavors-and-borrowing-semantics)                    | +              | +      | +         | +          | +     | +          | +      | +   | +          | +      | +          | +          | +           |
-| [LendingLimit](/docs/concepts/cluster_queue/#lendinglimit)                                                      | +              | +      | +         | +          | +     | +          | +      | +   | +          | +      | +          | +          | +           |
-| [All-or-nothing with ready Pods](/docs/concepts/workload/#all-or-nothing-semantics-for-job-resource-assignment) | +              | +      | +         | +          | +     | +          | +      | +   | +          | +      | +          | +          | +           |
-| [Fair Sharing](/docs/concepts/preemption/#fair-sharing)                                                         | +              | +      | +         | +          | +     | +          | +      | +   | +          | +      | +          | +          | +           |
-| [Topology Aware Scheduling](/docs/concepts/topology_aware_scheduling)                                           | +              | +      | +         | +          | +     | +          | +      | +   | +          | +      | +          | +          | +           |
+| Feature                                                                                                         | Batch&nbsp;Job | JobSet | PaddleJob | PytorchJob | TFJob | XGBoostJob | MPIJob | Pod | RayCluster | RayJob | AppWrapper | Deployment | StatefulSet | LeaderWorkerSet |
+|-----------------------------------------------------------------------------------------------------------------|----------------|--------|-----------|------------|-------|------------|--------|-----|------------|--------|------------|------------|-------------|-----------------|
+| [Dynamic Reclaim](/docs/concepts/workload/#dynamic-reclaim)                                                     | +              | +      |           |            |       |            |        | +   |            |        |            |            |             |                 |
+| [MultiKueue](/docs/concepts/multikueue/)                                                                        | +              | +      | +         | +          | +     | +          | +      |     | +          | +      | +          |            |             |                 |
+| [MultiKueueBatchJobWithManagedBy](/docs/concepts/multikueue/#multikueuebatchjobwithmanagedby-enabled)           | +              |        |           |            |       |            |        |     |            |        |            |            |             |                 |
+| [PartialAdmission](/docs/tasks/run/jobs/#partial-admission)                                                     | +              |        |           |            |       |            |        |     |            |        |            |            |             |                 |
+| [Workload Priority Class](/docs/concepts/workload_priority_class/)                                              | +              | +      | +         | +          | +     | +          | +      | +   | +          | +      | +          | +          | +           | +               |
+| [FlavorFungibility](/docs/concepts/cluster_queue/#flavorfungibility)                                            | +              | +      | +         | +          | +     | +          | +      | +   | +          | +      | +          | +          | +           | +               |
+| [ProvisioningACC](/docs/admission-check-controllers/provisioning/)                                              | +              | +      | +         | +          | +     | +          | +      | +   | +          | +      | +          | +          | +           | +               |
+| [QueueVisibility](/docs/tasks/manage/monitor_pending_workloads/pending_workloads_in_status/)                    | +              | +      | +         | +          | +     | +          | +      | +   | +          | +      | +          | +          | +           | +               |
+| [VisibilityOnDemand](/docs/tasks/manage/monitor_pending_workloads/pending_workloads_on_demand/)                 | +              | +      | +         | +          | +     | +          | +      | +   | +          | +      | +          | +          | +           | +               |
+| [PrioritySortingWithinCohort](/docs/concepts/cluster_queue/#flavors-and-borrowing-semantics)                    | +              | +      | +         | +          | +     | +          | +      | +   | +          | +      | +          | +          | +           | +               |
+| [LendingLimit](/docs/concepts/cluster_queue/#lendinglimit)                                                      | +              | +      | +         | +          | +     | +          | +      | +   | +          | +      | +          | +          | +           | +               |
+| [All-or-nothing with ready Pods](/docs/concepts/workload/#all-or-nothing-semantics-for-job-resource-assignment) | +              | +      | +         | +          | +     | +          | +      | +   | +          | +      | +          | +          | +           | +               |
+| [Fair Sharing](/docs/concepts/preemption/#fair-sharing)                                                         | +              | +      | +         | +          | +     | +          | +      | +   | +          | +      | +          | +          | +           | +               |
+| [Topology Aware Scheduling](/docs/concepts/topology_aware_scheduling)                                           | +              | +      | +         | +          | +     | +          | +      | +   | +          | +      | +          | +          | +           | +               |
 
 ## High-level Kueue operation
 



Test patch:
diff --git a/pkg/controller/jobs/leaderworkerset/leaderworkerset_reconciler_test.go b/pkg/controller/jobs/leaderworkerset/leaderworkerset_reconciler_test.go
index 1d16d4728..4ff4518ee 100644
--- a/pkg/controller/jobs/leaderworkerset/leaderworkerset_reconciler_test.go
+++ b/pkg/controller/jobs/leaderworkerset/leaderworkerset_reconciler_test.go
@@ -26,10 +26,12 @@ import (
 	"k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/types"
+	"k8s.io/utils/ptr"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/reconcile"
 	leaderworkersetv1 "sigs.k8s.io/lws/api/leaderworkerset/v1"
 
+	kueuealpha "sigs.k8s.io/kueue/apis/kueue/v1alpha1"
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
 	"sigs.k8s.io/kueue/pkg/controller/jobframework"
 	podcontroller "sigs.k8s.io/kueue/pkg/controller/jobs/pod"
@@ -273,6 +275,113 @@ func TestReconciler(t *testing.T) {
 				},
 			},
 		},
+		"should create prebuild workload with required topology annotation": {
+			leaderWorkerSet: leaderworkerset.MakeLeaderWorkerSet(testLWS, testNS).
+				UID(testUID).
+				Size(3).
+				LeaderTemplate(corev1.PodTemplateSpec{
+					ObjectMeta: metav1.ObjectMeta{
+						Annotations: map[string]string{
+							kueuealpha.PodSetRequiredTopologyAnnotation: "cloud.com/block",
+						},
+					},
+					Spec: corev1.PodSpec{
+						Containers: []corev1.Container{
+							{Name: "c", Image: "pause"},
+						},
+					},
+				}).
+				WorkerTemplate(corev1.PodTemplateSpec{
+					ObjectMeta: metav1.ObjectMeta{
+						Annotations: map[string]string{
+							kueuealpha.PodSetRequiredTopologyAnnotation: "cloud.com/block",
+						},
+					},
+					Spec: corev1.PodSpec{
+						Containers: []corev1.Container{
+							{Name: "c", Image: "pause"},
+						},
+					},
+				}).
+				Obj(),
+			wantLeaderWorkerSet: leaderworkerset.MakeLeaderWorkerSet(testLWS, testNS).
+				UID(testUID).
+				Size(3).
+				LeaderTemplate(corev1.PodTemplateSpec{
+					ObjectMeta: metav1.ObjectMeta{
+						Annotations: map[string]string{
+							kueuealpha.PodSetRequiredTopologyAnnotation: "cloud.com/block",
+						},
+					},
+					Spec: corev1.PodSpec{
+						Containers: []corev1.Container{
+							{Name: "c", Image: "pause"},
+						},
+					},
+				}).
+				WorkerTemplate(corev1.PodTemplateSpec{
+					ObjectMeta: metav1.ObjectMeta{
+						Annotations: map[string]string{
+							kueuealpha.PodSetRequiredTopologyAnnotation: "cloud.com/block",
+						},
+					},
+					Spec: corev1.PodSpec{
+						Containers: []corev1.Container{
+							{Name: "c", Image: "pause"},
+						},
+					},
+				}).
+				Obj(),
+			wantWorkloads: []kueue.Workload{
+				*utiltesting.MakeWorkload(GetWorkloadName(types.UID(testUID), testLWS, "0"), testNS).
+					Annotation(podcontroller.IsGroupWorkloadAnnotationKey, podcontroller.IsGroupWorkloadAnnotationValue).
+					Finalizers(kueue.ResourceInUseFinalizerName).
+					PodSets(
+						kueue.PodSet{
+							Name: leaderPodSetName,
+							Template: corev1.PodTemplateSpec{
+								Spec: corev1.PodSpec{
+									Containers: []corev1.Container{
+										{Name: "c", Image: "pause"},
+									},
+								},
+							},
+							Count: 1,
+							TopologyRequest: &kueue.PodSetTopologyRequest{
+								Required: ptr.To("cloud.com/block"),
+							},
+						},
+						kueue.PodSet{
+							Name: workerPodSetName,
+							Template: corev1.PodTemplateSpec{
+								Spec: corev1.PodSpec{
+									Containers: []corev1.Container{
+										{Name: "c", Image: "pause"},
+									},
+								},
+							},
+							Count: 2,
+							TopologyRequest: &kueue.PodSetTopologyRequest{
+								Required:      ptr.To("cloud.com/block"),
+								PodIndexLabel: ptr.To(leaderworkersetv1.WorkerIndexLabelKey),
+							},
+						},
+					).
+					Obj(),
+			},
+			wantEvents: []utiltesting.EventRecord{
+				{
+					Key:       types.NamespacedName{Name: testLWS, Namespace: testNS},
+					EventType: corev1.EventTypeNormal,
+					Reason:    jobframework.ReasonCreatedWorkload,
+					Message: fmt.Sprintf(
+						"Created Workload: %s/%s",
+						testNS,
+						GetWorkloadName(types.UID(testUID), testLWS, "0"),
+					),
+				},
+			},
+		},
 	}
 	for name, tc := range cases {
 		t.Run(name, func(t *testing.T) {
diff --git a/pkg/controller/jobs/leaderworkerset/leaderworkerset_webhook_test.go b/pkg/controller/jobs/leaderworkerset/leaderworkerset_webhook_test.go
index c50cf2e1c..0f70c7ad3 100644
--- a/pkg/controller/jobs/leaderworkerset/leaderworkerset_webhook_test.go
+++ b/pkg/controller/jobs/leaderworkerset/leaderworkerset_webhook_test.go
@@ -24,10 +24,12 @@ import (
 	"github.com/google/go-cmp/cmp/cmpopts"
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/util/validation/field"
 	"sigs.k8s.io/controller-runtime/pkg/webhook/admission"
 	leaderworkersetv1 "sigs.k8s.io/lws/api/leaderworkerset/v1"
 
+	kueuealpha "sigs.k8s.io/kueue/apis/kueue/v1alpha1"
 	"sigs.k8s.io/kueue/pkg/cache"
 	"sigs.k8s.io/kueue/pkg/controller/constants"
 	"sigs.k8s.io/kueue/pkg/controller/jobframework"
@@ -153,6 +155,58 @@ func TestValidateCreate(t *testing.T) {
 				},
 			}.ToAggregate(),
 		},
+		"valid topology request": {
+			lws: testingleaderworkerset.MakeLeaderWorkerSet("test-lws", "").
+				LeaderTemplate(corev1.PodTemplateSpec{}).
+				Queue("test-queue").
+				LeaderTemplate(corev1.PodTemplateSpec{
+					ObjectMeta: metav1.ObjectMeta{
+						Annotations: map[string]string{
+							kueuealpha.PodSetRequiredTopologyAnnotation: "cloud.com/block",
+						},
+					},
+				}).
+				WorkerTemplate(corev1.PodTemplateSpec{
+					ObjectMeta: metav1.ObjectMeta{
+						Annotations: map[string]string{
+							kueuealpha.PodSetRequiredTopologyAnnotation: "cloud.com/block",
+						},
+					},
+				}).
+				Obj(),
+		},
+		"invalid topology request": {
+			lws: testingleaderworkerset.MakeLeaderWorkerSet("test-lws", "").
+				LeaderTemplate(corev1.PodTemplateSpec{}).
+				Queue("test-queue").
+				LeaderTemplate(corev1.PodTemplateSpec{
+					ObjectMeta: metav1.ObjectMeta{
+						Annotations: map[string]string{
+							kueuealpha.PodSetRequiredTopologyAnnotation:  "cloud.com/block",
+							kueuealpha.PodSetPreferredTopologyAnnotation: "cloud.com/block",
+						},
+					},
+				}).
+				WorkerTemplate(corev1.PodTemplateSpec{
+					ObjectMeta: metav1.ObjectMeta{
+						Annotations: map[string]string{
+							kueuealpha.PodSetRequiredTopologyAnnotation:  "cloud.com/block",
+							kueuealpha.PodSetPreferredTopologyAnnotation: "cloud.com/block",
+						},
+					},
+				}).
+				Obj(),
+			wantErr: field.ErrorList{
+				&field.Error{
+					Type:  field.ErrorTypeInvalid,
+					Field: "spec.leaderWorkerTemplate.leaderTemplate.metadata.annotations",
+				},
+				&field.Error{
+					Type:  field.ErrorTypeInvalid,
+					Field: "spec.leaderWorkerTemplate.workerTemplate.metadata.annotations",
+				},
+			}.ToAggregate(),
+		},
 	}
 
 	for name, tc := range testCases {
@@ -519,6 +573,66 @@ func TestValidateUpdate(t *testing.T) {
 				},
 			}.ToAggregate(),
 		},
+		"set valid topology request": {
+			oldObj: testingleaderworkerset.MakeLeaderWorkerSet("test-lws", "").
+				LeaderTemplate(corev1.PodTemplateSpec{}).
+				WorkerTemplate(corev1.PodTemplateSpec{}).
+				Queue("test-queue").
+				Obj(),
+			newObj: testingleaderworkerset.MakeLeaderWorkerSet("test-lws", "").
+				LeaderTemplate(corev1.PodTemplateSpec{
+					ObjectMeta: metav1.ObjectMeta{
+						Annotations: map[string]string{
+							kueuealpha.PodSetRequiredTopologyAnnotation: "cloud.com/block",
+						},
+					},
+				}).
+				WorkerTemplate(corev1.PodTemplateSpec{
+					ObjectMeta: metav1.ObjectMeta{
+						Annotations: map[string]string{
+							kueuealpha.PodSetRequiredTopologyAnnotation: "cloud.com/block",
+						},
+					},
+				}).
+				Queue("test-queue").
+				Obj(),
+		},
+		"set invalid topology request": {
+			oldObj: testingleaderworkerset.MakeLeaderWorkerSet("test-lws", "").
+				LeaderTemplate(corev1.PodTemplateSpec{}).
+				WorkerTemplate(corev1.PodTemplateSpec{}).
+				Queue("test-queue").
+				Obj(),
+			newObj: testingleaderworkerset.MakeLeaderWorkerSet("test-lws", "").
+				LeaderTemplate(corev1.PodTemplateSpec{
+					ObjectMeta: metav1.ObjectMeta{
+						Annotations: map[string]string{
+							kueuealpha.PodSetRequiredTopologyAnnotation:  "cloud.com/block",
+							kueuealpha.PodSetPreferredTopologyAnnotation: "cloud.com/block",
+						},
+					},
+				}).
+				WorkerTemplate(corev1.PodTemplateSpec{
+					ObjectMeta: metav1.ObjectMeta{
+						Annotations: map[string]string{
+							kueuealpha.PodSetRequiredTopologyAnnotation:  "cloud.com/block",
+							kueuealpha.PodSetPreferredTopologyAnnotation: "cloud.com/block",
+						},
+					},
+				}).
+				Queue("test-queue").
+				Obj(),
+			wantErr: field.ErrorList{
+				&field.Error{
+					Type:  field.ErrorTypeInvalid,
+					Field: "spec.leaderWorkerTemplate.leaderTemplate.metadata.annotations",
+				},
+				&field.Error{
+					Type:  field.ErrorTypeInvalid,
+					Field: "spec.leaderWorkerTemplate.workerTemplate.metadata.annotations",
+				},
+			}.ToAggregate(),
+		},
 	}
 
 	for name, tc := range testCases {
diff --git a/test/e2e/tas/leaderworkerset_test.go b/test/e2e/tas/leaderworkerset_test.go
new file mode 100644
index 000000000..9878422ee
--- /dev/null
+++ b/test/e2e/tas/leaderworkerset_test.go
@@ -0,0 +1,169 @@
+/*
+Copyright 2025 The Kubernetes 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 tase2e
+
+import (
+	"fmt"
+
+	"github.com/onsi/ginkgo/v2"
+	"github.com/onsi/gomega"
+	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/api/resource"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/fields"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+	leaderworkersetv1 "sigs.k8s.io/lws/api/leaderworkerset/v1"
+
+	kueuealpha "sigs.k8s.io/kueue/apis/kueue/v1alpha1"
+	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
+	"sigs.k8s.io/kueue/pkg/util/testing"
+	leaderworkersettesting "sigs.k8s.io/kueue/pkg/util/testingjobs/leaderworkerset"
+	"sigs.k8s.io/kueue/test/util"
+)
+
+var _ = ginkgo.Describe("TopologyAwareScheduling for LeaderWorkerSet", func() {
+	var (
+		ns           *corev1.Namespace
+		topology     *kueuealpha.Topology
+		tasFlavor    *kueue.ResourceFlavor
+		localQueue   *kueue.LocalQueue
+		clusterQueue *kueue.ClusterQueue
+	)
+
+	ginkgo.BeforeEach(func() {
+		ns = &corev1.Namespace{ObjectMeta: metav1.ObjectMeta{GenerateName: "e2e-tas-lws-"}}
+		gomega.Expect(k8sClient.Create(ctx, ns)).To(gomega.Succeed())
+
+		topology = testing.MakeDefaultThreeLevelTopology("datacenter")
+		gomega.Expect(k8sClient.Create(ctx, topology)).Should(gomega.Succeed())
+
+		tasFlavor = testing.MakeResourceFlavor("tas-flavor").
+			NodeLabel(tasNodeGroupLabel, instanceType).
+			TopologyName(topology.Name).
+			Obj()
+		gomega.Expect(k8sClient.Create(ctx, tasFlavor)).Should(gomega.Succeed())
+
+		clusterQueue = testing.MakeClusterQueue("cluster-queue").
+			ResourceGroup(*testing.MakeFlavorQuotas("tas-flavor").Resource(extraResource, "8").Obj()).
+			Obj()
+		gomega.Expect(k8sClient.Create(ctx, clusterQueue)).Should(gomega.Succeed())
+		util.ExpectClusterQueuesToBeActive(ctx, k8sClient, clusterQueue)
+
+		localQueue = testing.MakeLocalQueue("test-queue", ns.Name).ClusterQueue("cluster-queue").Obj()
+		gomega.Expect(k8sClient.Create(ctx, localQueue)).Should(gomega.Succeed())
+		util.ExpectLocalQueuesToBeActive(ctx, k8sClient, localQueue)
+	})
+
+	ginkgo.AfterEach(func() {
+		gomega.Expect(util.DeleteNamespace(ctx, k8sClient, ns)).To(gomega.Succeed())
+		gomega.Expect(util.DeleteAllLeaderWorkerSetsInNamespace(ctx, k8sClient, ns)).To(gomega.Succeed())
+		util.ExpectObjectToBeDeleted(ctx, k8sClient, clusterQueue, true)
+		util.ExpectObjectToBeDeleted(ctx, k8sClient, tasFlavor, true)
+		util.ExpectObjectToBeDeleted(ctx, k8sClient, topology, true)
+		util.ExpectAllPodsInNamespaceDeleted(ctx, k8sClient, ns)
+	})
+
+	ginkgo.When("creating a LeaderWorkerSet", func() {
+		ginkgo.It("should place pods based on the ranks-ordering", func() {
+			const (
+				replicas = int32(2)
+				size     = int32(3)
+			)
+
+			podsTotalCount := replicas * size
+
+			lws := leaderworkersettesting.MakeLeaderWorkerSet("lws", ns.Name).
+				Replicas(replicas).
+				Size(size).
+				Queue(localQueue.Name).
+				WorkerTemplate(corev1.PodTemplateSpec{
+					ObjectMeta: metav1.ObjectMeta{
+						Annotations: map[string]string{
+							kueuealpha.PodSetRequiredTopologyAnnotation: testing.DefaultBlockTopologyLevel,
+						},
+					},
+					Spec: corev1.PodSpec{
+						Containers: []corev1.Container{
+							{
+								Name:  "c",
+								Image: util.E2eTestSleepImage,
+								Args:  []string{"10m"},
+								Resources: corev1.ResourceRequirements{
+									Limits: map[corev1.ResourceName]resource.Quantity{
+										extraResource: resource.MustParse("1"),
+									},
+									Requests: map[corev1.ResourceName]resource.Quantity{
+										extraResource: resource.MustParse("1"),
+									},
+								},
+							},
+						},
+					},
+				}).
+				Obj()
+			ginkgo.By("Creating a LeaderWorkerSet", func() {
+				gomega.Expect(k8sClient.Create(ctx, lws)).To(gomega.Succeed())
+			})
+
+			ginkgo.By("Waiting for replicas is ready", func() {
+				createdLeaderWorkerSet := &leaderworkersetv1.LeaderWorkerSet{}
+				gomega.Eventually(func(g gomega.Gomega) {
+					g.Expect(k8sClient.Get(ctx, client.ObjectKeyFromObject(lws), createdLeaderWorkerSet)).To(gomega.Succeed())
+					g.Expect(createdLeaderWorkerSet.Status.ReadyReplicas).To(gomega.Equal(replicas))
+					g.Expect(createdLeaderWorkerSet.Status.Conditions).To(gomega.ContainElement(
+						gomega.BeComparableTo(metav1.Condition{
+							Type:    "Available",
+							Status:  metav1.ConditionTrue,
+							Reason:  "AllGroupsReady",
+							Message: "All replicas are ready",
+						}, util.IgnoreConditionTimestampsAndObservedGeneration)),
+					)
+				}, util.LongTimeout, util.Interval).Should(gomega.Succeed())
+			})
+
+			pods := &corev1.PodList{}
+			ginkgo.By("ensure all pods are scheduled", func() {
+				listOpts := &client.ListOptions{
+					FieldSelector: fields.OneTermNotEqualSelector("spec.nodeName", ""),
+				}
+				gomega.Eventually(func(g gomega.Gomega) {
+					g.Expect(k8sClient.List(ctx, pods, client.InNamespace(ns.Name), listOpts)).To(gomega.Succeed())
+					g.Expect(pods.Items).Should(gomega.HaveLen(int(podsTotalCount)))
+				}, util.LongTimeout, util.Interval).Should(gomega.Succeed())
+			})
+
+			ginkgo.By("verify the assignment of pods are as expected with rank-based ordering", func() {
+				gomega.Expect(k8sClient.List(ctx, pods, client.InNamespace(ns.Name))).To(gomega.Succeed())
+				gotAssignment := make(map[string]string, podsTotalCount)
+				for _, pod := range pods.Items {
+					index := fmt.Sprintf("%s/%s", pod.Labels[leaderworkersetv1.GroupIndexLabelKey], pod.Labels[leaderworkersetv1.WorkerIndexLabelKey])
+					gotAssignment[index] = pod.Spec.NodeName
+				}
+				wantAssignment := map[string]string{
+					"0/0": "kind-worker",
+					"0/1": "kind-worker2",
+					"0/2": "kind-worker3",
+					"1/0": "kind-worker5",
+					"1/1": "kind-worker6",
+					"1/2": "kind-worker7",
+				}
+				gomega.Expect(gotAssignment).Should(gomega.BeComparableTo(wantAssignment))
+			})
+		},
+		)
+	})
+})