Repository: kubernetes-sigs/kueue

Issue:
TAS: configure PodSetTopologyRequests only when TAS enabled
**What would you like to be cleaned**:

This code is currenlty activated all the time: https://github.com/kubernetes-sigs/kueue/blob/a454f0e133682887c8043b10cc360bc24f7d3924/pkg/controller/jobframework/tas.go#L26

However, this should only be applied in case when TopologyAwareScheduling is enabled.

**Why is this needed**:

There should be no TAS related code active when TAS is disabled.

This also complicates testing the fix in https://github.com/kubernetes-sigs/kueue/issues/4573.

Hints from issue comments:
@mimowo: cc @mbobrovskyi@mykysha: /assign

Solution patch:
diff --git a/pkg/controller/jobs/appwrapper/appwrapper_controller.go b/pkg/controller/jobs/appwrapper/appwrapper_controller.go
index 39b6f7c3a..584bbd5a3 100644
--- a/pkg/controller/jobs/appwrapper/appwrapper_controller.go
+++ b/pkg/controller/jobs/appwrapper/appwrapper_controller.go
@@ -149,10 +149,17 @@ func (aw *AppWrapper) PodSets() ([]kueue.PodSet, error) {
 			}
 		}
 		podSets[psIndex] = kueue.PodSet{
-			Name:            kueue.NewPodSetReference(fmt.Sprintf("%s-%v", aw.Name, psIndex)),
-			Template:        *podSpecTemplates[psIndex],
-			Count:           awutils.Replicas(awPodSets[psIndex]),
-			TopologyRequest: jobframework.PodSetTopologyRequest(&(podSpecTemplates[psIndex].ObjectMeta), podIndexLabel, subGroupIndexLabel, subGroupCount),
+			Name:     kueue.NewPodSetReference(fmt.Sprintf("%s-%v", aw.Name, psIndex)),
+			Template: *podSpecTemplates[psIndex],
+			Count:    awutils.Replicas(awPodSets[psIndex]),
+		}
+		if features.Enabled(features.TopologyAwareScheduling) {
+			podSets[psIndex].TopologyRequest = jobframework.PodSetTopologyRequest(
+				&(podSpecTemplates[psIndex].ObjectMeta),
+				podIndexLabel,
+				subGroupIndexLabel,
+				subGroupCount,
+			)
 		}
 	}
 	return podSets, nil
diff --git a/pkg/controller/jobs/job/job_controller.go b/pkg/controller/jobs/job/job_controller.go
index 0570d026b..1bb430de3 100644
--- a/pkg/controller/jobs/job/job_controller.go
+++ b/pkg/controller/jobs/job/job_controller.go
@@ -249,15 +249,21 @@ func cleanManagedLabels(pt *corev1.PodTemplateSpec) *corev1.PodTemplateSpec {
 }
 
 func (j *Job) PodSets() ([]kueue.PodSet, error) {
+	podSet := kueue.PodSet{
+		Name:     kueue.DefaultPodSetName,
+		Template: *cleanManagedLabels(j.Spec.Template.DeepCopy()),
+		Count:    j.podsCount(),
+		MinCount: j.minPodsCount(),
+	}
+	if features.Enabled(features.TopologyAwareScheduling) {
+		podSet.TopologyRequest = jobframework.PodSetTopologyRequest(
+			&j.Spec.Template.ObjectMeta,
+			ptr.To(batchv1.JobCompletionIndexAnnotation),
+			nil, nil,
+		)
+	}
 	return []kueue.PodSet{
-		{
-			Name:     kueue.DefaultPodSetName,
-			Template: *cleanManagedLabels(j.Spec.Template.DeepCopy()),
-			Count:    j.podsCount(),
-			MinCount: j.minPodsCount(),
-			TopologyRequest: jobframework.PodSetTopologyRequest(&j.Spec.Template.ObjectMeta,
-				ptr.To(batchv1.JobCompletionIndexAnnotation), nil, nil),
-		},
+		podSet,
 	}, nil
 }
 
diff --git a/pkg/controller/jobs/jobset/jobset_controller.go b/pkg/controller/jobs/jobset/jobset_controller.go
index e65393d56..acae2f01a 100644
--- a/pkg/controller/jobs/jobset/jobset_controller.go
+++ b/pkg/controller/jobs/jobset/jobset_controller.go
@@ -124,9 +124,14 @@ func (j *JobSet) PodSets() ([]kueue.PodSet, error) {
 			Name:     kueue.NewPodSetReference(replicatedJob.Name),
 			Template: *replicatedJob.Template.Spec.Template.DeepCopy(),
 			Count:    podsCount(&replicatedJob),
-			TopologyRequest: jobframework.PodSetTopologyRequest(&replicatedJob.Template.Spec.Template.ObjectMeta,
-				ptr.To(batchv1.JobCompletionIndexAnnotation), ptr.To(jobsetapi.JobIndexKey),
-				ptr.To(replicatedJob.Replicas)),
+		}
+		if features.Enabled(features.TopologyAwareScheduling) {
+			podSets[index].TopologyRequest = jobframework.PodSetTopologyRequest(
+				&replicatedJob.Template.Spec.Template.ObjectMeta,
+				ptr.To(batchv1.JobCompletionIndexAnnotation),
+				ptr.To(jobsetapi.JobIndexKey),
+				ptr.To(replicatedJob.Replicas),
+			)
 		}
 	}
 	return podSets, nil
diff --git a/pkg/controller/jobs/kubeflow/kubeflowjob/kubeflowjob_controller.go b/pkg/controller/jobs/kubeflow/kubeflowjob/kubeflowjob_controller.go
index e31c38b5f..6a49860b9 100644
--- a/pkg/controller/jobs/kubeflow/kubeflowjob/kubeflowjob_controller.go
+++ b/pkg/controller/jobs/kubeflow/kubeflowjob/kubeflowjob_controller.go
@@ -106,8 +106,13 @@ func (j *KubeflowJob) PodSets() ([]kueue.PodSet, error) {
 			Name:     kueue.NewPodSetReference(string(replicaType)),
 			Template: *j.KFJobControl.ReplicaSpecs()[replicaType].Template.DeepCopy(),
 			Count:    podsCount(j.KFJobControl.ReplicaSpecs(), replicaType),
-			TopologyRequest: jobframework.PodSetTopologyRequest(&j.KFJobControl.ReplicaSpecs()[replicaType].Template.ObjectMeta,
-				ptr.To(kftraining.ReplicaIndexLabel), nil, nil),
+		}
+		if features.Enabled(features.TopologyAwareScheduling) {
+			podSets[index].TopologyRequest = jobframework.PodSetTopologyRequest(
+				&j.KFJobControl.ReplicaSpecs()[replicaType].Template.ObjectMeta,
+				ptr.To(kftraining.ReplicaIndexLabel),
+				nil, nil,
+			)
 		}
 	}
 	return podSets, nil
diff --git a/pkg/controller/jobs/leaderworkerset/leaderworkerset_reconciler.go b/pkg/controller/jobs/leaderworkerset/leaderworkerset_reconciler.go
index e49168e5c..8bce859b8 100644
--- a/pkg/controller/jobs/leaderworkerset/leaderworkerset_reconciler.go
+++ b/pkg/controller/jobs/leaderworkerset/leaderworkerset_reconciler.go
@@ -33,6 +33,7 @@ import (
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
 	"sigs.k8s.io/kueue/pkg/controller/jobframework"
 	podcontroller "sigs.k8s.io/kueue/pkg/controller/jobs/pod"
+	"sigs.k8s.io/kueue/pkg/features"
 	"sigs.k8s.io/kueue/pkg/workload"
 )
 
@@ -141,19 +142,20 @@ func (r *Reconciler) podSets(lws *leaderworkersetv1.LeaderWorkerSet) []kueue.Pod
 	podSets := make([]kueue.PodSet, 0, 2)
 
 	if lws.Spec.LeaderWorkerTemplate.LeaderTemplate != nil {
-		podSets = append(podSets, kueue.PodSet{
+		podSet := kueue.PodSet{
 			Name:  leaderPodSetName,
 			Count: 1,
 			Template: corev1.PodTemplateSpec{
 				Spec: *lws.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.DeepCopy(),
 			},
-			TopologyRequest: jobframework.PodSetTopologyRequest(
+		}
+		if features.Enabled(features.TopologyAwareScheduling) {
+			podSet.TopologyRequest = jobframework.PodSetTopologyRequest(
 				&lws.Spec.LeaderWorkerTemplate.LeaderTemplate.ObjectMeta,
-				nil,
-				nil,
-				nil,
-			),
-		})
+				nil, nil, nil,
+			)
+		}
+		podSets = append(podSets, podSet)
 	}
 
 	defaultPodSetName := kueue.DefaultPodSetName
@@ -166,19 +168,23 @@ func (r *Reconciler) podSets(lws *leaderworkersetv1.LeaderWorkerSet) []kueue.Pod
 		defaultPodSetCount--
 	}
 
-	podSets = append(podSets, kueue.PodSet{
+	podSet := kueue.PodSet{
 		Name:  defaultPodSetName,
 		Count: defaultPodSetCount,
 		Template: corev1.PodTemplateSpec{
 			Spec: *lws.Spec.LeaderWorkerTemplate.WorkerTemplate.Spec.DeepCopy(),
 		},
-		TopologyRequest: jobframework.PodSetTopologyRequest(
+	}
+
+	if features.Enabled(features.TopologyAwareScheduling) {
+		podSet.TopologyRequest = jobframework.PodSetTopologyRequest(
 			&lws.Spec.LeaderWorkerTemplate.WorkerTemplate.ObjectMeta,
 			ptr.To(leaderworkersetv1.WorkerIndexLabelKey),
-			nil,
-			nil,
-		),
-	})
+			nil, nil,
+		)
+	}
+
+	podSets = append(podSets, podSet)
 
 	return podSets
 }
