Repository: kubernetes-sigs/kueue
Pull request #4742: fix: configure PodSetTopologyRequests only when TAS enabled
Author: login: mykysha
Labels: name: kind/bug
description: Categorizes issue or PR as related to a bug.
name: release-note
description: Denotes a PR that will be considered when it comes time to generate release notes.
name: lgtm
description: "Looks good to me", indicates that a PR is ready to be merged.
name: size/XL
description: Denotes a PR that changes 500-999 lines, ignoring generated files.
name: approved
description: Indicates a PR has been approved by an approver from all required OWNERS files.
name: kind/cleanup
description: Categorizes issue or PR as related to cleaning up code, process, or technical debt.
name: cncf-cla: yes
description: Indicates the PR's author has signed the CNCF CLA.
Created: 2025-03-21T09:53:31Z | Merged: 2025-03-26T10:28:34Z

<!--  Thanks for sending a pull request!  Here are some tips for you:

1. If this is your first time, please read our contributor guidelines: https://git.k8s.io/community/contributors/guide/first-contribution.md#your-first-contribution and developer guide https://git.k8s.io/community/contributors/devel/development.md#development-guide
2. Please label this pull request according to what type of issue you are addressing, especially if this is a release targeted pull request. For reference on required PR/issue labels, read here:
https://git.k8s.io/community/contributors/devel/sig-release/release.md#issuepr-kind-label
3. Ensure you have added or ran the appropriate tests for your PR: https://git.k8s.io/community/contributors/devel/sig-testing/testing.md
4. If you want *faster* PR reviews, read how: https://git.k8s.io/community/contributors/guide/pull-requests.md#best-practices-for-faster-reviews
5. If the PR is unfinished, see how to mark it: https://git.k8s.io/community/contributors/guide/pull-requests.md#marking-unfinished-pull-requests
-->

#### What type of PR is this?
/kind cleanup
<!--
Add one of the following kinds:
/kind bug
/kind cleanup
/kind documentation
/kind feature

Optionally add one or more of the following kinds if applicable:
/kind api-change
/kind deprecation
/kind failing-test
/kind flake
/kind regression
-->

#### What this PR does / why we need it:
PodSetTopologyRequests are now configured only when TopologyAwareScheduling feature gate is enabled.

#### Which issue(s) this PR fixes:
<!--
*Automatically closes linked issue when PR is merged.
Usage: `Fixes #<issue number>`, or `Fixes (paste link of issue)`.
_If PR is about `failing-tests or flakes`, please post the related issues/tests in a comment and do not use `Fixes`_*
-->
Fixes #4590

#### Special notes for your reviewer:

#### Does this PR introduce a user-facing change?
<!--
If no, just write "NONE" in the release-note block below.
If yes, a release note is required:
Enter your extended release note in the block below. If the PR requires additional action from users switching to the new release, include the string "action required".
-->
```release-note
PodSetTopologyRequests are now configured only when TopologyAwareScheduling feature gate is enabled.
```

Linked issue #4590: 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.
Linked issue comments: user: login: mimowo
created_at: 2025-03-13T11:49:01Z
body: cc @mbobrovskyi
author_association: CONTRIBUTOR
user: login: mykysha
created_at: 2025-03-13T12:05:27Z
body: /assign
author_association: CONTRIBUTOR

Diff:
diff --git a/pkg/controller/jobs/appwrapper/appwrapper_controller.go b/pkg/controller/jobs/appwrapper/appwrapper_controller.go
index 39b6f7c3ada..584bbd5a344 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/appwrapper/appwrapper_controller_test.go b/pkg/controller/jobs/appwrapper/appwrapper_controller_test.go
index 25c90147c2b..93741f3639a 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.go b/pkg/controller/jobs/job/job_controller.go
index 0570d026b5e..1bb430de36e 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/job/job_controller_test.go b/pkg/controller/jobs/job/job_controller_test.go
index b5d4a38e01f..08868711355 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.go b/pkg/controller/jobs/jobset/jobset_controller.go
index e65393d564f..acae2f01ad1 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/jobset/jobset_controller_test.go b/pkg/controller/jobs/jobset/jobset_controller_test.go
index f352c38ff3a..9f88358461c 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 537a2a2d0d9..25cb7f73ab1 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 6d5a661c8a9..986548e9abf 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 dab1a8d3e53..10187c582f6 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 52a94a1856c..ff1436a86b1 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/kubeflow/kubeflowjob/kubeflowjob_controller.go b/pkg/controller/jobs/kubeflow/kubeflowjob/kubeflowjob_controller.go
index e31c38b5fd5..6a49860b930 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 e49168e5c5b..8bce859b83c 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/leaderworkerset/leaderworkerset_reconciler_test.go b/pkg/controller/jobs/leaderworkerset/leaderworkerset_reconciler_test.go
index ef9b7ecd15b..05a5dc358d8 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.go b/pkg/controller/jobs/mpijob/mpijob_controller.go
index ca0a634ef0b..2b0c0f95df8 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/mpijob/mpijob_controller_test.go b/pkg/controller/jobs/mpijob/mpijob_controller_test.go
index 35ae926bdb8..5be2f4c0ceb 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.go b/pkg/controller/jobs/pod/pod_controller.go
index c42a200df02..8f8f5cb843e 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/pod/pod_controller_test.go b/pkg/controller/jobs/pod/pod_controller_test.go
index deaa8c9ce46..6dad12ca9b7 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.go b/pkg/controller/jobs/raycluster/raycluster_controller.go
index dbc61b013b7..5da47add197 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/raycluster/raycluster_controller_test.go b/pkg/controller/jobs/raycluster/raycluster_controller_test.go
index 61948bf33db..5165d955d30 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.go b/pkg/controller/jobs/rayjob/rayjob_controller.go
index f29bf452b8e..6743045ce21 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/pkg/controller/jobs/rayjob/rayjob_controller_test.go b/pkg/controller/jobs/rayjob/rayjob_controller_test.go
index d70d5aa5940..2762880f8ac 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)