diff --git a/pkg/controller/jobs/mpijob/mpijob_controller.go b/pkg/controller/jobs/mpijob/mpijob_controller.go
index ca0a634ef..2b0c0f95d 100644
--- a/pkg/controller/jobs/mpijob/mpijob_controller.go
+++ b/pkg/controller/jobs/mpijob/mpijob_controller.go
@@ -121,10 +121,16 @@ func (j *MPIJob) PodSets() ([]kueue.PodSet, error) {
 	podSets := make([]kueue.PodSet, len(replicaTypes))
 	for index, mpiReplicaType := range replicaTypes {
 		podSets[index] = kueue.PodSet{
-			Name:            kueue.NewPodSetReference(string(mpiReplicaType)),
-			Template:        *j.Spec.MPIReplicaSpecs[mpiReplicaType].Template.DeepCopy(),
-			Count:           podsCount(&j.Spec, mpiReplicaType),
-			TopologyRequest: jobframework.PodSetTopologyRequest(&j.Spec.MPIReplicaSpecs[mpiReplicaType].Template.ObjectMeta, ptr.To(kfmpi.ReplicaIndexLabel), nil, nil),
+			Name:     kueue.NewPodSetReference(string(mpiReplicaType)),
+			Template: *j.Spec.MPIReplicaSpecs[mpiReplicaType].Template.DeepCopy(),
+			Count:    podsCount(&j.Spec, mpiReplicaType),
+		}
+		if features.Enabled(features.TopologyAwareScheduling) {
+			podSets[index].TopologyRequest = jobframework.PodSetTopologyRequest(
+				&j.Spec.MPIReplicaSpecs[mpiReplicaType].Template.ObjectMeta,
+				ptr.To(kfmpi.ReplicaIndexLabel),
+				nil, nil,
+			)
 		}
 	}
 	return podSets, nil
diff --git a/pkg/controller/jobs/pod/pod_controller.go b/pkg/controller/jobs/pod/pod_controller.go
index c42a200df..8f8f5cb84 100644
--- a/pkg/controller/jobs/pod/pod_controller.go
+++ b/pkg/controller/jobs/pod/pod_controller.go
@@ -49,6 +49,7 @@ import (
 	"sigs.k8s.io/kueue/pkg/constants"
 	"sigs.k8s.io/kueue/pkg/controller/jobframework"
 	podconstants "sigs.k8s.io/kueue/pkg/controller/jobs/pod/constants"
+	"sigs.k8s.io/kueue/pkg/features"
 	"sigs.k8s.io/kueue/pkg/podset"
 	clientutil "sigs.k8s.io/kueue/pkg/util/client"
 	"sigs.k8s.io/kueue/pkg/util/expectations"
@@ -641,14 +642,21 @@ func constructPodSets(p *corev1.Pod) []kueue.PodSet {
 }
 
 func constructPodSet(p *corev1.Pod) kueue.PodSet {
-	return kueue.PodSet{
+	podSet := kueue.PodSet{
 		Name:  kueue.DefaultPodSetName,
 		Count: 1,
 		Template: corev1.PodTemplateSpec{
 			Spec: *p.Spec.DeepCopy(),
 		},
-		TopologyRequest: jobframework.PodSetTopologyRequest(&p.ObjectMeta, ptr.To(kueuealpha.PodGroupPodIndexLabel), nil, nil),
 	}
+	if features.Enabled(features.TopologyAwareScheduling) {
+		podSet.TopologyRequest = jobframework.PodSetTopologyRequest(
+			&p.ObjectMeta,
+			ptr.To(kueuealpha.PodGroupPodIndexLabel),
+			nil, nil,
+		)
+	}
+	return podSet
 }
 
 func constructGroupPodSetsFast(pods []corev1.Pod, groupTotalCount int) ([]kueue.PodSet, error) {
diff --git a/pkg/controller/jobs/raycluster/raycluster_controller.go b/pkg/controller/jobs/raycluster/raycluster_controller.go
index dbc61b013..5da47add1 100644
--- a/pkg/controller/jobs/raycluster/raycluster_controller.go
+++ b/pkg/controller/jobs/raycluster/raycluster_controller.go
@@ -110,10 +110,16 @@ func (j *RayCluster) PodSets() ([]kueue.PodSet, error) {
 
 	// head
 	podSets[0] = kueue.PodSet{
-		Name:            headGroupPodSetName,
-		Template:        *j.Spec.HeadGroupSpec.Template.DeepCopy(),
-		Count:           1,
-		TopologyRequest: jobframework.PodSetTopologyRequest(&j.Spec.HeadGroupSpec.Template.ObjectMeta, nil, nil, nil),
+		Name:     headGroupPodSetName,
+		Template: *j.Spec.HeadGroupSpec.Template.DeepCopy(),
+		Count:    1,
+	}
+
+	if features.Enabled(features.TopologyAwareScheduling) {
+		podSets[0].TopologyRequest = jobframework.PodSetTopologyRequest(
+			&j.Spec.HeadGroupSpec.Template.ObjectMeta,
+			nil, nil, nil,
+		)
 	}
 
 	// workers
@@ -127,10 +133,15 @@ func (j *RayCluster) PodSets() ([]kueue.PodSet, error) {
 			count *= wgs.NumOfHosts
 		}
 		podSets[index+1] = kueue.PodSet{
-			Name:            kueue.NewPodSetReference(wgs.GroupName),
-			Template:        *wgs.Template.DeepCopy(),
-			Count:           count,
-			TopologyRequest: jobframework.PodSetTopologyRequest(&wgs.Template.ObjectMeta, nil, nil, nil),
+			Name:     kueue.NewPodSetReference(wgs.GroupName),
+			Template: *wgs.Template.DeepCopy(),
+			Count:    count,
+		}
+		if features.Enabled(features.TopologyAwareScheduling) {
+			podSets[index+1].TopologyRequest = jobframework.PodSetTopologyRequest(
+				&wgs.Template.ObjectMeta,
+				nil, nil, nil,
+			)
 		}
 	}
 	return podSets, nil
diff --git a/pkg/controller/jobs/rayjob/rayjob_controller.go b/pkg/controller/jobs/rayjob/rayjob_controller.go
index f29bf452b..6743045ce 100644
--- a/pkg/controller/jobs/rayjob/rayjob_controller.go
+++ b/pkg/controller/jobs/rayjob/rayjob_controller.go
@@ -118,12 +118,18 @@ func (j *RayJob) PodSets() ([]kueue.PodSet, error) {
 	podSets := make([]kueue.PodSet, 0)
 
 	// head
-	podSets = append(podSets, kueue.PodSet{
-		Name:            headGroupPodSetName,
-		Template:        *j.Spec.RayClusterSpec.HeadGroupSpec.Template.DeepCopy(),
-		Count:           1,
-		TopologyRequest: jobframework.PodSetTopologyRequest(&j.Spec.RayClusterSpec.HeadGroupSpec.Template.ObjectMeta, nil, nil, nil),
-	})
+	headPodSet := kueue.PodSet{
+		Name:     headGroupPodSetName,
+		Template: *j.Spec.RayClusterSpec.HeadGroupSpec.Template.DeepCopy(),
+		Count:    1,
+	}
+	if features.Enabled(features.TopologyAwareScheduling) {
+		headPodSet.TopologyRequest = jobframework.PodSetTopologyRequest(
+			&j.Spec.RayClusterSpec.HeadGroupSpec.Template.ObjectMeta,
+			nil, nil, nil,
+		)
+	}
+	podSets = append(podSets, headPodSet)
 
 	// workers
 	for index := range j.Spec.RayClusterSpec.WorkerGroupSpecs {
@@ -135,12 +141,15 @@ func (j *RayJob) PodSets() ([]kueue.PodSet, error) {
 		if wgs.NumOfHosts > 1 {
 			count *= wgs.NumOfHosts
 		}
-		podSets = append(podSets, kueue.PodSet{
-			Name:            kueue.NewPodSetReference(wgs.GroupName),
-			Template:        *wgs.Template.DeepCopy(),
-			Count:           count,
-			TopologyRequest: jobframework.PodSetTopologyRequest(&wgs.Template.ObjectMeta, nil, nil, nil),
-		})
+		workerPodSet := kueue.PodSet{
+			Name:     kueue.NewPodSetReference(wgs.GroupName),
+			Template: *wgs.Template.DeepCopy(),
+			Count:    count,
+		}
+		if features.Enabled(features.TopologyAwareScheduling) {
+			workerPodSet.TopologyRequest = jobframework.PodSetTopologyRequest(&wgs.Template.ObjectMeta, nil, nil, nil)
+		}
+		podSets = append(podSets, workerPodSet)
 	}
 
 	// submitter Job
@@ -153,7 +162,9 @@ func (j *RayJob) PodSets() ([]kueue.PodSet, error) {
 
 		// Create the TopologyRequest for the Submitter Job PodSet, based on the annotations
 		// in rayJob.Spec.SubmitterPodTemplate, which can be specified by the user.
-		submitterJobPodSet.TopologyRequest = jobframework.PodSetTopologyRequest(&submitterJobPodSet.Template.ObjectMeta, nil, nil, nil)
+		if features.Enabled(features.TopologyAwareScheduling) {
+			submitterJobPodSet.TopologyRequest = jobframework.PodSetTopologyRequest(&submitterJobPodSet.Template.ObjectMeta, nil, nil, nil)
+		}
 		podSets = append(podSets, submitterJobPodSet)
 	}
 
diff --git a/site/content/en/docs/concepts/resource_flavor.md b/site/content/en/docs/concepts/resource_flavor.md
index 1c07fef22..933744b60 100644
--- a/site/content/en/docs/concepts/resource_flavor.md
+++ b/site/content/en/docs/concepts/resource_flavor.md
@@ -3,110 +3,90 @@ title: "Resource Flavor"
 date: 2022-03-14
 weight: 2
 description: >
-  An object that you can define to describe what resources are available
-  in a cluster.
+  An object that defines available compute resources in a cluster and enables fine-grained resource management by associating workloads with specific node types.
 ---
 
 Resources in a cluster are typically not homogeneous. Resources could differ in:
 
-- pricing and availability (for example, spot versus on-demand VMs)
-- architecture (for example, x86 versus ARM CPUs)
-- brands and models (for example, Radeon 7000 versus Nvidia A100 versus T4 GPUs)
+- Pricing and availability (for example, spot versus on-demand VMs)
+- Architecture (for example, x86 versus ARM CPUs)
+- Brands and models (for example, Radeon 7000 versus Nvidia A100 versus T4 GPUs)
 
-A ResourceFlavor is an object that represents these resource variations and
-allows you to associate them with cluster nodes through labels, taints and tolerations.
+A ResourceFlavor is an object that represents these resource variations and allows you to associate them with cluster nodes through labels, taints and tolerations.
 
 {{% alert title="Note" color="primary" %}}
-If the resources in your cluster are homogeneous, you can use 
-an [empty ResourceFlavor](#empty-resourceflavor) instead of adding labels to custom ResourceFlavors.
+If the resources in your cluster are homogeneous, you can use an [empty ResourceFlavor](#empty-resourceflavor) instead of adding labels to custom ResourceFlavors.
 {{% /alert %}}
 
-A sample ResourceFlavor looks like the following:
-
-```yaml
-apiVersion: kueue.x-k8s.io/v1beta1
-kind: ResourceFlavor
-metadata:
-  name: "spot"
-spec:
-  nodeLabels:
-    instance-type: spot
-  nodeTaints:
-  - effect: NoSchedule
-    key: spot
-    value: "true"
-  tolerations:
-  - key: "spot-taint"
-    operator: "Exists"
-    effect: "NoSchedule"
-```
-
-You can use the `.metadata.name` field to reference a ResourceFlavor from a
-[ClusterQueue](/docs/concepts/cluster_queue) in the `.spec.resourceGroups[*].flavors[*].name` field.
-
-## ResourceFlavor labels
+## ResourceFlavor tolerations for automatic scheduling
 
 **Requires Kubernetes 1.23 or newer**
 
-To associate a ResourceFlavor with a subset of nodes of your cluster, you can
-configure the `.spec.nodeLabels` field with matching node labels that uniquely identify
-the nodes. If you are using [cluster autoscaler](https://github.com/kubernetes/autoscaler/tree/master/cluster-autoscaler)
-(or equivalent controllers), make sure that the controller is configured to add those labels when
-adding new nodes.
+This approach may be best for teams that wish to schedule pods onto the appropriate nodes automatically.
+However, a limitation arises when multiple types of specialized hardware are present, such as two different nvidia.com/gpu resources present in the cluster, i.e., T4 and A100 GPUs.
+The system may not differentiate between them, meaning the pods could be scheduled on any of both types of hardware.
+
+To associate a ResourceFlavor with a subset of nodes of your cluster, you can configure the `.spec.nodeLabels` field with matching node labels that uniquely identify the nodes.
+If you are using [cluster autoscaler](https://github.com/kubernetes/autoscaler/tree/master/cluster-autoscaler) (or equivalent controllers), make sure that the controller is configured to add those labels when adding new nodes.
 
-To guarantee that the Pods in the [Workload](/docs/concepts/workload) run on the nodes associated to the flavor
-that Kueue selected, Kueue performs the following steps:
+To guarantee that the Pods in the [Workload](/docs/concepts/workload) run on the nodes associated to the flavor that Kueue selected, Kueue performs the following steps:
 
-1. When admitting a Workload, Kueue evaluates the
-   [`.nodeSelector`](https://kubernetes.io/docs/concepts/scheduling-eviction/assign-pod-node/#nodeselector)
-   and [`.affinity.nodeAffinity.requiredDuringSchedulingIgnoredDuringExecution`](https://kubernetes.io/docs/concepts/scheduling-eviction/assign-pod-node/#node-affinity)
-   fields in the PodSpecs of your [Workload](/docs/concepts/workload) against the
-   ResourceFlavor labels.
-   `ResourceFlavors` that don't match the node affinity of the Workload
-   cannot be assigned to a Workload's podSet.
+1. When admitting a Workload, Kueue evaluates the [`.nodeSelector`](https://kubernetes.io/docs/concepts/scheduling-eviction/assign-pod-node/#nodeselector) and [`.affinity.nodeAffinity.requiredDuringSchedulingIgnoredDuringExecution`](https://kubernetes.io/docs/concepts/scheduling-eviction/assign-pod-node/#node-affinity) fields in the PodSpecs of your [Workload](/docs/concepts/workload) against the ResourceFlavor labels.
+`ResourceFlavors` that don't match the node affinity of the Workload cannot be assigned to a Workload's podSet.
 
 
 2. Once the Workload is admitted:
-   - Kueue adds the ResourceFlavor labels to the
-    `.nodeSelector` of the underlying Workload Pod templates. This occurs if the Workload
-     didn't specify the `ResourceFlavor` labels already as part of its nodeSelector.
-
-     For example, for a [batch/v1.Job](https://kubernetes.io/docs/concepts/workloads/controllers/job/),
-     Kueue adds the labels to the `.spec.template.spec.nodeSelector` field. This
-     guarantees that the workload's Pods can only be scheduled on the nodes
-     targeted by the flavor that Kueue assigned to the Workload.
-   
+   - Kueue adds the ResourceFlavor labels to the `.nodeSelector` of the underlying Workload Pod templates.
+   This occurs if the Workload didn't specify the `ResourceFlavor` labels already as part of its nodeSelector.
+
+     For example, for a [batch/v1.Job](https://kubernetes.io/docs/concepts/workloads/controllers/job/), Kueue adds the labels to the `.spec.template.spec.nodeSelector` field.
+     This guarantees that the Workload's Pods can only be scheduled on the nodes targeted by the flavor that Kueue assigned to the Workload.
+
    - Kueue adds the tolerations to the underlying Workload Pod templates.
 
-     For example, for a [batch/v1.Job](https://kubernetes.io/docs/concepts/workloads/controllers/job/),
-     Kueue adds the tolerations to the `.spec.template.spec.tolerations` field. This allows that the 
-     workloads Pods to be scheduled on nodes having specific taints.
+     For example, for a [batch/v1.Job](https://kubernetes.io/docs/concepts/workloads/controllers/job/), Kueue adds the tolerations to the `.spec.template.spec.tolerations` field.
+     This allows the Workload's Pods to be scheduled on nodes having specific taints.
 
-## ResourceFlavor taints
+A sample ResourceFlavor of this type looks like the following:
 
-To restrict the usage of a ResourceFlavor, you can configure the `.spec.nodeTaints` field.
-These taints should typically match the taints of the Nodes associated with the ResourceFlavor.
+{{< include "examples/admin/resource-flavor-tolerations.yaml" "yaml" >}}
+
+When defining a ResourceFlavor as above, you should set the following values:
+- The `.metadata.name` field, which is required to reference a ResourceFlavor from a [ClusterQueue](/docs/concepts/cluster_queue) in the `.spec.resourceGroups[*].flavors[*].name` field.
+- `spec.nodeLabels` associates the ResourceFlavor with a node or subset of nodes.
+- `spec.tolerations` adds the specified tolerations to the pods that require GPUs.
+
+
+## ResourceFlavor taints for user-selective scheduling
+
+This approach may be best for teams that wish to schedule their Workload to a specific hardware type selectively.
+An additional ResourceFlavor can be created per type of special hardware with a different set of taints and tolerations.
+The user can then add the tolerations to their Workload to schedule the pods onto the appropriate node.
+
+By adding the taint at the ResourceFlavor level, we ensure that only workloads that explicitly tolerate that taint can consume the quota.
 
 Taints on the ResourceFlavor work similarly to [Node taints](https://kubernetes.io/docs/concepts/scheduling-eviction/taint-and-toleration/),
-although supported `effect` is only `NoExecute` and `NoSchedule`.
-For Kueue to [admit](/docs/concepts#admission) a Workload to use the ResourceFlavor, the PodSpecs in the
-Workload should have a toleration for it. As opposed to the behavior for
-[ResourceFlavor labels](#resourceflavor-labels), Kueue does not add tolerations
-for the flavor taints.
+but only support the `NoExecute` and `NoSchedule` effects, while `PreferNoSchedule` is ignored.
+
+For Kueue to [admit](/docs/concepts#admission) a Workload to use the ResourceFlavor, the PodSpecs in the Workload should have a toleration for it.
+As opposed to the behavior for [ResourceFlavor tolerations for automatic scheduling](#ResourceFlavor-tolerations-for-automatic-scheduling), Kueue does not add tolerations for the flavor taints.
+
+A sample ResourceFlavor looks like the following:
+
+{{< include "examples/admin/resource-flavor-taints.yaml" "yaml" >}}
+
+When defining a ResourceFlavor as above, you should set the following values:
+- The `.metadata.name` field, which is required to reference a ResourceFlavor from a [ClusterQueue](/docs/concepts/cluster_queue) in the `.spec.resourceGroups[*].flavors[*].name` field.
+- `spec.nodeLabels` associates the ResourceFlavor with a node or subset of nodes.
+- `spec.nodeTaints` restricts usage of a ResourceFlavor.
+These taints should typically match the taints of the Nodes associated with the ResourceFlavor.
 
 ## Empty ResourceFlavor
 
-If your cluster has homogeneous resources, or if you don't need to manage
-quotas for the different flavors of a resource separately, you can create a
-ResourceFlavor without any labels or taints. Such ResourceFlavor is called an
-empty ResourceFlavor and its definition looks like the following:
-
-```yaml
-apiVersion: kueue.x-k8s.io/v1beta1
-kind: ResourceFlavor
-metadata:
-  name: default-flavor
-```
+If your cluster has homogeneous resources, or if you don't need to manage quotas for the different flavors of a resource separately, you can create a ResourceFlavor without any labels or taints.
+Such ResourceFlavor is called an empty ResourceFlavor and its definition looks like the following:
+
+{{< include "examples/admin/resource-flavor-empty.yaml" "yaml" >}}
 
 ## What's next?
 
diff --git a/site/static/examples/admin/resource-flavor-empty.yaml b/site/static/examples/admin/resource-flavor-empty.yaml
new file mode 100644
index 000000000..6cbccf3b7
--- /dev/null
+++ b/site/static/examples/admin/resource-flavor-empty.yaml
@@ -0,0 +1,4 @@
+apiVersion: kueue.x-k8s.io/v1beta1
+kind: ResourceFlavor
+metadata:
+  name: default-flavor
diff --git a/site/static/examples/admin/resource-flavor-taints.yaml b/site/static/examples/admin/resource-flavor-taints.yaml
new file mode 100644
index 000000000..f04d54e50
--- /dev/null
+++ b/site/static/examples/admin/resource-flavor-taints.yaml
@@ -0,0 +1,11 @@
+apiVersion: kueue.x-k8s.io/v1beta1
+kind: ResourceFlavor
+metadata:
+  name: "spot"
+spec:
+  nodeLabels:
+    instance-type: spot
+  nodeTaints:
+  - effect: NoSchedule ## Supported effects are NoSchedule and NoExecute, while PreferNoSchedule is ignored.
+    key: spot
+    value: "true"
diff --git a/site/static/examples/admin/resource-flavor-tolerations.yaml b/site/static/examples/admin/resource-flavor-tolerations.yaml
new file mode 100644
index 000000000..564cce5ef
--- /dev/null
+++ b/site/static/examples/admin/resource-flavor-tolerations.yaml
@@ -0,0 +1,11 @@
+apiVersion: kueue.x-k8s.io/v1beta1
+kind: ResourceFlavor
+metadata:
+  name: "spot"
+spec:
+  nodeLabels:
+    instance-type: spot
+  tolerations:
+  - key: "spot-taint" ## The key of the node taint.
+    operator: "Exists"
+    effect: "NoSchedule" ## Supported effects are NoSchedule, NoExecute, and PreferNoSchedule.



Test patch:
diff --git a/pkg/controller/jobs/appwrapper/appwrapper_controller_test.go b/pkg/controller/jobs/appwrapper/appwrapper_controller_test.go
index 25c90147c..93741f363 100644
--- a/pkg/controller/jobs/appwrapper/appwrapper_controller_test.go
+++ b/pkg/controller/jobs/appwrapper/appwrapper_controller_test.go
@@ -36,6 +36,7 @@ import (
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
 	controllerconsts "sigs.k8s.io/kueue/pkg/controller/constants"
 	"sigs.k8s.io/kueue/pkg/controller/jobframework"
+	"sigs.k8s.io/kueue/pkg/features"
 	utiltesting "sigs.k8s.io/kueue/pkg/util/testing"
 	testingappwrapper "sigs.k8s.io/kueue/pkg/util/testingjobs/appwrapper"
 	utiltestingjob "sigs.k8s.io/kueue/pkg/util/testingjobs/job"
@@ -88,8 +89,9 @@ func TestPodSets(t *testing.T) {
 		Obj()
 
 	testCases := map[string]struct {
-		job         *awv1beta2.AppWrapper
-		wantPodSets []kueue.PodSet
+		job                           *awv1beta2.AppWrapper
+		wantPodSets                   []kueue.PodSet
+		enableTopologyAwareScheduling bool
 	}{
 		"no annotations": {
 			job: testingappwrapper.MakeAppWrapper("aw", "ns").
@@ -107,6 +109,7 @@ func TestPodSets(t *testing.T) {
 					PodSpec(batchJob.Spec.Template.Spec).
 					Obj(),
 			},
+			enableTopologyAwareScheduling: false,
 		},
 		"with required and preferred topology annotation": {
 			job: testingappwrapper.MakeAppWrapper("aw", "ns").
@@ -127,11 +130,32 @@ func TestPodSets(t *testing.T) {
 					PodIndexLabel(ptr.To(kftraining.ReplicaIndexLabel)).
 					Obj(),
 			},
+
+			enableTopologyAwareScheduling: true,
+		},
+		"without required and preferred topology annotation if TAS is disabled": {
+			job: testingappwrapper.MakeAppWrapper("aw", "ns").
+				Component(pytorchJobTAS).
+				Obj(),
+
+			wantPodSets: []kueue.PodSet{
+				*utiltesting.MakePodSet("aw-0", 1).
+					PodSpec(pytorchJobTAS.Spec.PyTorchReplicaSpecs[kftraining.PyTorchJobReplicaTypeMaster].Template.Spec).
+					Annotations(map[string]string{kueuealpha.PodSetRequiredTopologyAnnotation: "cloud.com/rack"}).
+					Obj(),
+				*utiltesting.MakePodSet("aw-1", 4).
+					PodSpec(pytorchJobTAS.Spec.PyTorchReplicaSpecs[kftraining.PyTorchJobReplicaTypeWorker].Template.Spec).
+					Annotations(map[string]string{kueuealpha.PodSetPreferredTopologyAnnotation: "cloud.com/block"}).
+					Obj(),
+			},
+
+			enableTopologyAwareScheduling: false,
 		},
 	}
 
 	for name, tc := range testCases {
 		t.Run(name, func(t *testing.T) {
+			features.SetFeatureGateDuringTest(t, features.TopologyAwareScheduling, tc.enableTopologyAwareScheduling)
 			gotPodSets, err := fromObject(tc.job).PodSets()
 			if err != nil {
 				t.Fatalf("failed to get pod sets: %v", err)
diff --git a/pkg/controller/jobs/job/job_controller_test.go b/pkg/controller/jobs/job/job_controller_test.go
index b5d4a38e0..088687113 100644
--- a/pkg/controller/jobs/job/job_controller_test.go
+++ b/pkg/controller/jobs/job/job_controller_test.go
@@ -330,8 +330,9 @@ func TestPodSets(t *testing.T) {
 	jobTemplate := utiltestingjob.MakeJob("job", "ns")
 
 	cases := map[string]struct {
-		job         *Job
-		wantPodSets []kueue.PodSet
+		job                           *Job
+		wantPodSets                   []kueue.PodSet
+		enableTopologyAwareScheduling bool
 	}{
 		"no partial admission": {
 			job: (*Job)(jobTemplate.Clone().Parallelism(3).Obj()),
@@ -340,6 +341,7 @@ func TestPodSets(t *testing.T) {
 					PodSpec(*jobTemplate.Clone().Spec.Template.Spec.DeepCopy()).
 					Obj(),
 			},
+			enableTopologyAwareScheduling: false,
 		},
 		"partial admission": {
 			job: (*Job)(
@@ -354,6 +356,7 @@ func TestPodSets(t *testing.T) {
 					SetMinimumCount(2).
 					Obj(),
 			},
+			enableTopologyAwareScheduling: false,
 		},
 		"with required topology annotation": {
 			job: (*Job)(
@@ -370,6 +373,7 @@ func TestPodSets(t *testing.T) {
 					PodIndexLabel(ptr.To(batchv1.JobCompletionIndexAnnotation)).
 					Obj(),
 			},
+			enableTopologyAwareScheduling: true,
 		},
 		"with preferred topology annotation": {
 			job: (*Job)(
@@ -386,10 +390,42 @@ func TestPodSets(t *testing.T) {
 					PodIndexLabel(ptr.To(batchv1.JobCompletionIndexAnnotation)).
 					Obj(),
 			},
+			enableTopologyAwareScheduling: true,
+		},
+		"without preferred topology annotation if TAS is disabled": {
+			job: (*Job)(
+				jobTemplate.Clone().
+					Parallelism(3).
+					PodAnnotation(kueuealpha.PodSetPreferredTopologyAnnotation, "cloud.com/block").
+					Obj(),
+			),
+			wantPodSets: []kueue.PodSet{
+				*utiltesting.MakePodSet(kueue.DefaultPodSetName, 3).
+					PodSpec(jobTemplate.Clone().Spec.Template.Spec).
+					Annotations(map[string]string{kueuealpha.PodSetPreferredTopologyAnnotation: "cloud.com/block"}).
+					Obj(),
+			},
+			enableTopologyAwareScheduling: false,
+		},
+		"without required topology annotation if TAS is disabled": {
+			job: (*Job)(
+				jobTemplate.Clone().
+					Parallelism(3).
+					PodAnnotation(kueuealpha.PodSetRequiredTopologyAnnotation, "cloud.com/block").
+					Obj(),
+			),
+			wantPodSets: []kueue.PodSet{
+				*utiltesting.MakePodSet(kueue.DefaultPodSetName, 3).
+					PodSpec(jobTemplate.Clone().Spec.Template.Spec).
+					Annotations(map[string]string{kueuealpha.PodSetRequiredTopologyAnnotation: "cloud.com/block"}).
+					Obj(),
+			},
+			enableTopologyAwareScheduling: false,
 		},
 	}
 	for name, tc := range cases {
 		t.Run(name, func(t *testing.T) {
+			features.SetFeatureGateDuringTest(t, features.TopologyAwareScheduling, tc.enableTopologyAwareScheduling)
 			gotPodSets, err := tc.job.PodSets()
 			if err != nil {
 				t.Fatalf("unexpected error: %v", err)
diff --git a/pkg/controller/jobs/jobset/jobset_controller_test.go b/pkg/controller/jobs/jobset/jobset_controller_test.go
index f352c38ff..9f8835846 100644
--- a/pkg/controller/jobs/jobset/jobset_controller_test.go
+++ b/pkg/controller/jobs/jobset/jobset_controller_test.go
@@ -36,6 +36,7 @@ import (
 	"sigs.k8s.io/kueue/pkg/constants"
 	controllerconsts "sigs.k8s.io/kueue/pkg/controller/constants"
 	"sigs.k8s.io/kueue/pkg/controller/jobframework"
+	"sigs.k8s.io/kueue/pkg/features"
 	utiltesting "sigs.k8s.io/kueue/pkg/util/testing"
 	testingjobset "sigs.k8s.io/kueue/pkg/util/testingjobs/jobset"
 )
@@ -204,8 +205,9 @@ func TestPodSets(t *testing.T) {
 	jobSetTemplate := testingjobset.MakeJobSet("jobset", "ns")
 
 	testCases := map[string]struct {
-		jobSet      *JobSet
-		wantPodSets func(jobSet *JobSet) []kueue.PodSet
+		jobSet                        *JobSet
+		wantPodSets                   func(jobSet *JobSet) []kueue.PodSet
+		enableTopologyAwareScheduling bool
 	}{
 		"no annotations": {
 			jobSet: (*JobSet)(jobSetTemplate.Clone().
@@ -224,6 +226,7 @@ func TestPodSets(t *testing.T) {
 						Obj(),
 				}
 			},
+			enableTopologyAwareScheduling: false,
 		},
 		"with required topology annotation": {
 			jobSet: (*JobSet)(jobSetTemplate.Clone().
@@ -255,6 +258,7 @@ func TestPodSets(t *testing.T) {
 						Obj(),
 				}
 			},
+			enableTopologyAwareScheduling: true,
 		},
 		"with preferred topology annotation": {
 			jobSet: (*JobSet)(jobSetTemplate.Clone().
@@ -286,10 +290,53 @@ func TestPodSets(t *testing.T) {
 						Obj(),
 				}
 			},
+			enableTopologyAwareScheduling: true,
+		},
+		"without required and preferred topology annotation if TAS is disabled": {
+			jobSet: (*JobSet)(jobSetTemplate.Clone().
+				ReplicatedJobs(
+					testingjobset.ReplicatedJobRequirements{Name: "job1", Replicas: 2, Parallelism: 1, Completions: 1},
+					testingjobset.ReplicatedJobRequirements{
+						Name:        "job2",
+						Replicas:    2,
+						Parallelism: 1,
+						Completions: 1,
+						PodAnnotations: map[string]string{
+							kueuealpha.PodSetRequiredTopologyAnnotation: "cloud.com/block",
+						},
+					},
+					testingjobset.ReplicatedJobRequirements{
+						Name:        "job3",
+						Replicas:    3,
+						Parallelism: 2,
+						Completions: 3,
+						PodAnnotations: map[string]string{
+							kueuealpha.PodSetPreferredTopologyAnnotation: "cloud.com/block",
+						},
+					},
+				).
+				Obj()),
+			wantPodSets: func(jobSet *JobSet) []kueue.PodSet {
+				return []kueue.PodSet{
+					*utiltesting.MakePodSet(kueue.NewPodSetReference(jobSet.Spec.ReplicatedJobs[0].Name), 2).
+						PodSpec(*jobSet.Spec.ReplicatedJobs[0].Template.Spec.Template.Spec.DeepCopy()).
+						Obj(),
+					*utiltesting.MakePodSet(kueue.NewPodSetReference(jobSet.Spec.ReplicatedJobs[1].Name), 2).
+						PodSpec(*jobSet.Spec.ReplicatedJobs[1].Template.Spec.Template.Spec.DeepCopy()).
+						Annotations(map[string]string{kueuealpha.PodSetRequiredTopologyAnnotation: "cloud.com/block"}).
+						Obj(),
+					*utiltesting.MakePodSet(kueue.NewPodSetReference(jobSet.Spec.ReplicatedJobs[2].Name), 6).
+						PodSpec(*jobSet.Spec.ReplicatedJobs[2].Template.Spec.Template.Spec.DeepCopy()).
+						Annotations(map[string]string{kueuealpha.PodSetPreferredTopologyAnnotation: "cloud.com/block"}).
+						Obj(),
+				}
+			},
+			enableTopologyAwareScheduling: false,
 		},
 	}
 	for name, tc := range testCases {
 		t.Run(name, func(t *testing.T) {
+			features.SetFeatureGateDuringTest(t, features.TopologyAwareScheduling, tc.enableTopologyAwareScheduling)
 			gotPodSets, err := tc.jobSet.PodSets()
 			if err != nil {
 				t.Fatalf("unexpected error: %s", err)
diff --git a/pkg/controller/jobs/kubeflow/jobs/paddlejob/paddlejob_controller_test.go b/pkg/controller/jobs/kubeflow/jobs/paddlejob/paddlejob_controller_test.go
index 537a2a2d0..25cb7f73a 100644
--- a/pkg/controller/jobs/kubeflow/jobs/paddlejob/paddlejob_controller_test.go
+++ b/pkg/controller/jobs/kubeflow/jobs/paddlejob/paddlejob_controller_test.go
@@ -35,6 +35,7 @@ import (
 	kueuealpha "sigs.k8s.io/kueue/apis/kueue/v1alpha1"
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
 	"sigs.k8s.io/kueue/pkg/controller/jobframework"
+	"sigs.k8s.io/kueue/pkg/features"
 	utiltesting "sigs.k8s.io/kueue/pkg/util/testing"
 	testingpaddlejob "sigs.k8s.io/kueue/pkg/util/testingjobs/paddlejob"
 )
@@ -233,8 +234,9 @@ func TestOrderedReplicaTypes(t *testing.T) {
 
 func TestPodSets(t *testing.T) {
 	testCases := map[string]struct {
-		job         *kftraining.PaddleJob
-		wantPodSets func(job *kftraining.PaddleJob) []kueue.PodSet
+		job                           *kftraining.PaddleJob
+		wantPodSets                   func(job *kftraining.PaddleJob) []kueue.PodSet
+		enableTopologyAwareScheduling bool
 	}{
 		"no annotations": {
 			job: testingpaddlejob.MakePaddleJob("paddlejob", "ns").
@@ -259,6 +261,7 @@ func TestPodSets(t *testing.T) {
 						Obj(),
 				}
 			},
+			enableTopologyAwareScheduling: false,
 		},
 		"with required and preferred topology annotation": {
 			job: testingpaddlejob.MakePaddleJob("paddlejob", "ns").
@@ -295,10 +298,45 @@ func TestPodSets(t *testing.T) {
 						Obj(),
 				}
 			},
+			enableTopologyAwareScheduling: true,
+		},
+		"without required and preferred topology annotation if TAS is disabled": {
+			job: testingpaddlejob.MakePaddleJob("paddlejob", "ns").
+				PaddleReplicaSpecs(
+					testingpaddlejob.PaddleReplicaSpecRequirement{
+						ReplicaType:  kftraining.PaddleJobReplicaTypeMaster,
+						ReplicaCount: 1,
+						Annotations: map[string]string{
+							kueuealpha.PodSetRequiredTopologyAnnotation: "cloud.com/rack",
+						},
+					},
+					testingpaddlejob.PaddleReplicaSpecRequirement{
+						ReplicaType:  kftraining.PaddleJobReplicaTypeWorker,
+						ReplicaCount: 1,
+						Annotations: map[string]string{
+							kueuealpha.PodSetPreferredTopologyAnnotation: "cloud.com/block",
+						},
+					},
+				).
+				Obj(),
+			wantPodSets: func(job *kftraining.PaddleJob) []kueue.PodSet {
+				return []kueue.PodSet{
+					*utiltesting.MakePodSet(kueue.NewPodSetReference(string(kftraining.PaddleJobReplicaTypeMaster)), 1).
+						PodSpec(job.Spec.PaddleReplicaSpecs[kftraining.PaddleJobReplicaTypeMaster].Template.Spec).
+						Annotations(map[string]string{kueuealpha.PodSetRequiredTopologyAnnotation: "cloud.com/rack"}).
+						Obj(),
+					*utiltesting.MakePodSet(kueue.NewPodSetReference(string(kftraining.PaddleJobReplicaTypeWorker)), 1).
+						PodSpec(job.Spec.PaddleReplicaSpecs[kftraining.PaddleJobReplicaTypeWorker].Template.Spec).
+						Annotations(map[string]string{kueuealpha.PodSetPreferredTopologyAnnotation: "cloud.com/block"}).
+						Obj(),
+				}
+			},
+			enableTopologyAwareScheduling: false,
 		},
 	}
 	for name, tc := range testCases {
 		t.Run(name, func(t *testing.T) {
+			features.SetFeatureGateDuringTest(t, features.TopologyAwareScheduling, tc.enableTopologyAwareScheduling)
 			gotPodSets, err := fromObject(tc.job).PodSets()
 			if err != nil {
 				t.Fatalf("unexpected error: %v", err)
diff --git a/pkg/controller/jobs/kubeflow/jobs/pytorchjob/pytorchjob_controller_test.go b/pkg/controller/jobs/kubeflow/jobs/pytorchjob/pytorchjob_controller_test.go
index 6d5a661c8..986548e9a 100644
--- a/pkg/controller/jobs/kubeflow/jobs/pytorchjob/pytorchjob_controller_test.go
+++ b/pkg/controller/jobs/kubeflow/jobs/pytorchjob/pytorchjob_controller_test.go
@@ -27,6 +27,7 @@ import (
 
 	kueuealpha "sigs.k8s.io/kueue/apis/kueue/v1alpha1"
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
+	"sigs.k8s.io/kueue/pkg/features"
 	utiltesting "sigs.k8s.io/kueue/pkg/util/testing"
 	testingpytorchjob "sigs.k8s.io/kueue/pkg/util/testingjobs/pytorchjob"
 )
@@ -225,8 +226,9 @@ func TestOrderedReplicaTypes(t *testing.T) {
 
 func TestPodSets(t *testing.T) {
 	testCases := map[string]struct {
-		job         *kftraining.PyTorchJob
-		wantPodSets func(job *kftraining.PyTorchJob) []kueue.PodSet
+		job                           *kftraining.PyTorchJob
+		wantPodSets                   func(job *kftraining.PyTorchJob) []kueue.PodSet
+		enableTopologyAwareScheduling bool
 	}{
 		"no annotations": {
 			job: testingpytorchjob.MakePyTorchJob("pytorchjob", "ns").
@@ -251,6 +253,7 @@ func TestPodSets(t *testing.T) {
 						Obj(),
 				}
 			},
+			enableTopologyAwareScheduling: false,
 		},
 		"with required and preferred topology annotation": {
 			job: testingpytorchjob.MakePyTorchJob("pytorchjob", "ns").
@@ -287,10 +290,45 @@ func TestPodSets(t *testing.T) {
 						Obj(),
 				}
 			},
+			enableTopologyAwareScheduling: true,
+		},
+		"without required and preferred topology annotation if TAS is disabled": {
+			job: testingpytorchjob.MakePyTorchJob("pytorchjob", "ns").
+				PyTorchReplicaSpecs(
+					testingpytorchjob.PyTorchReplicaSpecRequirement{
+						ReplicaType:  kftraining.PyTorchJobReplicaTypeMaster,
+						ReplicaCount: 1,
+						Annotations: map[string]string{
+							kueuealpha.PodSetRequiredTopologyAnnotation: "cloud.com/rack",
+						},
+					},
+					testingpytorchjob.PyTorchReplicaSpecRequirement{
+						ReplicaType:  kftraining.PyTorchJobReplicaTypeWorker,
+						ReplicaCount: 1,
+						Annotations: map[string]string{
+							kueuealpha.PodSetPreferredTopologyAnnotation: "cloud.com/block",
+						},
+					},
+				).
+				Obj(),
+			wantPodSets: func(job *kftraining.PyTorchJob) []kueue.PodSet {
+				return []kueue.PodSet{
+					*utiltesting.MakePodSet(kueue.NewPodSetReference(string(kftraining.PyTorchJobReplicaTypeMaster)), 1).
+						PodSpec(job.Spec.PyTorchReplicaSpecs[kftraining.PyTorchJobReplicaTypeMaster].Template.Spec).
+						Annotations(map[string]string{kueuealpha.PodSetRequiredTopologyAnnotation: "cloud.com/rack"}).
+						Obj(),
+					*utiltesting.MakePodSet(kueue.NewPodSetReference(string(kftraining.PyTorchJobReplicaTypeWorker)), 1).
+						PodSpec(job.Spec.PyTorchReplicaSpecs[kftraining.PyTorchJobReplicaTypeWorker].Template.Spec).
+						Annotations(map[string]string{kueuealpha.PodSetPreferredTopologyAnnotation: "cloud.com/block"}).
+						Obj(),
+				}
+			},
+			enableTopologyAwareScheduling: false,
 		},
 	}
 	for name, tc := range testCases {
 		t.Run(name, func(t *testing.T) {
+			features.SetFeatureGateDuringTest(t, features.TopologyAwareScheduling, tc.enableTopologyAwareScheduling)
 			gotPodSets, err := fromObject(tc.job).PodSets()
 			if err != nil {
 				t.Fatalf("unexpected error: %v", err)
diff --git a/pkg/controller/jobs/kubeflow/jobs/tfjob/tfjob_controller_test.go b/pkg/controller/jobs/kubeflow/jobs/tfjob/tfjob_controller_test.go
index dab1a8d3e..10187c582 100644
--- a/pkg/controller/jobs/kubeflow/jobs/tfjob/tfjob_controller_test.go
+++ b/pkg/controller/jobs/kubeflow/jobs/tfjob/tfjob_controller_test.go
@@ -27,6 +27,7 @@ import (
 
 	kueuealpha "sigs.k8s.io/kueue/apis/kueue/v1alpha1"
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
+	"sigs.k8s.io/kueue/pkg/features"
 	utiltesting "sigs.k8s.io/kueue/pkg/util/testing"
 	testingtfjob "sigs.k8s.io/kueue/pkg/util/testingjobs/tfjob"
 )
@@ -243,8 +244,9 @@ func TestOrderedReplicaTypes(t *testing.T) {
 
 func TestPodSets(t *testing.T) {
 	testCases := map[string]struct {
-		job         *kftraining.TFJob
-		wantPodSets func(job *kftraining.TFJob) []kueue.PodSet
+		job                           *kftraining.TFJob
+		wantPodSets                   func(job *kftraining.TFJob) []kueue.PodSet
+		enableTopologyAwareScheduling bool
 	}{
 		"no annotations": {
 			job: testingtfjob.MakeTFJob("tfjob", "ns").TFReplicaSpecsDefault().Obj(),
@@ -261,6 +263,7 @@ func TestPodSets(t *testing.T) {
 						Obj(),
 				}
 			},
+			enableTopologyAwareScheduling: false,
 		},
 		"with required and preferred topology annotation": {
 			job: testingtfjob.MakeTFJob("tfjob", "ns").
@@ -304,10 +307,52 @@ func TestPodSets(t *testing.T) {
 						Obj(),
 				}
 			},
+			enableTopologyAwareScheduling: true,
+		},
+		"without required and preferred topology annotation if TAS is disabled": {
+			job: testingtfjob.MakeTFJob("tfjob", "ns").
+				TFReplicaSpecs(
+					testingtfjob.TFReplicaSpecRequirement{
+						ReplicaType:  kftraining.TFJobReplicaTypeChief,
+						ReplicaCount: 1,
+						Annotations: map[string]string{
+							kueuealpha.PodSetRequiredTopologyAnnotation: "cloud.com/rack",
+						},
+					},
+					testingtfjob.TFReplicaSpecRequirement{
+						ReplicaType:  kftraining.TFJobReplicaTypePS,
+						ReplicaCount: 1,
+						Annotations: map[string]string{
+							kueuealpha.PodSetPreferredTopologyAnnotation: "cloud.com/block",
+						},
+					},
+					testingtfjob.TFReplicaSpecRequirement{
+						ReplicaType:  kftraining.TFJobReplicaTypeWorker,
+						ReplicaCount: 1,
+					},
+				).
+				Obj(),
+			wantPodSets: func(job *kftraining.TFJob) []kueue.PodSet {
+				return []kueue.PodSet{
+					*utiltesting.MakePodSet(kueue.NewPodSetReference(string(kftraining.TFJobReplicaTypeChief)), 1).
+						PodSpec(job.Spec.TFReplicaSpecs[kftraining.TFJobReplicaTypeChief].Template.Spec).
+						Annotations(map[string]string{kueuealpha.PodSetRequiredTopologyAnnotation: "cloud.com/rack"}).
+						Obj(),
+					*utiltesting.MakePodSet(kueue.NewPodSetReference(string(kftraining.TFJobReplicaTypePS)), 1).
+						PodSpec(job.Spec.TFReplicaSpecs[kftraining.TFJobReplicaTypePS].Template.Spec).
+						Annotations(map[string]string{kueuealpha.PodSetPreferredTopologyAnnotation: "cloud.com/block"}).
+						Obj(),
+					*utiltesting.MakePodSet(kueue.NewPodSetReference(string(kftraining.TFJobReplicaTypeWorker)), 1).
+						PodSpec(job.Spec.TFReplicaSpecs[kftraining.TFJobReplicaTypeWorker].Template.Spec).
+						Obj(),
+				}
+			},
+			enableTopologyAwareScheduling: false,
 		},
 	}
 	for name, tc := range testCases {
 		t.Run(name, func(t *testing.T) {
+			features.SetFeatureGateDuringTest(t, features.TopologyAwareScheduling, tc.enableTopologyAwareScheduling)
 			gotPodSets, err := fromObject(tc.job).PodSets()
 			if err != nil {
 				t.Fatalf("unexpected error: %v", err)
diff --git a/pkg/controller/jobs/kubeflow/jobs/xgboostjob/xgboostjob_controller_test.go b/pkg/controller/jobs/kubeflow/jobs/xgboostjob/xgboostjob_controller_test.go
index 52a94a185..ff1436a86 100644
--- a/pkg/controller/jobs/kubeflow/jobs/xgboostjob/xgboostjob_controller_test.go
+++ b/pkg/controller/jobs/kubeflow/jobs/xgboostjob/xgboostjob_controller_test.go
@@ -35,6 +35,7 @@ import (
 	kueuealpha "sigs.k8s.io/kueue/apis/kueue/v1alpha1"
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
 	"sigs.k8s.io/kueue/pkg/controller/jobframework"
+	"sigs.k8s.io/kueue/pkg/features"
 	utiltesting "sigs.k8s.io/kueue/pkg/util/testing"
 	testingxgboostjob "sigs.k8s.io/kueue/pkg/util/testingjobs/xgboostjob"
 )
@@ -231,8 +232,9 @@ func TestOrderedReplicaTypes(t *testing.T) {
 
 func TestPodSets(t *testing.T) {
 	testCases := map[string]struct {
-		job         *kftraining.XGBoostJob
-		wantPodSets func(job *kftraining.XGBoostJob) []kueue.PodSet
+		job                           *kftraining.XGBoostJob
+		wantPodSets                   func(job *kftraining.XGBoostJob) []kueue.PodSet
+		enableTopologyAwareScheduling bool
 	}{
 		"no annotations": {
 			job: testingxgboostjob.MakeXGBoostJob("xgboostjob", "ns").
@@ -257,6 +259,7 @@ func TestPodSets(t *testing.T) {
 						Obj(),
 				}
 			},
+			enableTopologyAwareScheduling: false,
 		},
 		"with required and preferred topology annotation": {
 			job: testingxgboostjob.MakeXGBoostJob("xgboostjob", "ns").
@@ -293,10 +296,45 @@ func TestPodSets(t *testing.T) {
 						Obj(),
 				}
 			},
+			enableTopologyAwareScheduling: true,
+		},
+		"without required and preferred topology annotation if TAS is disabled": {
+			job: testingxgboostjob.MakeXGBoostJob("xgboostjob", "ns").
+				XGBReplicaSpecs(
+					testingxgboostjob.XGBReplicaSpecRequirement{
+						ReplicaType:  kftraining.XGBoostJobReplicaTypeMaster,
+						ReplicaCount: 1,
+						Annotations: map[string]string{
+							kueuealpha.PodSetRequiredTopologyAnnotation: "cloud.com/rack",
+						},
+					},
+					testingxgboostjob.XGBReplicaSpecRequirement{
+						ReplicaType:  kftraining.XGBoostJobReplicaTypeWorker,
+						ReplicaCount: 1,
+						Annotations: map[string]string{
+							kueuealpha.PodSetPreferredTopologyAnnotation: "cloud.com/block",
+						},
+					},
+				).
+				Obj(),
+			wantPodSets: func(job *kftraining.XGBoostJob) []kueue.PodSet {
+				return []kueue.PodSet{
+					*utiltesting.MakePodSet(kueue.NewPodSetReference(string(kftraining.XGBoostJobReplicaTypeMaster)), 1).
+						PodSpec(job.Spec.XGBReplicaSpecs[kftraining.XGBoostJobReplicaTypeMaster].Template.Spec).
+						Annotations(map[string]string{kueuealpha.PodSetRequiredTopologyAnnotation: "cloud.com/rack"}).
+						Obj(),
+					*utiltesting.MakePodSet(kueue.NewPodSetReference(string(kftraining.XGBoostJobReplicaTypeWorker)), 1).
+						PodSpec(job.Spec.XGBReplicaSpecs[kftraining.XGBoostJobReplicaTypeWorker].Template.Spec).
+						Annotations(map[string]string{kueuealpha.PodSetPreferredTopologyAnnotation: "cloud.com/block"}).
+						Obj(),
+				}
+			},
+			enableTopologyAwareScheduling: false,
 		},
 	}
 	for name, tc := range testCases {
 		t.Run(name, func(t *testing.T) {
+			features.SetFeatureGateDuringTest(t, features.TopologyAwareScheduling, tc.enableTopologyAwareScheduling)
 			gotPodSets, err := fromObject(tc.job).PodSets()
 			if err != nil {
 				t.Fatalf("unexpected error: %v", err)
diff --git a/pkg/controller/jobs/leaderworkerset/leaderworkerset_reconciler_test.go b/pkg/controller/jobs/leaderworkerset/leaderworkerset_reconciler_test.go
index ef9b7ecd1..05a5dc358 100644
--- a/pkg/controller/jobs/leaderworkerset/leaderworkerset_reconciler_test.go
+++ b/pkg/controller/jobs/leaderworkerset/leaderworkerset_reconciler_test.go
@@ -36,6 +36,7 @@ import (
 	"sigs.k8s.io/kueue/pkg/constants"
 	"sigs.k8s.io/kueue/pkg/controller/jobframework"
 	podconstants "sigs.k8s.io/kueue/pkg/controller/jobs/pod/constants"
+	"sigs.k8s.io/kueue/pkg/features"
 	utiltesting "sigs.k8s.io/kueue/pkg/util/testing"
 	"sigs.k8s.io/kueue/pkg/util/testingjobs/leaderworkerset"
 )
@@ -55,14 +56,15 @@ var (
 
 func TestReconciler(t *testing.T) {
 	cases := map[string]struct {
-		labelKeysToCopy         []string
-		leaderWorkerSet         *leaderworkersetv1.LeaderWorkerSet
-		workloads               []kueue.Workload
-		workloadPriorityClasses []kueue.WorkloadPriorityClass
-		wantLeaderWorkerSet     *leaderworkersetv1.LeaderWorkerSet
-		wantWorkloads           []kueue.Workload
-		wantEvents              []utiltesting.EventRecord
-		wantErr                 error
+		labelKeysToCopy               []string
+		leaderWorkerSet               *leaderworkersetv1.LeaderWorkerSet
+		workloads                     []kueue.Workload
+		workloadPriorityClasses       []kueue.WorkloadPriorityClass
+		wantLeaderWorkerSet           *leaderworkersetv1.LeaderWorkerSet
+		wantWorkloads                 []kueue.Workload
+		wantEvents                    []utiltesting.EventRecord
+		wantErr                       error
+		enableTopologyAwareScheduling bool
 	}{
 		"should create prebuilt workload": {
 			leaderWorkerSet:     leaderworkerset.MakeLeaderWorkerSet(testLWS, testNS).UID(testUID).Obj(),
@@ -99,6 +101,7 @@ func TestReconciler(t *testing.T) {
 					),
 				},
 			},
+			enableTopologyAwareScheduling: false,
 		},
 		"should create prebuilt workload with leader template": {
 			leaderWorkerSet: leaderworkerset.MakeLeaderWorkerSet(testLWS, testNS).
@@ -168,6 +171,7 @@ func TestReconciler(t *testing.T) {
 					),
 				},
 			},
+			enableTopologyAwareScheduling: false,
 		},
 		"should create prebuilt workloads with leader template": {
 			leaderWorkerSet: leaderworkerset.MakeLeaderWorkerSet(testLWS, testNS).
@@ -280,6 +284,7 @@ func TestReconciler(t *testing.T) {
 					),
 				},
 			},
+			enableTopologyAwareScheduling: false,
 		},
 		"should create prebuilt workload with required topology annotation": {
 			leaderWorkerSet: leaderworkerset.MakeLeaderWorkerSet(testLWS, testNS).
@@ -388,6 +393,109 @@ func TestReconciler(t *testing.T) {
 					),
 				},
 			},
+			enableTopologyAwareScheduling: true,
+		},
+		"should create prebuilt workload without required topology annotation is TAS is disabled": {
+			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(podconstants.IsGroupWorkloadAnnotationKey, podconstants.IsGroupWorkloadAnnotationValue).
+					Finalizers(kueue.ResourceInUseFinalizerName).
+					PodSets(
+						kueue.PodSet{
+							Name: leaderPodSetName,
+							Template: corev1.PodTemplateSpec{
+								Spec: corev1.PodSpec{
+									Containers: []corev1.Container{
+										{Name: "c", Image: "pause"},
+									},
+								},
+							},
+							Count: 1,
+						},
+						kueue.PodSet{
+							Name: workerPodSetName,
+							Template: corev1.PodTemplateSpec{
+								Spec: corev1.PodSpec{
+									Containers: []corev1.Container{
+										{Name: "c", Image: "pause"},
+									},
+								},
+							},
+							Count: 2,
+						},
+					).
+					Priority(0).
+					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"),
+					),
+				},
+			},
+			enableTopologyAwareScheduling: false,
 		},
 		"should create prebuilt workload with workload priority": {
 			leaderWorkerSet: leaderworkerset.MakeLeaderWorkerSet(testLWS, testNS).
@@ -435,10 +543,12 @@ func TestReconciler(t *testing.T) {
 					),
 				},
 			},
+			enableTopologyAwareScheduling: false,
 		},
 	}
 	for name, tc := range cases {
 		t.Run(name, func(t *testing.T) {
+			features.SetFeatureGateDuringTest(t, features.TopologyAwareScheduling, tc.enableTopologyAwareScheduling)
 			ctx, _ := utiltesting.ContextWithLog(t)
 			clientBuilder := utiltesting.NewClientBuilder(leaderworkersetv1.AddToScheme)
 
diff --git a/pkg/controller/jobs/mpijob/mpijob_controller_test.go b/pkg/controller/jobs/mpijob/mpijob_controller_test.go
index 35ae926bd..5be2f4c0c 100644
--- a/pkg/controller/jobs/mpijob/mpijob_controller_test.go
+++ b/pkg/controller/jobs/mpijob/mpijob_controller_test.go
@@ -34,6 +34,7 @@ import (
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
 	"sigs.k8s.io/kueue/pkg/constants"
 	"sigs.k8s.io/kueue/pkg/controller/jobframework"
+	"sigs.k8s.io/kueue/pkg/features"
 	utiltesting "sigs.k8s.io/kueue/pkg/util/testing"
 	testingmpijob "sigs.k8s.io/kueue/pkg/util/testingjobs/mpijob"
 )
@@ -196,8 +197,9 @@ func TestPodSets(t *testing.T) {
 	)
 
 	testCases := map[string]struct {
-		job         *MPIJob
-		wantPodSets []kueue.PodSet
+		job                           *MPIJob
+		wantPodSets                   []kueue.PodSet
+		enableTopologyAwareScheduling bool
 	}{
 		"no annotations": {
 			job: (*MPIJob)(jobTemplate.DeepCopy()),
@@ -209,6 +211,7 @@ func TestPodSets(t *testing.T) {
 					PodSpec(jobTemplate.Clone().Spec.MPIReplicaSpecs[kfmpi.MPIReplicaTypeWorker].Template.Spec).
 					Obj(),
 			},
+			enableTopologyAwareScheduling: false,
 		},
 		"with required topology annotation": {
 			job: (*MPIJob)(jobTemplate.Clone().
@@ -238,6 +241,7 @@ func TestPodSets(t *testing.T) {
 					PodIndexLabel(ptr.To(kfmpi.ReplicaIndexLabel)).
 					Obj(),
 			},
+			enableTopologyAwareScheduling: true,
 		},
 		"with preferred topology annotation": {
 			job: (*MPIJob)(jobTemplate.Clone().
@@ -267,10 +271,64 @@ func TestPodSets(t *testing.T) {
 					PodIndexLabel(ptr.To(kfmpi.ReplicaIndexLabel)).
 					Obj(),
 			},
+			enableTopologyAwareScheduling: true,
+		},
+		"without required and preferred topology annotation if TAS is disabled": {
+			job: (*MPIJob)(jobTemplate.Clone().
+				PodAnnotation(
+					kfmpi.MPIReplicaTypeLauncher,
+					kueuealpha.PodSetRequiredTopologyAnnotation,
+					"cloud.com/block",
+				).
+				PodAnnotation(
+					kfmpi.MPIReplicaTypeWorker,
+					kueuealpha.PodSetRequiredTopologyAnnotation,
+					"cloud.com/block",
+				).
+				Obj(),
+			),
+			wantPodSets: []kueue.PodSet{
+				*utiltesting.MakePodSet(kueue.NewPodSetReference(string(kfmpi.MPIReplicaTypeLauncher)), 1).
+					PodSpec(jobTemplate.Clone().Spec.MPIReplicaSpecs[kfmpi.MPIReplicaTypeLauncher].Template.Spec).
+					Annotations(map[string]string{kueuealpha.PodSetRequiredTopologyAnnotation: "cloud.com/block"}).
+					Obj(),
+				*utiltesting.MakePodSet(kueue.NewPodSetReference(string(kfmpi.MPIReplicaTypeWorker)), 3).
+					PodSpec(jobTemplate.Clone().Spec.MPIReplicaSpecs[kfmpi.MPIReplicaTypeWorker].Template.Spec).
+					Annotations(map[string]string{kueuealpha.PodSetRequiredTopologyAnnotation: "cloud.com/block"}).
+					Obj(),
+			},
+			enableTopologyAwareScheduling: false,
+		},
+		"without preferred topology annotation if TAS is disabled": {
+			job: (*MPIJob)(jobTemplate.Clone().
+				PodAnnotation(
+					kfmpi.MPIReplicaTypeLauncher,
+					kueuealpha.PodSetPreferredTopologyAnnotation,
+					"cloud.com/block",
+				).
+				PodAnnotation(
+					kfmpi.MPIReplicaTypeWorker,
+					kueuealpha.PodSetPreferredTopologyAnnotation,
+					"cloud.com/block",
+				).
+				Obj(),
+			),
+			wantPodSets: []kueue.PodSet{
+				*utiltesting.MakePodSet(kueue.NewPodSetReference(string(kfmpi.MPIReplicaTypeLauncher)), 1).
+					PodSpec(jobTemplate.Clone().Spec.MPIReplicaSpecs[kfmpi.MPIReplicaTypeLauncher].Template.Spec).
+					Annotations(map[string]string{kueuealpha.PodSetPreferredTopologyAnnotation: "cloud.com/block"}).
+					Obj(),
+				*utiltesting.MakePodSet(kueue.NewPodSetReference(string(kfmpi.MPIReplicaTypeWorker)), 3).
+					PodSpec(jobTemplate.Clone().Spec.MPIReplicaSpecs[kfmpi.MPIReplicaTypeWorker].Template.Spec).
+					Annotations(map[string]string{kueuealpha.PodSetPreferredTopologyAnnotation: "cloud.com/block"}).
+					Obj(),
+			},
+			enableTopologyAwareScheduling: false,
 		},
 	}
 	for name, tc := range testCases {
 		t.Run(name, func(t *testing.T) {
+			features.SetFeatureGateDuringTest(t, features.TopologyAwareScheduling, tc.enableTopologyAwareScheduling)
 			gotPodSets, err := tc.job.PodSets()
 			if err != nil {
 				t.Fatalf("unexpected error: %v", err)
diff --git a/pkg/controller/jobs/pod/pod_controller_test.go b/pkg/controller/jobs/pod/pod_controller_test.go
index 1fba8882d..c8ca3906c 100644
--- a/pkg/controller/jobs/pod/pod_controller_test.go
+++ b/pkg/controller/jobs/pod/pod_controller_test.go
@@ -43,6 +43,7 @@ import (
 	controllerconsts "sigs.k8s.io/kueue/pkg/controller/constants"
 	"sigs.k8s.io/kueue/pkg/controller/jobframework"
 	podconstants "sigs.k8s.io/kueue/pkg/controller/jobs/pod/constants"
+	"sigs.k8s.io/kueue/pkg/features"
 	"sigs.k8s.io/kueue/pkg/podset"
 	utilpod "sigs.k8s.io/kueue/pkg/util/pod"
 	utiltesting "sigs.k8s.io/kueue/pkg/util/testing"
@@ -130,8 +131,9 @@ func TestRun(t *testing.T) {
 
 func TestPodSets(t *testing.T) {
 	testCases := map[string]struct {
-		pod         *Pod
-		wantPodSets func(pod *Pod) []kueue.PodSet
+		pod                           *Pod
+		wantPodSets                   func(pod *Pod) []kueue.PodSet
+		enableTopologyAwareScheduling bool
 	}{
 		"no annotations": {
 			pod: FromObject(testingpod.MakePod("pod", "ns").Obj()),
@@ -142,6 +144,7 @@ func TestPodSets(t *testing.T) {
 						Obj(),
 				}
 			},
+			enableTopologyAwareScheduling: false,
 		},
 		"with required topology annotation": {
 			pod: FromObject(testingpod.MakePod("pod", "ns").
@@ -157,6 +160,7 @@ func TestPodSets(t *testing.T) {
 						Obj(),
 				}
 			},
+			enableTopologyAwareScheduling: true,
 		},
 		"with required topology preferred": {
 			pod: FromObject(testingpod.MakePod("pod", "ns").
@@ -172,10 +176,40 @@ func TestPodSets(t *testing.T) {
 						Obj(),
 				}
 			},
+			enableTopologyAwareScheduling: true,
+		},
+		"without required topology annotation if TAS is disabled": {
+			pod: FromObject(testingpod.MakePod("pod", "ns").
+				Annotation(kueuealpha.PodSetRequiredTopologyAnnotation, "cloud.com/block").
+				Obj(),
+			),
+			wantPodSets: func(pod *Pod) []kueue.PodSet {
+				return []kueue.PodSet{
+					*utiltesting.MakePodSet(kueue.DefaultPodSetName, 1).
+						PodSpec(*pod.pod.Spec.DeepCopy()).
+						Obj(),
+				}
+			},
+			enableTopologyAwareScheduling: false,
+		},
+		"without preferred topology annotation if TAS is disabled": {
+			pod: FromObject(testingpod.MakePod("pod", "ns").
+				Annotation(kueuealpha.PodSetPreferredTopologyAnnotation, "cloud.com/block").
+				Obj(),
+			),
+			wantPodSets: func(pod *Pod) []kueue.PodSet {
+				return []kueue.PodSet{
+					*utiltesting.MakePodSet(kueue.DefaultPodSetName, 1).
+						PodSpec(*pod.pod.Spec.DeepCopy()).
+						Obj(),
+				}
+			},
+			enableTopologyAwareScheduling: false,
 		},
 	}
 	for name, tc := range testCases {
 		t.Run(name, func(t *testing.T) {
+			features.SetFeatureGateDuringTest(t, features.TopologyAwareScheduling, tc.enableTopologyAwareScheduling)
 			gotPodSets, err := tc.pod.PodSets()
 			if err != nil {
 				t.Fatalf("unexpected error: %v", err)
diff --git a/pkg/controller/jobs/raycluster/raycluster_controller_test.go b/pkg/controller/jobs/raycluster/raycluster_controller_test.go
index 61948bf33..5165d955d 100644
--- a/pkg/controller/jobs/raycluster/raycluster_controller_test.go
+++ b/pkg/controller/jobs/raycluster/raycluster_controller_test.go
@@ -36,6 +36,7 @@ import (
 	kueuealpha "sigs.k8s.io/kueue/apis/kueue/v1alpha1"
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
 	"sigs.k8s.io/kueue/pkg/controller/jobframework"
+	"sigs.k8s.io/kueue/pkg/features"
 	"sigs.k8s.io/kueue/pkg/podset"
 	utiltesting "sigs.k8s.io/kueue/pkg/util/testing"
 	testingrayutil "sigs.k8s.io/kueue/pkg/util/testingjobs/raycluster"
@@ -57,8 +58,9 @@ var (
 
 func TestPodSets(t *testing.T) {
 	testCases := map[string]struct {
-		rayCluster  *RayCluster
-		wantPodSets func(rayJob *RayCluster) []kueue.PodSet
+		rayCluster                    *RayCluster
+		wantPodSets                   func(rayJob *RayCluster) []kueue.PodSet
+		enableTopologyAwareScheduling bool
 	}{
 		"no annotations": {
 			rayCluster: (*RayCluster)(testingrayutil.MakeCluster("raycluster", "ns").
@@ -98,6 +100,7 @@ func TestPodSets(t *testing.T) {
 						Obj(),
 				}
 			},
+			enableTopologyAwareScheduling: false,
 		},
 		"with required topology annotation": {
 			rayCluster: (*RayCluster)(testingrayutil.MakeCluster("raycluster", "ns").
@@ -151,6 +154,7 @@ func TestPodSets(t *testing.T) {
 						Obj(),
 				}
 			},
+			enableTopologyAwareScheduling: true,
 		},
 		"with preferred topology annotation": {
 			rayCluster: (*RayCluster)(testingrayutil.MakeCluster("raycluster", "ns").
@@ -204,10 +208,80 @@ func TestPodSets(t *testing.T) {
 						Obj(),
 				}
 			},
+			enableTopologyAwareScheduling: true,
+		},
+		"without required and preferred topology annotation if TAS is disabled": {
+			rayCluster: (*RayCluster)(testingrayutil.MakeCluster("raycluster", "ns").
+				WithHeadGroupSpec(
+					rayv1.HeadGroupSpec{
+						Template: corev1.PodTemplateSpec{
+							ObjectMeta: metav1.ObjectMeta{
+								Annotations: map[string]string{
+									kueuealpha.PodSetRequiredTopologyAnnotation: "cloud.com/block",
+								},
+							},
+							Spec: corev1.PodSpec{Containers: []corev1.Container{{Name: "head_c"}}},
+						},
+					},
+				).
+				WithWorkerGroups(
+					rayv1.WorkerGroupSpec{
+						GroupName: "group1",
+						Template: corev1.PodTemplateSpec{
+							ObjectMeta: metav1.ObjectMeta{
+								Annotations: map[string]string{
+									kueuealpha.PodSetRequiredTopologyAnnotation: "cloud.com/block",
+								},
+							},
+							Spec: corev1.PodSpec{Containers: []corev1.Container{{Name: "group1_c"}}},
+						},
+					},
+					rayv1.WorkerGroupSpec{
+						GroupName: "group2",
+						Replicas:  ptr.To[int32](3),
+						Template: corev1.PodTemplateSpec{
+							ObjectMeta: metav1.ObjectMeta{
+								Annotations: map[string]string{
+									kueuealpha.PodSetPreferredTopologyAnnotation: "cloud.com/block",
+								},
+							},
+							Spec: corev1.PodSpec{Containers: []corev1.Container{{Name: "group2_c"}}},
+						},
+					},
+					rayv1.WorkerGroupSpec{
+						GroupName: "group3",
+						Replicas:  ptr.To[int32](3),
+						Template: corev1.PodTemplateSpec{
+							Spec: corev1.PodSpec{Containers: []corev1.Container{{Name: "group2_c"}}},
+						},
+					},
+				).
+				Obj()),
+			wantPodSets: func(rayJob *RayCluster) []kueue.PodSet {
+				return []kueue.PodSet{
+					*utiltesting.MakePodSet(headGroupPodSetName, 1).
+						PodSpec(*rayJob.Spec.HeadGroupSpec.Template.Spec.DeepCopy()).
+						Annotations(rayJob.Spec.HeadGroupSpec.Template.Annotations).
+						Obj(),
+					*utiltesting.MakePodSet("group1", 1).
+						PodSpec(*rayJob.Spec.WorkerGroupSpecs[0].Template.Spec.DeepCopy()).
+						Annotations(rayJob.Spec.WorkerGroupSpecs[0].Template.Annotations).
+						Obj(),
+					*utiltesting.MakePodSet("group2", 3).
+						PodSpec(*rayJob.Spec.WorkerGroupSpecs[1].Template.Spec.DeepCopy()).
+						Annotations(rayJob.Spec.WorkerGroupSpecs[1].Template.Annotations).
+						Obj(),
+					*utiltesting.MakePodSet("group3", 3).
+						PodSpec(*rayJob.Spec.WorkerGroupSpecs[2].Template.Spec.DeepCopy()).
+						Obj(),
+				}
+			},
+			enableTopologyAwareScheduling: false,
 		},
 	}
 	for name, tc := range testCases {
 		t.Run(name, func(t *testing.T) {
+			features.SetFeatureGateDuringTest(t, features.TopologyAwareScheduling, tc.enableTopologyAwareScheduling)
 			gotPodSets, err := tc.rayCluster.PodSets()
 			if err != nil {
 				t.Fatalf("unexpected error: %v", err)
diff --git a/pkg/controller/jobs/rayjob/rayjob_controller_test.go b/pkg/controller/jobs/rayjob/rayjob_controller_test.go
index d70d5aa59..2762880f8 100644
--- a/pkg/controller/jobs/rayjob/rayjob_controller_test.go
+++ b/pkg/controller/jobs/rayjob/rayjob_controller_test.go
@@ -29,6 +29,7 @@ import (
 
 	kueuealpha "sigs.k8s.io/kueue/apis/kueue/v1alpha1"
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
+	"sigs.k8s.io/kueue/pkg/features"
 	"sigs.k8s.io/kueue/pkg/podset"
 	utiltesting "sigs.k8s.io/kueue/pkg/util/testing"
 	testingrayutil "sigs.k8s.io/kueue/pkg/util/testingjobs/rayjob"
@@ -36,8 +37,9 @@ import (
 
 func TestPodSets(t *testing.T) {
 	testCases := map[string]struct {
-		rayJob      *RayJob
-		wantPodSets func(rayJob *RayJob) []kueue.PodSet
+		rayJob                        *RayJob
+		wantPodSets                   func(rayJob *RayJob) []kueue.PodSet
+		enableTopologyAwareScheduling bool
 	}{
 		"no annotations": {
 			rayJob: (*RayJob)(testingrayutil.MakeJob("rayjob", "ns").
@@ -77,6 +79,7 @@ func TestPodSets(t *testing.T) {
 						Obj(),
 				}
 			},
+			enableTopologyAwareScheduling: false,
 		},
 		"with required topology annotation": {
 			rayJob: (*RayJob)(testingrayutil.MakeJob("rayjob", "ns").
@@ -130,6 +133,7 @@ func TestPodSets(t *testing.T) {
 						Obj(),
 				}
 			},
+			enableTopologyAwareScheduling: true,
 		},
 		"with preferred topology annotation": {
 			rayJob: (*RayJob)(testingrayutil.MakeJob("rayjob", "ns").
@@ -183,6 +187,65 @@ func TestPodSets(t *testing.T) {
 						Obj(),
 				}
 			},
+			enableTopologyAwareScheduling: true,
+		},
+		"without required and preferred topology annotation if TAS is disabled": {
+			rayJob: (*RayJob)(testingrayutil.MakeJob("rayjob", "ns").
+				WithHeadGroupSpec(
+					rayv1.HeadGroupSpec{
+						Template: corev1.PodTemplateSpec{
+							ObjectMeta: metav1.ObjectMeta{
+								Annotations: map[string]string{
+									kueuealpha.PodSetRequiredTopologyAnnotation: "cloud.com/block",
+								},
+							},
+							Spec: corev1.PodSpec{Containers: []corev1.Container{{Name: "head_c"}}},
+						},
+					},
+				).
+				WithWorkerGroups(
+					rayv1.WorkerGroupSpec{
+						GroupName: "group1",
+						Template: corev1.PodTemplateSpec{
+							ObjectMeta: metav1.ObjectMeta{
+								Annotations: map[string]string{
+									kueuealpha.PodSetRequiredTopologyAnnotation: "cloud.com/block",
+								},
+							},
+							Spec: corev1.PodSpec{Containers: []corev1.Container{{Name: "group1_c"}}},
+						},
+					},
+					rayv1.WorkerGroupSpec{
+						GroupName: "group2",
+						Replicas:  ptr.To[int32](3),
+						Template: corev1.PodTemplateSpec{
+							ObjectMeta: metav1.ObjectMeta{
+								Annotations: map[string]string{
+									kueuealpha.PodSetPreferredTopologyAnnotation: "cloud.com/block",
+								},
+							},
+							Spec: corev1.PodSpec{Containers: []corev1.Container{{Name: "group2_c"}}},
+						},
+					},
+				).
+				Obj()),
+			wantPodSets: func(rayJob *RayJob) []kueue.PodSet {
+				return []kueue.PodSet{
+					*utiltesting.MakePodSet(headGroupPodSetName, 1).
+						PodSpec(*rayJob.Spec.RayClusterSpec.HeadGroupSpec.Template.Spec.DeepCopy()).
+						Annotations(rayJob.Spec.RayClusterSpec.HeadGroupSpec.Template.Annotations).
+						Obj(),
+					*utiltesting.MakePodSet(kueue.NewPodSetReference(rayJob.Spec.RayClusterSpec.WorkerGroupSpecs[0].GroupName), 1).
+						PodSpec(*rayJob.Spec.RayClusterSpec.WorkerGroupSpecs[0].Template.Spec.DeepCopy()).
+						Annotations(rayJob.Spec.RayClusterSpec.WorkerGroupSpecs[0].Template.Annotations).
+						Obj(),
+					*utiltesting.MakePodSet(kueue.NewPodSetReference(rayJob.Spec.RayClusterSpec.WorkerGroupSpecs[1].GroupName), 3).
+						PodSpec(*rayJob.Spec.RayClusterSpec.WorkerGroupSpecs[1].Template.Spec.DeepCopy()).
+						Annotations(rayJob.Spec.RayClusterSpec.WorkerGroupSpecs[1].Template.Annotations).
+						Obj(),
+				}
+			},
+			enableTopologyAwareScheduling: false,
 		},
 		"with submitter topology annotation not specified": {
 			rayJob: (*RayJob)(testingrayutil.MakeJob("rayjob", "ns").
@@ -231,6 +294,7 @@ func TestPodSets(t *testing.T) {
 						Obj(),
 				}
 			},
+			enableTopologyAwareScheduling: true,
 		},
 		"with submitter topology annotation specified": {
 			rayJob: (*RayJob)(testingrayutil.MakeJob("rayjob", "ns").
@@ -296,6 +360,7 @@ func TestPodSets(t *testing.T) {
 						Obj(),
 				}
 			},
+			enableTopologyAwareScheduling: true,
 		},
 		"with NumOfHosts > 1": {
 			rayJob: (*RayJob)(testingrayutil.MakeJob("rayjob", "ns").
@@ -337,6 +402,7 @@ func TestPodSets(t *testing.T) {
 						Obj(),
 				}
 			},
+			enableTopologyAwareScheduling: false,
 		},
 		"with default job submitter": {
 			rayJob: (*RayJob)(testingrayutil.MakeJob("rayjob", "ns").
@@ -380,6 +446,7 @@ func TestPodSets(t *testing.T) {
 						Obj(),
 				}
 			},
+			enableTopologyAwareScheduling: false,
 		},
 		"with submitter job pod template override": {
 			rayJob: (*RayJob)(testingrayutil.MakeJob("rayjob", "ns").
@@ -448,10 +515,12 @@ func TestPodSets(t *testing.T) {
 						Obj(),
 				}
 			},
+			enableTopologyAwareScheduling: false,
 		},
 	}
 	for name, tc := range testCases {
 		t.Run(name, func(t *testing.T) {
+			features.SetFeatureGateDuringTest(t, features.TopologyAwareScheduling, tc.enableTopologyAwareScheduling)
 			gotPodSets, err := tc.rayJob.PodSets()
 			if err != nil {
 				t.Fatalf("unexpected error: %v", err)