Reviews:
user: login: mimowo
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-21T11:01:56Z
user: login: mykysha
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-21T11:07:35Z
user: login: mykysha
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-26T08:50:12Z

Review comments:
diff_hunk: @@ -132,6 +133,7 @@ func TestPodSets(t *testing.T) {
 
 	for name, tc := range testCases {
 		t.Run(name, func(t *testing.T) {
+			features.SetFeatureGateDuringTest(t, features.TopologyAwareScheduling, true)
path: pkg/controller/jobs/appwrapper/appwrapper_controller_test.go
user: login: mimowo
body: IIUC the tests would pass even without changing the code. Let's introduce a new param enableTopologyAwareScheduling, and add a sanity test when `enableTopologyAwareScheduling=false`.

Same remark for other tests.
created_at: 2025-03-21T11:01:55Z
author_association: CONTRIBUTOR
diff_hunk: @@ -132,6 +133,7 @@ func TestPodSets(t *testing.T) {
 
 	for name, tc := range testCases {
 		t.Run(name, func(t *testing.T) {
+			features.SetFeatureGateDuringTest(t, features.TopologyAwareScheduling, true)
path: pkg/controller/jobs/appwrapper/appwrapper_controller_test.go
user: login: mykysha
body: I've added these changes only to the tests that were failing, so enabling the feature gate is required. Will add a `enableTopologyAwareScheduling=false` case
created_at: 2025-03-21T11:07:35Z
author_association: CONTRIBUTOR
diff_hunk: @@ -132,6 +133,7 @@ func TestPodSets(t *testing.T) {
 
 	for name, tc := range testCases {
 		t.Run(name, func(t *testing.T) {
+			features.SetFeatureGateDuringTest(t, features.TopologyAwareScheduling, true)
path: pkg/controller/jobs/appwrapper/appwrapper_controller_test.go
user: login: mykysha
body: Added `enableTopologyAwareScheduling` to the affected tests, as well as sanity tests where it is intentionally disabled.
created_at: 2025-03-26T08:50:12Z
author_association: CONTRIBUTOR

PR comments:
user: login: k8s-ci-robot
created_at: 2025-03-21T09:53:34Z
body: Skipping CI for Draft Pull Request.
If you want CI signal for your change, please convert it to an actual PR.
You can still manually trigger a test run with `/test all`
author_association: CONTRIBUTOR
user: login: netlify[bot]
created_at: 2025-03-21T09:53:50Z
body: ### <span aria-hidden="true">✅</span> Deploy Preview for *kubernetes-sigs-kueue* ready!


|  Name | Link |
|:-:|------------------------|
|<span aria-hidden="true">🔨</span> Latest commit | e316d4d68f5d326c67e3c77bef6d7192da2fc8e5 |
|<span aria-hidden="true">🔍</span> Latest deploy log | https://app.netlify.com/sites/kubernetes-sigs-kueue/deploys/67e3bf68f738d5000879345e |
|<span aria-hidden="true">😎</span> Deploy Preview | [https://deploy-preview-4742--kubernetes-sigs-kueue.netlify.app](https://deploy-preview-4742--kubernetes-sigs-kueue.netlify.app) |
|<span aria-hidden="true">📱</span> Preview on mobile | <details><summary> Toggle QR Code... </summary><br /><br />![QR Code](https://app.netlify.com/qr-code/eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJ1cmwiOiJodHRwczovL2RlcGxveS1wcmV2aWV3LTQ3NDItLWt1YmVybmV0ZXMtc2lncy1rdWV1ZS5uZXRsaWZ5LmFwcCJ9.BUc1Qc-jv3QB13x-AGgerkPT56PfcVgG8rYMYiGB-e4)<br /><br />_Use your smartphone camera to open QR code link._</details> |
---
<!-- [kubernetes-sigs-kueue Preview](https://deploy-preview-4742--kubernetes-sigs-kueue.netlify.app) -->
_To edit notification comments on pull requests, go to your [Netlify site configuration](https://app.netlify.com/sites/kubernetes-sigs-kueue/configuration/notifications#deploy-webhooks)._
author_association: NONE
user: login: mykysha
created_at: 2025-03-21T09:54:18Z
body: /test all
author_association: CONTRIBUTOR
user: login: mykysha
created_at: 2025-03-21T10:24:05Z
body: /cc @mbobrovskyi 
author_association: CONTRIBUTOR
user: login: mbobrovskyi
created_at: 2025-03-21T10:52:01Z
body: /lgtm
Thanks!
author_association: CONTRIBUTOR
user: login: k8s-ci-robot
created_at: 2025-03-21T10:52:09Z
body: LGTM label has been added.  <details>Git tree hash: 790520aad0daf98c98486594c41ae591e12c91e0</details>
author_association: CONTRIBUTOR
user: login: mimowo
created_at: 2025-03-26T10:03:53Z
body: /lgtm
/approve

I think this is actually a small bug. However, it is quite unlikely to hit users, because it would only set the PodSetTopologyRequest, if the TAS-specific annotations were used. Still, I'm leaning to mark it as a bug.
/remove-kind cleanuo
/kind bug

/cherry-pick release-0.11 release-0.10

author_association: CONTRIBUTOR
user: login: k8s-infra-cherrypick-robot
created_at: 2025-03-26T10:03:55Z
body: @mimowo: once the present PR merges, I will cherry-pick it on top of `release-0.11` in a new PR and assign it to you.

<details>

In response to [this](https://github.com/kubernetes-sigs/kueue/pull/4742#issuecomment-2753850828):

>/lgtm
>/approve
>
>I think this is actually a small bug. However, it is quite unlikely to hit users, because it would only set the PodSetTopologyRequest, if the TAS-specific annotations were used. Still, I'm leaning to mark it as a bug.
>/remove-kind cleanuo
>/kind bug
>
>/cherry-pick release-0.11 release-0.10
>


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>
author_association: CONTRIBUTOR
user: login: k8s-ci-robot
created_at: 2025-03-26T10:03:56Z
body: @mimowo: Those labels are not set on the issue: `kind/cleanuo`

<details>

In response to [this](https://github.com/kubernetes-sigs/kueue/pull/4742#issuecomment-2753850828):

>/lgtm
>/approve
>
>I think this is actually a small bug. However, it is quite unlikely to hit users, because it would only set the PodSetTopologyRequest, if the TAS-specific annotations were used. Still, I'm leaning to mark it as a bug.
>/remove-kind cleanuo
>/kind bug
>
>/cherry-pick release-0.11 release-0.10
>


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>
author_association: CONTRIBUTOR
user: login: k8s-ci-robot
created_at: 2025-03-26T10:04:00Z
body: LGTM label has been added.  <details>Git tree hash: 0df8e75754491193500c8bd69b31bd292465516c</details>
author_association: CONTRIBUTOR
user: login: k8s-ci-robot
created_at: 2025-03-26T10:04:02Z
body: [APPROVALNOTIFIER] This PR is **APPROVED**

This pull-request has been approved by: *<a href="https://github.com/kubernetes-sigs/kueue/pull/4742#issuecomment-2753850828" title="Approved">mimowo</a>*, *<a href="https://github.com/kubernetes-sigs/kueue/pull/4742#" title="Author self-approved">mykysha</a>*

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=kubernetes-sigs%2Fkueue).

The pull request process is described [here](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process)

<details >
Needs approval from an approver in each of these files:

- ~~[OWNERS](https://github.com/kubernetes-sigs/kueue/blob/main/OWNERS)~~ [mimowo]

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":[]} -->
author_association: CONTRIBUTOR
user: login: k8s-infra-cherrypick-robot
created_at: 2025-03-26T10:29:14Z
body: @mimowo: new pull request created: #4797

<details>

In response to [this](https://github.com/kubernetes-sigs/kueue/pull/4742#issuecomment-2753850828):

>/lgtm
>/approve
>
>I think this is actually a small bug. However, it is quite unlikely to hit users, because it would only set the PodSetTopologyRequest, if the TAS-specific annotations were used. Still, I'm leaning to mark it as a bug.
>/remove-kind cleanuo
>/kind bug
>
>/cherry-pick release-0.11 release-0.10
>


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>
author_association: CONTRIBUTOR