Repository: kubernetes-sigs/kueue

Issue:
Flaky Integration Test: Scheduler when Preemption is enabled in fairsharing and there are large values of quota and weights Queue can reclaim its nominal quota
<!-- Please use this template while reporting a bug and provide as much info as possible. Not doing so may result in your bug not being addressed in a timely manner. Thanks!

If the matter is security related, please disclose it privately via https://kubernetes.io/security/
-->


**What happened**:

`Scheduler Fair Sharing Suite: [It] Scheduler when Preemption is enabled in fairsharing and there are large values of quota and weights Queue can reclaim its nominal quota [slow]` failed in release-0.12 periodic Job.

```shell
{Timed out after 10.000s.
The function passed to Eventually failed at /home/prow/go/src/kubernetes-sigs/kueue/test/util/util_scheduling.go:67 with:
Not enough workloads evicted
Expected
    <int>: 0
to equal
    <int>: 3 failed [FAILED] Timed out after 10.000s.
The function passed to Eventually failed at /home/prow/go/src/kubernetes-sigs/kueue/test/util/util_scheduling.go:67 with:
Not enough workloads evicted
Expected
    <int>: 0
to equal
    <int>: 3
In [It] at: /home/prow/go/src/kubernetes-sigs/kueue/test/integration/singlecluster/scheduler/fairsharing/fair_sharing_test.go:499 @ 09/13/25 07:51:18.383
}
```

**What you expected to happen**:

No errors.

**How to reproduce it (as minimally and precisely as possible)**:

<img width="3168" height="482" alt="Image" src="https://github.com/user-attachments/assets/7259e200-1e88-4336-b761-da6a3d415555" />

https://prow.k8s.io/view/gs/kubernetes-ci-logs/logs/periodic-kueue-test-integration-release-0-12/1966767439517585408

**Anything else we need to know?**:

**Environment**:
- Kubernetes version (use `kubectl version`):
- Kueue version (use `git describe --tags --dirty --always`):
- Cloud provider or hardware configuration:
- OS (e.g: `cat /etc/os-release`):
- Kernel (e.g. `uname -a`):
- Install tools:
- Others:

Discussion:
@tenzen-y: /kind flake

IIUC, we added this case in https://github.com/kubernetes-sigs/kueue/pull/6617

cc @pajakd @PBundyra @mimowo: https://prow.k8s.io/view/gs/kubernetes-ci-logs/logs/periodic-kueue-test-integration-release-0-12/1969848044652335104@pajakd: /assign@mbobrovskyi: It also happens in main: https://prow.k8s.io/view/gs/kubernetes-ci-logs/pr-logs/pull/kubernetes-sigs_kueue/7122/pull-kueue-test-integration-extended-main/1973390165825032192.@mbobrovskyi: /retitle Flaky Integration Test: Scheduler when Preemption is enabled in fairsharing and there are large values of quota and weights Queue can reclaim its nominal quota@pajakd: I looked at the logs and I think there is the following race condition:
https://github.com/kubernetes-sigs/kueue/blob/0d6a4aee95cb221f172cb82b88d6ebb68f1c57e0/test/integration/singlecluster/scheduler/fairsharing/fair_sharing_test.go#L837-L846

- The last of the 10 workloads in queue "a" is created and reserves the quota.
- Before it gets admitted, the workload in queue "b" gets created. 
- I can see in the logs that the workload in "b" evicts 2 workloads (for some reason the error in the test failure says 0 evictions).

IMO to deflake the test we should wait for the actual admission of the 10 workloads.@mimowo: Is it becuase the metric `ReservingActiveWorkloadsMetric` is bumped before the workloads are actually reserving the quota? 

If this is the case, then it looks like an issue on its own, as we try to bump metrics only once the transition happens. 

Still, the metrics fix may not be trivial and the issue is likely low priority. Finally, testing the metric is not the intention of the test.

So, +1 on the pragmatic flakiness fix.

Solution patch:
diff --git a/internal/mocks/controller/jobframework/interface.go b/internal/mocks/controller/jobframework/interface.go
index ec0ca319b..874865ec4 100644
--- a/internal/mocks/controller/jobframework/interface.go
+++ b/internal/mocks/controller/jobframework/interface.go
@@ -10,6 +10,7 @@
 package mocks
 
 import (
+	context "context"
 	reflect "reflect"
 
 	gomock "go.uber.org/mock/gomock"
@@ -46,9 +47,9 @@ func (m *MockGenericJob) EXPECT() *MockGenericJobMockRecorder {
 }
 
 // Finished mocks base method.
-func (m *MockGenericJob) Finished() (string, bool, bool) {
+func (m *MockGenericJob) Finished(ctx context.Context) (string, bool, bool) {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "Finished")
+	ret := m.ctrl.Call(m, "Finished", ctx)
 	ret0, _ := ret[0].(string)
 	ret1, _ := ret[1].(bool)
 	ret2, _ := ret[2].(bool)
@@ -56,9 +57,9 @@ func (m *MockGenericJob) Finished() (string, bool, bool) {
 }
 
 // Finished indicates an expected call of Finished.
-func (mr *MockGenericJobMockRecorder) Finished() *gomock.Call {
+func (mr *MockGenericJobMockRecorder) Finished(ctx any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Finished", reflect.TypeOf((*MockGenericJob)(nil).Finished))
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Finished", reflect.TypeOf((*MockGenericJob)(nil).Finished), ctx)
 }
 
 // GVK mocks base method.
@@ -118,32 +119,32 @@ func (mr *MockGenericJobMockRecorder) Object() *gomock.Call {
 }
 
 // PodSets mocks base method.
-func (m *MockGenericJob) PodSets() ([]v1beta1.PodSet, error) {
+func (m *MockGenericJob) PodSets(ctx context.Context) ([]v1beta1.PodSet, error) {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "PodSets")
+	ret := m.ctrl.Call(m, "PodSets", ctx)
 	ret0, _ := ret[0].([]v1beta1.PodSet)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
 // PodSets indicates an expected call of PodSets.
-func (mr *MockGenericJobMockRecorder) PodSets() *gomock.Call {
+func (mr *MockGenericJobMockRecorder) PodSets(ctx any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PodSets", reflect.TypeOf((*MockGenericJob)(nil).PodSets))
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PodSets", reflect.TypeOf((*MockGenericJob)(nil).PodSets), ctx)
 }
 
 // PodsReady mocks base method.
-func (m *MockGenericJob) PodsReady() bool {
+func (m *MockGenericJob) PodsReady(ctx context.Context) bool {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "PodsReady")
+	ret := m.ctrl.Call(m, "PodsReady", ctx)
 	ret0, _ := ret[0].(bool)
 	return ret0
 }
 
 // PodsReady indicates an expected call of PodsReady.
-func (mr *MockGenericJobMockRecorder) PodsReady() *gomock.Call {
+func (mr *MockGenericJobMockRecorder) PodsReady(ctx any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PodsReady", reflect.TypeOf((*MockGenericJob)(nil).PodsReady))
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PodsReady", reflect.TypeOf((*MockGenericJob)(nil).PodsReady), ctx)
 }
 
 // RestorePodSetsInfo mocks base method.
@@ -161,17 +162,17 @@ func (mr *MockGenericJobMockRecorder) RestorePodSetsInfo(podSetsInfo any) *gomoc
 }
 
 // RunWithPodSetsInfo mocks base method.
-func (m *MockGenericJob) RunWithPodSetsInfo(podSetsInfo []podset.PodSetInfo) error {
+func (m *MockGenericJob) RunWithPodSetsInfo(ctx context.Context, podSetsInfo []podset.PodSetInfo) error {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "RunWithPodSetsInfo", podSetsInfo)
+	ret := m.ctrl.Call(m, "RunWithPodSetsInfo", ctx, podSetsInfo)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
 // RunWithPodSetsInfo indicates an expected call of RunWithPodSetsInfo.
-func (mr *MockGenericJobMockRecorder) RunWithPodSetsInfo(podSetsInfo any) *gomock.Call {
+func (mr *MockGenericJobMockRecorder) RunWithPodSetsInfo(ctx, podSetsInfo any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RunWithPodSetsInfo", reflect.TypeOf((*MockGenericJob)(nil).RunWithPodSetsInfo), podSetsInfo)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RunWithPodSetsInfo", reflect.TypeOf((*MockGenericJob)(nil).RunWithPodSetsInfo), ctx, podSetsInfo)
 }
 
 // Suspend mocks base method.
@@ -211,33 +212,33 @@ func (m *MockJobWithCustomValidation) EXPECT() *MockJobWithCustomValidationMockR
 }
 
 // ValidateOnCreate mocks base method.
-func (m *MockJobWithCustomValidation) ValidateOnCreate() (field.ErrorList, error) {
+func (m *MockJobWithCustomValidation) ValidateOnCreate(ctx context.Context) (field.ErrorList, error) {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "ValidateOnCreate")
+	ret := m.ctrl.Call(m, "ValidateOnCreate", ctx)
 	ret0, _ := ret[0].(field.ErrorList)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
 // ValidateOnCreate indicates an expected call of ValidateOnCreate.
-func (mr *MockJobWithCustomValidationMockRecorder) ValidateOnCreate() *gomock.Call {
+func (mr *MockJobWithCustomValidationMockRecorder) ValidateOnCreate(ctx any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ValidateOnCreate", reflect.TypeOf((*MockJobWithCustomValidation)(nil).ValidateOnCreate))
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ValidateOnCreate", reflect.TypeOf((*MockJobWithCustomValidation)(nil).ValidateOnCreate), ctx)
 }
 
 // ValidateOnUpdate mocks base method.
-func (m *MockJobWithCustomValidation) ValidateOnUpdate(oldJob jobframework.GenericJob) (field.ErrorList, error) {
+func (m *MockJobWithCustomValidation) ValidateOnUpdate(ctx context.Context, oldJob jobframework.GenericJob) (field.ErrorList, error) {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "ValidateOnUpdate", oldJob)
+	ret := m.ctrl.Call(m, "ValidateOnUpdate", ctx, oldJob)
 	ret0, _ := ret[0].(field.ErrorList)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
 // ValidateOnUpdate indicates an expected call of ValidateOnUpdate.
-func (mr *MockJobWithCustomValidationMockRecorder) ValidateOnUpdate(oldJob any) *gomock.Call {
+func (mr *MockJobWithCustomValidationMockRecorder) ValidateOnUpdate(ctx, oldJob any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ValidateOnUpdate", reflect.TypeOf((*MockJobWithCustomValidation)(nil).ValidateOnUpdate), oldJob)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ValidateOnUpdate", reflect.TypeOf((*MockJobWithCustomValidation)(nil).ValidateOnUpdate), ctx, oldJob)
 }
 
 // MockJobWithManagedBy is a mock of JobWithManagedBy interface.
diff --git a/pkg/controller/jobframework/base_webhook.go b/pkg/controller/jobframework/base_webhook.go
index e813a5fac..8879530f0 100644
--- a/pkg/controller/jobframework/base_webhook.go
+++ b/pkg/controller/jobframework/base_webhook.go
@@ -83,7 +83,7 @@ func (w *BaseWebhook) ValidateCreate(ctx context.Context, obj runtime.Object) (a
 	log.V(5).Info("Validating create")
 	allErrs := ValidateJobOnCreate(job)
 	if jobWithValidation, ok := job.(JobWithCustomValidation); ok {
-		validationErrs, err := jobWithValidation.ValidateOnCreate()
+		validationErrs, err := jobWithValidation.ValidateOnCreate(ctx)
 		if err != nil {
 			return nil, err
 		}
@@ -100,7 +100,7 @@ func (w *BaseWebhook) ValidateUpdate(ctx context.Context, oldObj, newObj runtime
 	log.Info("Validating update")
 	allErrs := ValidateJobOnUpdate(oldJob, newJob, w.Queues.DefaultLocalQueueExist)
 	if jobWithValidation, ok := newJob.(JobWithCustomValidation); ok {
-		validationErrs, err := jobWithValidation.ValidateOnUpdate(oldJob)
+		validationErrs, err := jobWithValidation.ValidateOnUpdate(ctx, oldJob)
 		if err != nil {
 			return nil, err
 		}
diff --git a/pkg/controller/jobframework/interface.go b/pkg/controller/jobframework/interface.go
index ca4733c6b..28293f4c4 100644
--- a/pkg/controller/jobframework/interface.go
+++ b/pkg/controller/jobframework/interface.go
@@ -46,20 +46,20 @@ type GenericJob interface {
 	// Suspend will suspend the job.
 	Suspend()
 	// RunWithPodSetsInfo will inject the node affinity and podSet counts extracting from workload to job and unsuspend it.
-	RunWithPodSetsInfo(podSetsInfo []podset.PodSetInfo) error
+	RunWithPodSetsInfo(ctx context.Context, podSetsInfo []podset.PodSetInfo) error
 	// RestorePodSetsInfo will restore the original node affinity and podSet counts of the job.
 	// Returns whether any change was done.
 	RestorePodSetsInfo(podSetsInfo []podset.PodSetInfo) bool
 	// Finished means whether the job is completed/failed or not,
 	// condition represents the workload finished condition.
 	// Observed generation of the workload is set by the jobframework.
-	Finished() (message string, success, finished bool)
+	Finished(ctx context.Context) (message string, success, finished bool)
 	// PodSets will build workload podSets corresponding to the job.
-	PodSets() ([]kueue.PodSet, error)
+	PodSets(ctx context.Context) ([]kueue.PodSet, error)
 	// IsActive returns true if there are any running pods.
 	IsActive() bool
 	// PodsReady instructs whether job derived pods are all ready now.
-	PodsReady() bool
+	PodsReady(ctx context.Context) bool
 	// GVK returns GVK (Group Version Kind) for the job.
 	GVK() schema.GroupVersionKind
 }
@@ -74,7 +74,7 @@ type JobWithPodLabelSelector interface {
 
 type JobWithReclaimablePods interface {
 	// ReclaimablePods returns the list of reclaimable pods.
-	ReclaimablePods() ([]kueue.ReclaimablePod, error)
+	ReclaimablePods(ctx context.Context) ([]kueue.ReclaimablePod, error)
 }
 
 type StopReason string
@@ -102,7 +102,7 @@ type JobWithFinalize interface {
 // JobWithSkip interface should be implemented by generic jobs,
 // when reconciliation should be skipped depending on the job's state
 type JobWithSkip interface {
-	Skip() bool
+	Skip(ctx context.Context) bool
 }
 
 type JobWithPriorityClass interface {
@@ -114,9 +114,9 @@ type JobWithPriorityClass interface {
 // for Jobs that use BaseWebhook.
 type JobWithCustomValidation interface {
 	// ValidateOnCreate returns list of webhook create validation errors.
-	ValidateOnCreate() (field.ErrorList, error)
+	ValidateOnCreate(ctx context.Context) (field.ErrorList, error)
 	// ValidateOnUpdate returns list of webhook update validation errors.
-	ValidateOnUpdate(oldJob GenericJob) (field.ErrorList, error)
+	ValidateOnUpdate(ctx context.Context, oldJob GenericJob) (field.ErrorList, error)
 }
 
 // ComposableJob interface should be implemented by generic jobs that
diff --git a/pkg/controller/jobframework/reconciler.go b/pkg/controller/jobframework/reconciler.go
index f5cbcd774..8a5823755 100644
--- a/pkg/controller/jobframework/reconciler.go
+++ b/pkg/controller/jobframework/reconciler.go
@@ -266,7 +266,7 @@ func (r *JobReconciler) ReconcileGenericJob(ctx context.Context, req ctrl.Reques
 	}
 
 	if jws, implements := job.(JobWithSkip); implements {
-		if jws.Skip() {
+		if jws.Skip(ctx) {
 			return ctrl.Result{}, nil
 		}
 	}
@@ -357,7 +357,7 @@ func (r *JobReconciler) ReconcileGenericJob(ctx context.Context, req ctrl.Reques
 
 	// if this is a non-toplevel job, suspend the job if its ancestor's workload is not found or not admitted.
 	if !isTopLevelJob {
-		_, _, finished := job.Finished()
+		_, _, finished := job.Finished(ctx)
 		if !finished && !job.IsSuspended() {
 			if ancestorWorkload, err := r.getWorkloadForObject(ctx, ancestorJob); err != nil {
 				log.Error(err, "couldn't get an ancestor job workload")
@@ -433,7 +433,7 @@ func (r *JobReconciler) ReconcileGenericJob(ctx context.Context, req ctrl.Reques
 	}
 
 	// 2. handle job is finished.
-	if message, success, finished := job.Finished(); finished {
+	if message, success, finished := job.Finished(ctx); finished {
 		log.V(3).Info("The workload is already finished")
 		if wl != nil && !apimeta.IsStatusConditionTrue(wl.Status.Conditions, kueue.WorkloadFinished) {
 			reason := kueue.WorkloadFinishedReasonSucceeded
@@ -477,7 +477,7 @@ func (r *JobReconciler) ReconcileGenericJob(ctx context.Context, req ctrl.Reques
 	// 4. update reclaimable counts if implemented by the job
 	if jobRecl, implementsReclaimable := job.(JobWithReclaimablePods); implementsReclaimable {
 		log.V(3).Info("update reclaimable counts if implemented by the job")
-		reclPods, err := jobRecl.ReclaimablePods()
+		reclPods, err := jobRecl.ReclaimablePods(ctx)
 		if err != nil {
 			log.Error(err, "Getting reclaimable pods")
 			return ctrl.Result{}, err
@@ -497,7 +497,7 @@ func (r *JobReconciler) ReconcileGenericJob(ctx context.Context, req ctrl.Reques
 	// handle a job when waitForPodsReady is enabled, and it is the main job
 	if r.waitForPodsReady {
 		log.V(3).Info("Handling a job when waitForPodsReady is enabled")
-		condition := generatePodsReadyCondition(log, job, wl, r.clock)
+		condition := generatePodsReadyCondition(ctx, job, wl, r.clock)
 		if !workload.HasConditionWithTypeAndReason(wl, &condition) {
 			log.V(3).Info("Updating the PodsReady condition", "reason", condition.Reason, "status", condition.Status)
 			apimeta.SetStatusCondition(&wl.Status.Conditions, condition)
@@ -826,7 +826,7 @@ func (r *JobReconciler) ensureOneWorkload(ctx context.Context, job GenericJob, o
 
 	// If workload slicing is enabled for this job, use the slice-based processing path.
 	if workloadSliceEnabled(job) {
-		podSets, err := job.PodSets()
+		podSets, err := job.PodSets(ctx)
 		if err != nil {
 			return nil, fmt.Errorf("failed to retrieve pod sets from job: %w", err)
 		}
@@ -880,7 +880,7 @@ func (r *JobReconciler) ensureOneWorkload(ctx context.Context, job GenericJob, o
 			w = toDelete[0]
 		}
 
-		if _, _, finished := job.Finished(); !finished {
+		if _, _, finished := job.Finished(ctx); !finished {
 			var msg string
 			if w == nil {
 				msg = "Missing Workload; unable to restore pod templates"
@@ -1044,7 +1044,7 @@ func EquivalentToWorkload(ctx context.Context, c client.Client, job GenericJob,
 		return false, nil
 	}
 
-	getPodSets, err := job.PodSets()
+	getPodSets, err := job.PodSets(ctx)
 	if err != nil {
 		return false, err
 	}
@@ -1097,7 +1097,7 @@ func (r *JobReconciler) startJob(ctx context.Context, job GenericJob, object cli
 		}
 	} else {
 		if err := clientutil.Patch(ctx, r.client, object, func() (client.Object, bool, error) {
-			return object, true, job.RunWithPodSetsInfo(info)
+			return object, true, job.RunWithPodSetsInfo(ctx, info)
 		}); err != nil {
 			return err
 		}
@@ -1191,7 +1191,7 @@ func ConstructWorkload(ctx context.Context, c client.Client, job GenericJob, lab
 	log := ctrl.LoggerFrom(ctx)
 	object := job.Object()
 
-	podSets, err := job.PodSets()
+	podSets, err := job.PodSets(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -1381,7 +1381,8 @@ func (r *JobReconciler) ignoreUnretryableError(log logr.Logger, err error) error
 	return err
 }
 
-func generatePodsReadyCondition(log logr.Logger, job GenericJob, wl *kueue.Workload, clock clock.Clock) metav1.Condition {
+func generatePodsReadyCondition(ctx context.Context, job GenericJob, wl *kueue.Workload, clock clock.Clock) metav1.Condition {
+	log := ctrl.LoggerFrom(ctx)
 	const (
 		notReadyMsg           = "Not all pods are ready or succeeded"
 		waitingForRecoveryMsg = "At least one pod has failed, waiting for recovery"
@@ -1397,7 +1398,7 @@ func generatePodsReadyCondition(log logr.Logger, job GenericJob, wl *kueue.Workl
 	}
 
 	podsReadyCond := apimeta.FindStatusCondition(wl.Status.Conditions, kueue.WorkloadPodsReady)
-	podsReady := job.PodsReady()
+	podsReady := job.PodsReady(ctx)
 	log.V(3).Info("Generating PodsReady condition",
 		"Current PodsReady condition", podsReadyCond,
 		"Pods are ready", podsReady)
diff --git a/pkg/controller/jobs/appwrapper/appwrapper_controller.go b/pkg/controller/jobs/appwrapper/appwrapper_controller.go
index c3c7c6fc7..cc3bac8cb 100644
--- a/pkg/controller/jobs/appwrapper/appwrapper_controller.go
+++ b/pkg/controller/jobs/appwrapper/appwrapper_controller.go
@@ -115,10 +115,11 @@ func (j *AppWrapper) GVK() schema.GroupVersionKind {
 	return gvk
 }
 
-func (j *AppWrapper) PodSets() ([]kueue.PodSet, error) {
+func (j *AppWrapper) PodSets(ctx context.Context) ([]kueue.PodSet, error) {
+	log := ctrl.LoggerFrom(ctx)
 	podSpecTemplates, awPodSets, err := awutils.GetComponentPodSpecs((*awv1beta2.AppWrapper)(j))
 	if err != nil {
-		ctrl.Log.Error(err, "Error returned from awutils.GetComponentPodSpecs", "appwrapper", j)
+		log.Error(err, "Error returned from awutils.GetComponentPodSpecs", "appwrapper", j)
 		return nil, err
 	}
 	podSets := make([]kueue.PodSet, len(podSpecTemplates))
@@ -136,7 +137,7 @@ func (j *AppWrapper) PodSets() ([]kueue.PodSet, error) {
 			if count, err := strconv.Atoi(annotation); err == nil {
 				subGroupCount = ptr.To(int32(count))
 			} else {
-				ctrl.Log.Error(err, "Malformed annotation ignored",
+				log.Error(err, "Malformed annotation ignored",
 					"annotationKey", awutils.PodSetAnnotationTASSubGroupCount,
 					"annotationValue", annotation)
 			}
@@ -159,7 +160,7 @@ func (j *AppWrapper) PodSets() ([]kueue.PodSet, error) {
 	return podSets, nil
 }
 
-func (j *AppWrapper) RunWithPodSetsInfo(podSetsInfo []podset.PodSetInfo) error {
+func (j *AppWrapper) RunWithPodSetsInfo(ctx context.Context, podSetsInfo []podset.PodSetInfo) error {
 	awPodSetsInfo := make([]awv1beta2.AppWrapperPodSetInfo, len(podSetsInfo))
 	for idx := range podSetsInfo {
 		awPodSetsInfo[idx].Annotations = podSetsInfo[idx].Annotations
@@ -180,7 +181,7 @@ func (j *AppWrapper) RestorePodSetsInfo(podSetsInfo []podset.PodSetInfo) bool {
 	return awutils.ClearPodSetInfos((*awv1beta2.AppWrapper)(j))
 }
 
-func (j *AppWrapper) Finished() (message string, success, finished bool) {
+func (j *AppWrapper) Finished(ctx context.Context) (message string, success, finished bool) {
 	switch j.Status.Phase {
 	case awv1beta2.AppWrapperSucceeded:
 		return "AppWrapper finished successfully", true, true
@@ -195,7 +196,7 @@ func (j *AppWrapper) Finished() (message string, success, finished bool) {
 	return "", false, false
 }
 
-func (j *AppWrapper) PodsReady() bool {
+func (j *AppWrapper) PodsReady(ctx context.Context) bool {
 	return meta.IsStatusConditionTrue(j.Status.Conditions, string(awv1beta2.PodsReady))
 }
 
diff --git a/pkg/controller/jobs/job/job_controller.go b/pkg/controller/jobs/job/job_controller.go
index 1e8ca4b21..3f9713054 100644
--- a/pkg/controller/jobs/job/job_controller.go
+++ b/pkg/controller/jobs/job/job_controller.go
@@ -210,7 +210,7 @@ func (j *Job) PodLabelSelector() string {
 	return fmt.Sprintf("%s=%s", batchv1.JobNameLabel, j.Name)
 }
 
-func (j *Job) ReclaimablePods() ([]kueue.ReclaimablePod, error) {
+func (j *Job) ReclaimablePods(ctx context.Context) ([]kueue.ReclaimablePod, error) {
 	parallelism := ptr.Deref(j.Spec.Parallelism, 1)
 	if parallelism == 1 || j.Status.Succeeded == 0 {
 		return nil, nil
@@ -243,7 +243,7 @@ func cleanManagedLabels(pt *corev1.PodTemplateSpec) *corev1.PodTemplateSpec {
 	return pt
 }
 
-func (j *Job) PodSets() ([]kueue.PodSet, error) {
+func (j *Job) PodSets(ctx context.Context) ([]kueue.PodSet, error) {
 	podSet := kueue.PodSet{
 		Name:     kueue.DefaultPodSetName,
 		Template: *cleanManagedLabels(j.Spec.Template.DeepCopy()),
@@ -264,7 +264,7 @@ func (j *Job) PodSets() ([]kueue.PodSet, error) {
 	}, nil
 }
 
-func (j *Job) RunWithPodSetsInfo(podSetsInfo []podset.PodSetInfo) error {
+func (j *Job) RunWithPodSetsInfo(ctx context.Context, podSetsInfo []podset.PodSetInfo) error {
 	j.Spec.Suspend = ptr.To(false)
 	if len(podSetsInfo) != 1 {
 		return podset.BadPodSetsInfoLenError(1, len(podSetsInfo))
@@ -305,7 +305,7 @@ func (j *Job) RestorePodSetsInfo(podSetsInfo []podset.PodSetInfo) bool {
 	return changed
 }
 
-func (j *Job) Finished() (message string, success, finished bool) {
+func (j *Job) Finished(ctx context.Context) (message string, success, finished bool) {
 	for _, c := range j.Status.Conditions {
 		if (c.Type == batchv1.JobComplete || c.Type == batchv1.JobFailed) && c.Status == corev1.ConditionTrue {
 			return c.Message, c.Type != batchv1.JobFailed, true
@@ -315,7 +315,7 @@ func (j *Job) Finished() (message string, success, finished bool) {
 	return "", true, false
 }
 
-func (j *Job) PodsReady() bool {
+func (j *Job) PodsReady(ctx context.Context) bool {
 	ready := ptr.Deref(j.Status.Ready, 0)
 	uncountedTerminatedSucceeded := 0
 	if j.Status.UncountedTerminatedPods != nil {
diff --git a/pkg/controller/jobs/job/job_webhook.go b/pkg/controller/jobs/job/job_webhook.go
index 9b110cc10..cdb57bd67 100644
--- a/pkg/controller/jobs/job/job_webhook.go
+++ b/pkg/controller/jobs/job/job_webhook.go
@@ -122,20 +122,20 @@ func (w *JobWebhook) ValidateCreate(ctx context.Context, obj runtime.Object) (ad
 	job := fromObject(obj)
 	log := ctrl.LoggerFrom(ctx).WithName("job-webhook")
 	log.V(5).Info("Validating create")
-	validationErrs, err := w.validateCreate(job)
+	validationErrs, err := w.validateCreate(ctx, job)
 	if err != nil {
 		return nil, err
 	}
 	return nil, validationErrs.ToAggregate()
 }
 
-func (w *JobWebhook) validateCreate(job *Job) (field.ErrorList, error) {
+func (w *JobWebhook) validateCreate(ctx context.Context, job *Job) (field.ErrorList, error) {
 	var allErrs field.ErrorList
 	allErrs = append(allErrs, jobframework.ValidateJobOnCreate(job)...)
 	allErrs = append(allErrs, w.validatePartialAdmissionCreate(job)...)
 	allErrs = append(allErrs, w.validateSyncCompletionCreate(job)...)
 	if features.Enabled(features.TopologyAwareScheduling) {
-		validationErrs, err := w.validateTopologyRequest(job)
+		validationErrs, err := w.validateTopologyRequest(ctx, job)
 		if err != nil {
 			return nil, err
 		}
@@ -182,14 +182,14 @@ func (w *JobWebhook) ValidateUpdate(ctx context.Context, oldObj, newObj runtime.
 	newJob := fromObject(newObj)
 	log := ctrl.LoggerFrom(ctx).WithName("job-webhook")
 	log.V(5).Info("Validating update")
-	validationErrs, err := w.validateUpdate(oldJob, newJob)
+	validationErrs, err := w.validateUpdate(ctx, oldJob, newJob)
 	if err != nil {
 		return nil, err
 	}
 	return nil, validationErrs.ToAggregate()
 }
 
-func (w *JobWebhook) validateUpdate(oldJob, newJob *Job) (field.ErrorList, error) {
+func (w *JobWebhook) validateUpdate(ctx context.Context, oldJob, newJob *Job) (field.ErrorList, error) {
 	var allErrs field.ErrorList
 	allErrs = append(allErrs, jobframework.ValidateJobOnCreate(newJob)...)
 	if newJob.Annotations[JobMinParallelismAnnotation] != oldJob.Annotations[JobMinParallelismAnnotation] {
@@ -199,7 +199,7 @@ func (w *JobWebhook) validateUpdate(oldJob, newJob *Job) (field.ErrorList, error
 	allErrs = append(allErrs, jobframework.ValidateJobOnUpdate(oldJob, newJob, w.queues.DefaultLocalQueueExist)...)
 	allErrs = append(allErrs, validatePartialAdmissionUpdate(oldJob, newJob)...)
 	if features.Enabled(features.TopologyAwareScheduling) {
-		validationErrs, err := w.validateTopologyRequest(newJob)
+		validationErrs, err := w.validateTopologyRequest(ctx, newJob)
 		if err != nil {
 			return nil, err
 		}
@@ -221,13 +221,13 @@ func validatePartialAdmissionUpdate(oldJob, newJob *Job) field.ErrorList {
 	return allErrs
 }
 
-func (w *JobWebhook) validateTopologyRequest(job *Job) (field.ErrorList, error) {
+func (w *JobWebhook) validateTopologyRequest(ctx context.Context, job *Job) (field.ErrorList, error) {
 	validationErrs := jobframework.ValidateTASPodSetRequest(replicaMetaPath, &job.Spec.Template.ObjectMeta)
 	if validationErrs != nil {
 		return validationErrs, nil
 	}
 
-	podSets, err := job.PodSets()
+	podSets, err := job.PodSets(ctx)
 	if err != nil {
 		return nil, err
 	}
diff --git a/pkg/controller/jobs/jobset/jobset_controller.go b/pkg/controller/jobs/jobset/jobset_controller.go
index 2dbc2bf86..6745d7bbd 100644
--- a/pkg/controller/jobs/jobset/jobset_controller.go
+++ b/pkg/controller/jobs/jobset/jobset_controller.go
@@ -111,7 +111,7 @@ func (j *JobSet) PodLabelSelector() string {
 	return fmt.Sprintf("%s=%s", jobsetapi.JobSetNameKey, j.Name)
 }
 
-func (j *JobSet) PodSets() ([]kueue.PodSet, error) {
+func (j *JobSet) PodSets(ctx context.Context) ([]kueue.PodSet, error) {
 	podSets := make([]kueue.PodSet, len(j.Spec.ReplicatedJobs))
 	for index, replicatedJob := range j.Spec.ReplicatedJobs {
 		podSets[index] = kueue.PodSet{
@@ -134,7 +134,7 @@ func (j *JobSet) PodSets() ([]kueue.PodSet, error) {
 	return podSets, nil
 }
 
-func (j *JobSet) RunWithPodSetsInfo(podSetsInfo []podset.PodSetInfo) error {
+func (j *JobSet) RunWithPodSetsInfo(ctx context.Context, podSetsInfo []podset.PodSetInfo) error {
 	j.Spec.Suspend = ptr.To(false)
 	if len(podSetsInfo) != len(j.Spec.ReplicatedJobs) {
 		return podset.BadPodSetsInfoLenError(len(j.Spec.ReplicatedJobs), len(podSetsInfo))
@@ -165,7 +165,7 @@ func (j *JobSet) RestorePodSetsInfo(podSetsInfo []podset.PodSetInfo) bool {
 	return changed
 }
 
-func (j *JobSet) Finished() (message string, success, finished bool) {
+func (j *JobSet) Finished(ctx context.Context) (message string, success, finished bool) {
 	if c := apimeta.FindStatusCondition(j.Status.Conditions, string(jobsetapi.JobSetCompleted)); c != nil && c.Status == metav1.ConditionTrue {
 		return c.Message, true, true
 	}
@@ -175,7 +175,7 @@ func (j *JobSet) Finished() (message string, success, finished bool) {
 	return message, success, false
 }
 
-func (j *JobSet) PodsReady() bool {
+func (j *JobSet) PodsReady(ctx context.Context) bool {
 	var replicas int32
 	for _, replicatedJob := range j.Spec.ReplicatedJobs {
 		replicas += replicatedJob.Replicas
@@ -187,7 +187,7 @@ func (j *JobSet) PodsReady() bool {
 	return replicas == readyReplicas
 }
 
-func (j *JobSet) ReclaimablePods() ([]kueue.ReclaimablePod, error) {
+func (j *JobSet) ReclaimablePods(ctx context.Context) ([]kueue.ReclaimablePod, error) {
 	if len(j.Status.ReplicatedJobsStatus) == 0 {
 		return nil, nil
 	}
diff --git a/pkg/controller/jobs/jobset/jobset_webhook.go b/pkg/controller/jobs/jobset/jobset_webhook.go
index e409cfec2..9fa8531d4 100644
--- a/pkg/controller/jobs/jobset/jobset_webhook.go
+++ b/pkg/controller/jobs/jobset/jobset_webhook.go
@@ -95,7 +95,7 @@ func (w *JobSetWebhook) ValidateCreate(ctx context.Context, obj runtime.Object)
 	jobSet := fromObject(obj)
 	log := ctrl.LoggerFrom(ctx).WithName("jobset-webhook")
 	log.Info("Validating create")
-	validationErrs, err := w.validateCreate(jobSet)
+	validationErrs, err := w.validateCreate(ctx, jobSet)
 	if err != nil {
 		return nil, err
 	}
@@ -108,17 +108,17 @@ func (w *JobSetWebhook) ValidateUpdate(ctx context.Context, oldObj, newObj runti
 	newJobSet := fromObject(newObj)
 	log := ctrl.LoggerFrom(ctx).WithName("jobset-webhook")
 	log.Info("Validating update")
-	validationErrs, err := w.validateUpdate(oldJobSet, newJobSet)
+	validationErrs, err := w.validateUpdate(ctx, oldJobSet, newJobSet)
 	if err != nil {
 		return nil, err
 	}
 	return nil, validationErrs.ToAggregate()
 }
 
-func (w *JobSetWebhook) validateUpdate(oldJob, newJob *JobSet) (field.ErrorList, error) {
+func (w *JobSetWebhook) validateUpdate(ctx context.Context, oldJob, newJob *JobSet) (field.ErrorList, error) {
 	var allErrs field.ErrorList
 	allErrs = append(allErrs, jobframework.ValidateJobOnUpdate(oldJob, newJob, w.queues.DefaultLocalQueueExist)...)
-	validationErrs, err := w.validateCreate(newJob)
+	validationErrs, err := w.validateCreate(ctx, newJob)
 	if err != nil {
 		return nil, err
 	}
@@ -126,11 +126,11 @@ func (w *JobSetWebhook) validateUpdate(oldJob, newJob *JobSet) (field.ErrorList,
 	return allErrs, nil
 }
 
-func (w *JobSetWebhook) validateCreate(jobSet *JobSet) (field.ErrorList, error) {
+func (w *JobSetWebhook) validateCreate(ctx context.Context, jobSet *JobSet) (field.ErrorList, error) {
 	var allErrs field.ErrorList
 	allErrs = append(allErrs, jobframework.ValidateJobOnCreate(jobSet)...)
 	if features.Enabled(features.TopologyAwareScheduling) {
-		validationErrs, err := w.validateTopologyRequest(jobSet)
+		validationErrs, err := w.validateTopologyRequest(ctx, jobSet)
 		if err != nil {
 			return nil, err
 		}
@@ -139,10 +139,10 @@ func (w *JobSetWebhook) validateCreate(jobSet *JobSet) (field.ErrorList, error)
 	return allErrs, nil
 }
 
-func (w *JobSetWebhook) validateTopologyRequest(jobSet *JobSet) (field.ErrorList, error) {
+func (w *JobSetWebhook) validateTopologyRequest(ctx context.Context, jobSet *JobSet) (field.ErrorList, error) {
 	var allErrs field.ErrorList
 
-	podSets, podSetsErr := jobSet.PodSets()
+	podSets, podSetsErr := jobSet.PodSets(ctx)
 
 	for i, rj := range jobSet.Spec.ReplicatedJobs {
 		replicaMetaPath := replicatedJobsPath.Index(i).Child("template", "metadata")
diff --git a/pkg/controller/jobs/kubeflow/kubeflowjob/kubeflowjob_controller.go b/pkg/controller/jobs/kubeflow/kubeflowjob/kubeflowjob_controller.go
index 73e1ecb96..712919dd6 100644
--- a/pkg/controller/jobs/kubeflow/kubeflowjob/kubeflowjob_controller.go
+++ b/pkg/controller/jobs/kubeflow/kubeflowjob/kubeflowjob_controller.go
@@ -17,6 +17,8 @@ limitations under the License.
 package kubeflowjob
 
 import (
+	"context"
+
 	kftraining "github.com/kubeflow/training-operator/pkg/apis/kubeflow.org/v1"
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/runtime/schema"
@@ -52,7 +54,7 @@ func (j *KubeflowJob) Suspend() {
 	j.KFJobControl.RunPolicy().Suspend = ptr.To(true)
 }
 
-func (j *KubeflowJob) RunWithPodSetsInfo(podSetsInfo []podset.PodSetInfo) error {
+func (j *KubeflowJob) RunWithPodSetsInfo(ctx context.Context, podSetsInfo []podset.PodSetInfo) error {
 	j.KFJobControl.RunPolicy().Suspend = ptr.To(false)
 	orderedReplicaTypes := j.OrderedReplicaTypes()
 
@@ -83,7 +85,7 @@ func (j *KubeflowJob) RestorePodSetsInfo(podSetsInfo []podset.PodSetInfo) bool {
 	return changed
 }
 
-func (j *KubeflowJob) Finished() (message string, success, finished bool) {
+func (j *KubeflowJob) Finished(ctx context.Context) (message string, success, finished bool) {
 	if j.KFJobControl.JobStatus() == nil {
 		return "", false, false
 	}
@@ -97,7 +99,7 @@ func (j *KubeflowJob) Finished() (message string, success, finished bool) {
 	return "", true, false
 }
 
-func (j *KubeflowJob) PodSets() ([]kueue.PodSet, error) {
+func (j *KubeflowJob) PodSets(ctx context.Context) ([]kueue.PodSet, error) {
 	replicaTypes := j.OrderedReplicaTypes()
 	podSets := make([]kueue.PodSet, len(replicaTypes))
 	for index, replicaType := range replicaTypes {
@@ -128,7 +130,7 @@ func (j *KubeflowJob) IsActive() bool {
 	return false
 }
 
-func (j *KubeflowJob) PodsReady() bool {
+func (j *KubeflowJob) PodsReady(ctx context.Context) bool {
 	for _, c := range j.KFJobControl.JobStatus().Conditions {
 		if c.Type == kftraining.JobRunning && c.Status == corev1.ConditionTrue {
 			return true
@@ -180,12 +182,12 @@ func (j *KubeflowJob) OrderedReplicaTypes() []kftraining.ReplicaType {
 	return result
 }
 
-func (j *KubeflowJob) ValidateOnCreate() (field.ErrorList, error) {
+func (j *KubeflowJob) ValidateOnCreate(ctx context.Context) (field.ErrorList, error) {
 	if !features.Enabled(features.TopologyAwareScheduling) {
 		return nil, nil
 	}
 
-	podSets, podSetsErr := j.PodSets()
+	podSets, podSetsErr := j.PodSets(ctx)
 
 	var allErrs field.ErrorList
 	replicaTypes := j.OrderedReplicaTypes()
@@ -215,8 +217,8 @@ func (j *KubeflowJob) ValidateOnCreate() (field.ErrorList, error) {
 	return nil, podSetsErr
 }
 
-func (j *KubeflowJob) ValidateOnUpdate(_ jobframework.GenericJob) (field.ErrorList, error) {
-	return j.ValidateOnCreate()
+func (j *KubeflowJob) ValidateOnUpdate(ctx context.Context, _ jobframework.GenericJob) (field.ErrorList, error) {
+	return j.ValidateOnCreate(ctx)
 }
 
 func podsCount(replicaSpecs map[kftraining.ReplicaType]*kftraining.ReplicaSpec, replicaType kftraining.ReplicaType) int32 {
diff --git a/pkg/controller/jobs/mpijob/mpijob_controller.go b/pkg/controller/jobs/mpijob/mpijob_controller.go
index c752c310d..bfbffd882 100644
--- a/pkg/controller/jobs/mpijob/mpijob_controller.go
+++ b/pkg/controller/jobs/mpijob/mpijob_controller.go
@@ -109,7 +109,7 @@ func (j *MPIJob) PodLabelSelector() string {
 	return fmt.Sprintf("%s=%s,%s=%s", kfmpi.JobNameLabel, j.Name, kfmpi.OperatorNameLabel, kfmpi.OperatorName)
 }
 
-func (j *MPIJob) PodSets() ([]kueue.PodSet, error) {
+func (j *MPIJob) PodSets(ctx context.Context) ([]kueue.PodSet, error) {
 	replicaTypes := orderedReplicaTypes(&j.Spec)
 	podSets := make([]kueue.PodSet, len(replicaTypes))
 	for index, mpiReplicaType := range replicaTypes {
@@ -131,7 +131,7 @@ func (j *MPIJob) PodSets() ([]kueue.PodSet, error) {
 	return podSets, nil
 }
 
-func (j *MPIJob) RunWithPodSetsInfo(podSetsInfo []podset.PodSetInfo) error {
+func (j *MPIJob) RunWithPodSetsInfo(ctx context.Context, podSetsInfo []podset.PodSetInfo) error {
 	j.Spec.RunPolicy.Suspend = ptr.To(false)
 	orderedReplicaTypes := orderedReplicaTypes(&j.Spec)
 
@@ -163,7 +163,7 @@ func (j *MPIJob) RestorePodSetsInfo(podSetsInfo []podset.PodSetInfo) bool {
 	return changed
 }
 
-func (j *MPIJob) Finished() (message string, success, finished bool) {
+func (j *MPIJob) Finished(ctx context.Context) (message string, success, finished bool) {
 	for _, c := range j.Status.Conditions {
 		if (c.Type == kfmpi.JobSucceeded || c.Type == kfmpi.JobFailed) && c.Status == corev1.ConditionTrue {
 			return c.Message, c.Type != kfmpi.JobFailed, true
@@ -191,7 +191,7 @@ func (j *MPIJob) PriorityClass() string {
 	return ""
 }
 
-func (j *MPIJob) PodsReady() bool {
+func (j *MPIJob) PodsReady(ctx context.Context) bool {
 	for _, c := range j.Status.Conditions {
 		if c.Type == kfmpi.JobRunning && c.Status == corev1.ConditionTrue {
 			return true
diff --git a/pkg/controller/jobs/mpijob/mpijob_webhook.go b/pkg/controller/jobs/mpijob/mpijob_webhook.go
index c82edec06..6c4187be7 100644
--- a/pkg/controller/jobs/mpijob/mpijob_webhook.go
+++ b/pkg/controller/jobs/mpijob/mpijob_webhook.go
@@ -100,7 +100,7 @@ func (w *MpiJobWebhook) ValidateCreate(ctx context.Context, obj runtime.Object)
 	mpiJob := fromObject(obj)
 	log := ctrl.LoggerFrom(ctx).WithName("mpijob-webhook")
 	log.Info("Validating create")
-	validationErrs, err := w.validateCommon(mpiJob)
+	validationErrs, err := w.validateCommon(ctx, mpiJob)
 	if err != nil {
 		return nil, err
 	}
@@ -117,7 +117,7 @@ func (w *MpiJobWebhook) ValidateUpdate(ctx context.Context, oldObj, newObj runti
 	log := ctrl.LoggerFrom(ctx).WithName("mpijob-webhook")
 	log.Info("Validating update")
 	allErrs := jobframework.ValidateJobOnUpdate(oldMpiJob, newMpiJob, w.queues.DefaultLocalQueueExist)
-	validationErrs, err := w.validateCommon(newMpiJob)
+	validationErrs, err := w.validateCommon(ctx, newMpiJob)
 	if err != nil {
 		return nil, err
 	}
@@ -133,11 +133,11 @@ func (w *MpiJobWebhook) ValidateDelete(context.Context, runtime.Object) (admissi
 	return nil, nil
 }
 
-func (w *MpiJobWebhook) validateCommon(mpiJob *MPIJob) (field.ErrorList, error) {
+func (w *MpiJobWebhook) validateCommon(ctx context.Context, mpiJob *MPIJob) (field.ErrorList, error) {
 	var allErrs field.ErrorList
 	allErrs = jobframework.ValidateJobOnCreate(mpiJob)
 	if features.Enabled(features.TopologyAwareScheduling) {
-		validationErrs, err := w.validateTopologyRequest(mpiJob)
+		validationErrs, err := w.validateTopologyRequest(ctx, mpiJob)
 		if err != nil {
 			return nil, err
 		}
@@ -146,10 +146,10 @@ func (w *MpiJobWebhook) validateCommon(mpiJob *MPIJob) (field.ErrorList, error)
 	return allErrs, nil
 }
 
-func (w *MpiJobWebhook) validateTopologyRequest(mpiJob *MPIJob) (field.ErrorList, error) {
+func (w *MpiJobWebhook) validateTopologyRequest(ctx context.Context, mpiJob *MPIJob) (field.ErrorList, error) {
 	var allErrs field.ErrorList
 
-	podSets, podSetsErr := mpiJob.PodSets()
+	podSets, podSetsErr := mpiJob.PodSets(ctx)
 
 	for replicaType, replicaSpec := range mpiJob.Spec.MPIReplicaSpecs {
 		replicaMetaPath := mpiReplicaSpecsPath.Key(string(replicaType)).Child("template", "metadata")
diff --git a/pkg/controller/jobs/pod/pod_controller.go b/pkg/controller/jobs/pod/pod_controller.go
index c635bf0b2..6291b58fe 100644
--- a/pkg/controller/jobs/pod/pod_controller.go
+++ b/pkg/controller/jobs/pod/pod_controller.go
@@ -321,7 +321,7 @@ func (p *Pod) IsTopLevel() bool {
 }
 
 // RunWithPodSetsInfo will inject the node affinity and podSet counts extracting from workload to job and unsuspend it.
-func (p *Pod) RunWithPodSetsInfo(_ []podset.PodSetInfo) error {
+func (p *Pod) RunWithPodSetsInfo(_ context.Context, _ []podset.PodSetInfo) error {
 	// Not implemented because this is not called when JobWithCustomRun is implemented.
 	return errors.New("RunWithPodSetsInfo is not implemented for the Pod object")
 }
@@ -334,11 +334,11 @@ func (p *Pod) RestorePodSetsInfo(_ []podset.PodSetInfo) bool {
 
 // Finished means whether the job is completed/failed or not,
 // condition represents the workload finished condition.
-func (p *Pod) Finished() (message string, success, finished bool) {
+func (p *Pod) Finished(ctx context.Context) (message string, success, finished bool) {
 	if p.isServing() {
 		return "", true, false
 	}
-
+	log := ctrl.LoggerFrom(ctx)
 	finished = true
 	success = true
 
@@ -354,7 +354,7 @@ func (p *Pod) Finished() (message string, success, finished bool) {
 
 	groupTotalCount, err := p.groupTotalCount()
 	if err != nil {
-		ctrl.Log.V(2).Error(err, "failed to check if pod group is finished")
+		log.V(2).Error(err, "failed to check if pod group is finished")
 		message = "failed to check if pod group is finished"
 		return message, success, false
 	}
@@ -380,7 +380,7 @@ func (p *Pod) Finished() (message string, success, finished bool) {
 }
 
 // PodSets will build workload podSets corresponding to the job.
-func (p *Pod) PodSets() ([]kueue.PodSet, error) {
+func (p *Pod) PodSets(ctx context.Context) ([]kueue.PodSet, error) {
 	if !p.isGroup {
 		return constructPodSets(&p.pod)
 	} else {
@@ -435,7 +435,7 @@ func hasPodReadyTrue(conds []corev1.PodCondition) bool {
 }
 
 // PodsReady instructs whether job derived pods are all ready now.
-func (p *Pod) PodsReady() bool {
+func (p *Pod) PodsReady(ctx context.Context) bool {
 	if !p.isGroup {
 		return hasPodReadyTrue(p.pod.Status.Conditions)
 	}
@@ -560,14 +560,16 @@ func (p *Pod) Finalize(ctx context.Context, c client.Client) error {
 	})
 }
 
-func (p *Pod) Skip() bool {
+func (p *Pod) Skip(ctx context.Context) bool {
+	log := ctrl.LoggerFrom(ctx)
 	// Skip pod reconciliation, if pod is found, and it's managed label is not set or incorrect.
 	if v, ok := p.pod.GetLabels()[constants.ManagedByKueueLabelKey]; p.isFound && (!ok || v != constants.ManagedByKueueLabelValue) {
+		log.V(3).Info("Skipping pod, not managed by Kueue", constants.ManagedByKueueLabelKey, v, "labelSet", ok)
 		return true
 	}
 	if jobframework.HasImplicitlyEnabledFramework(p.pod.GroupVersionKind()) &&
 		p.pod.GetAnnotations()[podconstants.SuspendedByParentAnnotation] == "" {
-		ctrl.Log.V(3).Info("Pod Integration was implicitly enabled but object lacks parent annotation, skipping")
+		log.V(3).Info("Pod Integration was implicitly enabled but object lacks parent annotation, skipping")
 		return true
 	}
 	return false
@@ -1061,7 +1063,7 @@ func (p *Pod) ConstructComposableWorkload(ctx context.Context, c client.Client,
 		p.list.Items = activePods[:len(activePods)-excessPodsCount]
 	}
 
-	podSets, err := p.PodSets()
+	podSets, err := p.PodSets(ctx)
 	if err != nil {
 		if jobframework.IsUnretryableError(err) {
 			r.Eventf(p.Object(), corev1.EventTypeWarning, jobframework.ReasonErrWorkloadCompose, err.Error())
@@ -1271,7 +1273,7 @@ func (p *Pod) isReclaimable() bool {
 	return p.isGroup && !p.isServing()
 }
 
-func (p *Pod) ReclaimablePods() ([]kueue.ReclaimablePod, error) {
+func (p *Pod) ReclaimablePods(ctx context.Context) ([]kueue.ReclaimablePod, error) {
 	if !p.isReclaimable() {
 		return []kueue.ReclaimablePod{}, nil
 	}
diff --git a/pkg/controller/jobs/raycluster/raycluster_controller.go b/pkg/controller/jobs/raycluster/raycluster_controller.go
index baf700805..dc3c2f02c 100644
--- a/pkg/controller/jobs/raycluster/raycluster_controller.go
+++ b/pkg/controller/jobs/raycluster/raycluster_controller.go
@@ -101,7 +101,7 @@ func (j *RayCluster) PodLabelSelector() string {
 	return fmt.Sprintf("%s=%s", rayutils.RayClusterLabelKey, j.Name)
 }
 
-func (j *RayCluster) PodSets() ([]kueue.PodSet, error) {
+func (j *RayCluster) PodSets(ctx context.Context) ([]kueue.PodSet, error) {
 	// len = workerGroups + head
 	podSets := make([]kueue.PodSet, len(j.Spec.WorkerGroupSpecs)+1)
 
@@ -148,7 +148,7 @@ func (j *RayCluster) PodSets() ([]kueue.PodSet, error) {
 	return podSets, nil
 }
 
-func (j *RayCluster) RunWithPodSetsInfo(podSetsInfo []podset.PodSetInfo) error {
+func (j *RayCluster) RunWithPodSetsInfo(ctx context.Context, podSetsInfo []podset.PodSetInfo) error {
 	expectedLen := len(j.Spec.WorkerGroupSpecs) + 1
 	if len(podSetsInfo) != expectedLen {
 		return podset.BadPodSetsInfoLenError(expectedLen, len(podSetsInfo))
@@ -193,12 +193,12 @@ func (j *RayCluster) RestorePodSetsInfo(podSetsInfo []podset.PodSetInfo) bool {
 	return changed
 }
 
-func (j *RayCluster) Finished() (message string, success, finished bool) {
+func (j *RayCluster) Finished(ctx context.Context) (message string, success, finished bool) {
 	// Technically a RayCluster is never "finished"
 	return j.Status.Reason, j.Status.State != rayv1.Failed, false
 }
 
-func (j *RayCluster) PodsReady() bool {
+func (j *RayCluster) PodsReady(ctx context.Context) bool {
 	return j.Status.State == rayv1.Ready
 }
 
diff --git a/pkg/controller/jobs/raycluster/raycluster_webhook.go b/pkg/controller/jobs/raycluster/raycluster_webhook.go
index 071478ed8..31e4bcd67 100644
--- a/pkg/controller/jobs/raycluster/raycluster_webhook.go
+++ b/pkg/controller/jobs/raycluster/raycluster_webhook.go
@@ -117,7 +117,7 @@ func (w *RayClusterWebhook) ValidateCreate(ctx context.Context, obj runtime.Obje
 	job := obj.(*rayv1.RayCluster)
 	log := ctrl.LoggerFrom(ctx).WithName("raycluster-webhook")
 	log.V(10).Info("Validating create")
-	validationErrs, err := w.validateCreate(job)
+	validationErrs, err := w.validateCreate(ctx, job)
 	if err != nil {
 		return nil, err
 	}
@@ -129,7 +129,7 @@ func isAnElasticJob(job *rayv1.RayCluster) bool {
 	return features.Enabled(features.ElasticJobsViaWorkloadSlices) && workloadslicing.Enabled(job.GetObjectMeta())
 }
 
-func (w *RayClusterWebhook) validateCreate(job *rayv1.RayCluster) (field.ErrorList, error) {
+func (w *RayClusterWebhook) validateCreate(ctx context.Context, job *rayv1.RayCluster) (field.ErrorList, error) {
 	var allErrors field.ErrorList
 	kueueJob := (*RayCluster)(job)
 
@@ -159,7 +159,7 @@ func (w *RayClusterWebhook) validateCreate(job *rayv1.RayCluster) (field.ErrorLi
 
 	allErrors = append(allErrors, jobframework.ValidateJobOnCreate(kueueJob)...)
 	if features.Enabled(features.TopologyAwareScheduling) {
-		validationErrs, err := w.validateTopologyRequest(kueueJob)
+		validationErrs, err := w.validateTopologyRequest(ctx, kueueJob)
 		if err != nil {
 			return nil, err
 		}
@@ -193,10 +193,10 @@ func validateElasticJob(job *rayv1.RayCluster) field.ErrorList {
 	return allErrors
 }
 
-func (w *RayClusterWebhook) validateTopologyRequest(rayJob *RayCluster) (field.ErrorList, error) {
+func (w *RayClusterWebhook) validateTopologyRequest(ctx context.Context, rayJob *RayCluster) (field.ErrorList, error) {
 	var allErrs field.ErrorList
 
-	podSets, podSetsErr := rayJob.PodSets()
+	podSets, podSetsErr := rayJob.PodSets(ctx)
 
 	allErrs = append(allErrs, jobframework.ValidateTASPodSetRequest(headGroupMetaPath, &rayJob.Spec.HeadGroupSpec.Template.ObjectMeta)...)
 
@@ -232,7 +232,7 @@ func (w *RayClusterWebhook) ValidateUpdate(ctx context.Context, oldObj, newObj r
 	if w.manageJobsWithoutQueueName || jobframework.QueueName((*RayCluster)(newJob)) != "" {
 		log.Info("Validating update")
 		allErrors := jobframework.ValidateJobOnUpdate((*RayCluster)(oldJob), (*RayCluster)(newJob), w.queues.DefaultLocalQueueExist)
-		validationErrs, err := w.validateCreate(newJob)
+		validationErrs, err := w.validateCreate(ctx, newJob)
 		if err != nil {
 			return nil, err
 		}
diff --git a/pkg/controller/jobs/rayjob/rayjob_controller.go b/pkg/controller/jobs/rayjob/rayjob_controller.go
index 2c2f39fd0..4e9e3303e 100644
--- a/pkg/controller/jobs/rayjob/rayjob_controller.go
+++ b/pkg/controller/jobs/rayjob/rayjob_controller.go
@@ -111,7 +111,7 @@ func (j *RayJob) PodLabelSelector() string {
 	return ""
 }
 
-func (j *RayJob) PodSets() ([]kueue.PodSet, error) {
+func (j *RayJob) PodSets(ctx context.Context) ([]kueue.PodSet, error) {
 	podSets := make([]kueue.PodSet, 0)
 
 	// head
@@ -178,7 +178,7 @@ func (j *RayJob) PodSets() ([]kueue.PodSet, error) {
 	return podSets, nil
 }
 
-func (j *RayJob) RunWithPodSetsInfo(podSetsInfo []podset.PodSetInfo) error {
+func (j *RayJob) RunWithPodSetsInfo(ctx context.Context, podSetsInfo []podset.PodSetInfo) error {
 	expectedLen := len(j.Spec.RayClusterSpec.WorkerGroupSpecs) + 1
 	if j.Spec.SubmissionMode == rayv1.K8sJobMode {
 		expectedLen++
@@ -249,14 +249,14 @@ func (j *RayJob) RestorePodSetsInfo(podSetsInfo []podset.PodSetInfo) bool {
 	return changed
 }
 
-func (j *RayJob) Finished() (message string, success, finished bool) {
+func (j *RayJob) Finished(ctx context.Context) (message string, success, finished bool) {
 	message = j.Status.Message
 	success = j.Status.JobStatus == rayv1.JobStatusSucceeded
 	finished = j.Status.JobDeploymentStatus == rayv1.JobDeploymentStatusFailed || j.Status.JobDeploymentStatus == rayv1.JobDeploymentStatusComplete
 	return message, success, finished
 }
 
-func (j *RayJob) PodsReady() bool {
+func (j *RayJob) PodsReady(ctx context.Context) bool {
 	return j.Status.RayClusterStatus.State == rayv1.Ready
 }
 
diff --git a/pkg/controller/jobs/rayjob/rayjob_webhook.go b/pkg/controller/jobs/rayjob/rayjob_webhook.go
index e9f3c0658..0bc95c2c6 100644
--- a/pkg/controller/jobs/rayjob/rayjob_webhook.go
+++ b/pkg/controller/jobs/rayjob/rayjob_webhook.go
@@ -96,14 +96,14 @@ func (w *RayJobWebhook) ValidateCreate(ctx context.Context, obj runtime.Object)
 	job := obj.(*rayv1.RayJob)
 	log := ctrl.LoggerFrom(ctx).WithName("rayjob-webhook")
 	log.Info("Validating create")
-	validationErrs, err := w.validateCreate(job)
+	validationErrs, err := w.validateCreate(ctx, job)
 	if err != nil {
 		return nil, err
 	}
 	return nil, validationErrs.ToAggregate()
 }
 
-func (w *RayJobWebhook) validateCreate(job *rayv1.RayJob) (field.ErrorList, error) {
+func (w *RayJobWebhook) validateCreate(ctx context.Context, job *rayv1.RayJob) (field.ErrorList, error) {
 	var allErrors field.ErrorList
 	kueueJob := (*RayJob)(job)
 
@@ -144,7 +144,7 @@ func (w *RayJobWebhook) validateCreate(job *rayv1.RayJob) (field.ErrorList, erro
 
 	allErrors = append(allErrors, jobframework.ValidateJobOnCreate(kueueJob)...)
 	if features.Enabled(features.TopologyAwareScheduling) {
-		validationErrs, err := w.validateTopologyRequest(job)
+		validationErrs, err := w.validateTopologyRequest(ctx, job)
 		if err != nil {
 			return nil, err
 		}
@@ -154,13 +154,13 @@ func (w *RayJobWebhook) validateCreate(job *rayv1.RayJob) (field.ErrorList, erro
 	return allErrors, nil
 }
 
-func (w *RayJobWebhook) validateTopologyRequest(rayJob *rayv1.RayJob) (field.ErrorList, error) {
+func (w *RayJobWebhook) validateTopologyRequest(ctx context.Context, rayJob *rayv1.RayJob) (field.ErrorList, error) {
 	var allErrs field.ErrorList
 	if rayJob.Spec.RayClusterSpec == nil {
 		return allErrs, nil
 	}
 
-	podSets, podSetsErr := (*RayJob)(rayJob).PodSets()
+	podSets, podSetsErr := (*RayJob)(rayJob).PodSets(ctx)
 
 	allErrs = append(allErrs, jobframework.ValidateTASPodSetRequest(headGroupMetaPath, &rayJob.Spec.RayClusterSpec.HeadGroupSpec.Template.ObjectMeta)...)
 
@@ -196,7 +196,7 @@ func (w *RayJobWebhook) ValidateUpdate(ctx context.Context, oldObj, newObj runti
 	if w.manageJobsWithoutQueueName || jobframework.QueueName((*RayJob)(newJob)) != "" {
 		log.Info("Validating update")
 		allErrors := jobframework.ValidateJobOnUpdate((*RayJob)(oldJob), (*RayJob)(newJob), w.queues.DefaultLocalQueueExist)
-		validationErrs, err := w.validateCreate(newJob)
+		validationErrs, err := w.validateCreate(ctx, newJob)
 		if err != nil {
 			return nil, err
 		}
diff --git a/pkg/controller/jobs/trainjob/trainjob_controller.go b/pkg/controller/jobs/trainjob/trainjob_controller.go
index 74b1bfcf6..61232b1a7 100644
--- a/pkg/controller/jobs/trainjob/trainjob_controller.go
+++ b/pkg/controller/jobs/trainjob/trainjob_controller.go
@@ -81,7 +81,6 @@ func init() {
 // +kubebuilder:rbac:groups=trainer.kubeflow.org,resources=clustertrainingruntimes,verbs=get;list;watch
 
 type trainJobReconciler struct {
-	ctx              context.Context
 	jr               *jobframework.JobReconciler
 	runtimes         map[string]kftrainerruntime.Runtime
 	client           client.Client
@@ -97,7 +96,6 @@ func NewReconciler(ctx context.Context, client client.Client, indexer client.Fie
 		return nil, err
 	}
 	reconciler = trainJobReconciler{
-		ctx:      ctx,
 		jr:       jobframework.NewReconciler(client, eventRecorder, opts...),
 		client:   client,
 		runtimes: runtimes,
@@ -163,7 +161,7 @@ func (t *TrainJob) PodLabelSelector() string {
 	return fmt.Sprintf("%s=%s", jobsetapi.JobSetNameKey, t.Name)
 }
 
-func getChildJobSet(t *TrainJob) (*jobsetapi.JobSet, error) {
+func getChildJobSet(ctx context.Context, t *TrainJob) (*jobsetapi.JobSet, error) {
 	runtimeRefGK := kftrainerruntime.RuntimeRefToRuntimeRegistryKey(t.Spec.RuntimeRef)
 	runtime, ok := reconciler.runtimes[runtimeRefGK]
 	if !ok {
@@ -171,7 +169,7 @@ func getChildJobSet(t *TrainJob) (*jobsetapi.JobSet, error) {
 	}
 
 	trainJob := (*kftrainer.TrainJob)(t)
-	trSpec, err := getRuntimeSpec(trainJob)
+	trSpec, err := getRuntimeSpec(ctx, trainJob)
 	if err != nil {
 		return nil, fmt.Errorf("runtime '%s' not found", trainJob.Spec.RuntimeRef.Name)
 	}
@@ -189,7 +187,7 @@ func getChildJobSet(t *TrainJob) (*jobsetapi.JobSet, error) {
 		WithSpec(jobSetSpec)).Initializer(trainJob).Trainer(info, trainJob).PodLabels(info.Scheduler.PodLabels).Build()
 
 	// convert to jobset with the defaults set
-	return jobsetApplyToJobset(jobsetApply)
+	return jobsetApplyToJobset(ctx, jobsetApply)
 }
 
 // TODO: Remove mutex once RuntimeInfo gets thread safe - https://github.com/kubeflow/trainer/issues/2873
@@ -199,7 +197,7 @@ func (r *trainJobReconciler) runtimeInfo(runtime kftrainerruntime.Runtime, train
 	return runtime.RuntimeInfo(trainJob, trSpec.Template, trSpec.MLPolicy, trSpec.PodGroupPolicy)
 }
 
-func jobsetApplyToJobset(jobsetApply *jobsetapplyapi.JobSetApplyConfiguration) (*jobsetapi.JobSet, error) {
+func jobsetApplyToJobset(ctx context.Context, jobsetApply *jobsetapplyapi.JobSetApplyConfiguration) (*jobsetapi.JobSet, error) {
 	jsonData, err := json.Marshal(jobsetApply)
 	if err != nil {
 		return nil, err
@@ -213,7 +211,7 @@ func jobsetApplyToJobset(jobsetApply *jobsetapplyapi.JobSetApplyConfiguration) (
 	// Run a dry-run patch to set the jobset defaults
 	// Defaults must be applied here because Kueue later compares podsets to match workloads.
 	// Workloads coming from the API server are already defaulted, so without defaulting this JobSet, matching would fail.
-	if err = reconciler.client.Patch(reconciler.ctx, jobset, client.Apply, &client.PatchOptions{
+	if err = reconciler.client.Patch(ctx, jobset, client.Apply, &client.PatchOptions{
 		FieldManager: "defaulter",
 		DryRun:       []string{metav1.DryRunAll},
 	}); err != nil {
@@ -222,17 +220,17 @@ func jobsetApplyToJobset(jobsetApply *jobsetapplyapi.JobSetApplyConfiguration) (
 	return jobset, nil
 }
 
-func getRuntimeSpec(trainJob *kftrainer.TrainJob) (*kftrainer.TrainingRuntimeSpec, error) {
+func getRuntimeSpec(ctx context.Context, trainJob *kftrainer.TrainJob) (*kftrainer.TrainingRuntimeSpec, error) {
 	if trainjobutil.RuntimeRefIsClusterTrainingRuntime(trainJob.Spec.RuntimeRef) {
 		var ctr kftrainer.ClusterTrainingRuntime
-		err := reconciler.client.Get(reconciler.ctx, client.ObjectKey{Name: trainJob.Spec.RuntimeRef.Name}, &ctr)
+		err := reconciler.client.Get(ctx, client.ObjectKey{Name: trainJob.Spec.RuntimeRef.Name}, &ctr)
 		if err != nil {
 			return nil, err
 		}
 		return &ctr.Spec, nil
 	} else {
 		var tr kftrainer.TrainingRuntime
-		err := reconciler.client.Get(reconciler.ctx, client.ObjectKey{Namespace: trainJob.Namespace, Name: trainJob.Spec.RuntimeRef.Name}, &tr)
+		err := reconciler.client.Get(ctx, client.ObjectKey{Namespace: trainJob.Namespace, Name: trainJob.Spec.RuntimeRef.Name}, &tr)
 		if err != nil {
 			return nil, err
 		}
@@ -240,16 +238,16 @@ func getRuntimeSpec(trainJob *kftrainer.TrainJob) (*kftrainer.TrainingRuntimeSpe
 	}
 }
 
-func (t *TrainJob) PodSets() ([]kueue.PodSet, error) {
-	jobset, err := getChildJobSet(t)
+func (t *TrainJob) PodSets(ctx context.Context) ([]kueue.PodSet, error) {
+	jobset, err := getChildJobSet(ctx, t)
 	if err != nil {
 		return nil, err
 	}
-	return (*workloadjobset.JobSet)(jobset).PodSets()
+	return (*workloadjobset.JobSet)(jobset).PodSets(ctx)
 }
 
-func (t *TrainJob) RunWithPodSetsInfo(podSetsInfo []podset.PodSetInfo) error {
-	jobset, err := getChildJobSet(t)
+func (t *TrainJob) RunWithPodSetsInfo(ctx context.Context, podSetsInfo []podset.PodSetInfo) error {
+	jobset, err := getChildJobSet(ctx, t)
 	if err != nil {
 		return err
 	}
@@ -281,8 +279,7 @@ func (t *TrainJob) RunWithPodSetsInfo(podSetsInfo []podset.PodSetInfo) error {
 		})
 	}
 	// Update the podSpecOverrides while the job is suspended, since is a requirement from the trainjob admission webhook
-	// TODO: Ideally we should be using the parent function context here
-	err = reconciler.client.Update(reconciler.ctx, t.Object())
+	err = reconciler.client.Update(ctx, t.Object())
 	if err != nil {
 		return err
 	}
@@ -329,7 +326,7 @@ func (t *TrainJob) RestorePodSetsInfo(_ []podset.PodSetInfo) bool {
 	return true
 }
 
-func (t *TrainJob) Finished() (message string, success, finished bool) {
+func (t *TrainJob) Finished(ctx context.Context) (message string, success, finished bool) {
 	if c := apimeta.FindStatusCondition(t.Status.Conditions, kftrainer.TrainJobComplete); c != nil && c.Status == metav1.ConditionTrue {
 		return c.Message, true, true
 	}
@@ -339,8 +336,8 @@ func (t *TrainJob) Finished() (message string, success, finished bool) {
 	return message, success, false
 }
 
-func (t *TrainJob) PodsReady() bool {
-	jobset, err := getChildJobSet(t)
+func (t *TrainJob) PodsReady(ctx context.Context) bool {
+	jobset, err := getChildJobSet(ctx, t)
 	if err != nil {
 		return false
 	}
@@ -356,11 +353,11 @@ func (t *TrainJob) PodsReady() bool {
 	return replicas == readyReplicas
 }
 
-func (t *TrainJob) ReclaimablePods() ([]kueue.ReclaimablePod, error) {
+func (t *TrainJob) ReclaimablePods(ctx context.Context) ([]kueue.ReclaimablePod, error) {
 	if len(t.Status.JobsStatus) == 0 {
 		return nil, nil
 	}
-	jobset, err := getChildJobSet(t)
+	jobset, err := getChildJobSet(ctx, t)
 	if err != nil {
 		return nil, err
 	}
diff --git a/pkg/scheduler/scheduler.go b/pkg/scheduler/scheduler.go
index 31bcd2da6..d46ed4544 100644
--- a/pkg/scheduler/scheduler.go
+++ b/pkg/scheduler/scheduler.go
@@ -82,9 +82,6 @@ type Scheduler struct {
 	// schedulingCycle identifies the number of scheduling
 	// attempts since the last restart.
 	schedulingCycle int64
-
-	// Stubs.
-	patchAdmission func(ctx context.Context, original, updated *kueue.Workload) error
 }
 
 type options struct {
@@ -150,7 +147,6 @@ func New(queues *qcache.Manager, cache *schdcache.Cache, cl client.Client, recor
 		clock:                   options.clock,
 		admissionFairSharing:    options.admissionFairSharing,
 	}
-	s.patchAdmission = s.patchAdmissionStatus
 	return s
 }
 
@@ -646,7 +642,9 @@ func (s *Scheduler) admit(ctx context.Context, e *entry, cq *schdcache.ClusterQu
 	}
 
 	s.admissionRoutineWrapper.Run(func() {
-		err := s.patchAdmission(ctx, origWorkload, newWorkload)
+		err := workload.PatchAdmissionStatus(ctx, s.client, origWorkload, s.clock, func() (*kueue.Workload, bool, error) {
+			return newWorkload, true, nil
+		}, workload.WithLoose())
 		if err == nil {
 			// Record metrics and events for quota reservation and admission
 			s.recordWorkloadAdmissionMetrics(newWorkload, e.Obj, admission)
@@ -672,12 +670,6 @@ func (s *Scheduler) admit(ctx context.Context, e *entry, cq *schdcache.ClusterQu
 	return nil
 }
 
-func (s *Scheduler) patchAdmissionStatus(ctx context.Context, wOrig, w *kueue.Workload) error {
-	return workload.PatchAdmissionStatus(ctx, s.client, wOrig, s.clock, func() (*kueue.Workload, bool, error) {
-		return w, true, nil
-	}, workload.WithLoose())
-}
-
 type entryOrdering struct {
 	entries          []entry
 	workloadOrdering workload.Ordering



Test patch:
diff --git a/pkg/controller/jobframework/base_webhook_test.go b/pkg/controller/jobframework/base_webhook_test.go
index 74de632a9..c36f2255e 100644
--- a/pkg/controller/jobframework/base_webhook_test.go
+++ b/pkg/controller/jobframework/base_webhook_test.go
@@ -248,7 +248,7 @@ func TestValidateOnCreate(t *testing.T) {
 				MockJobWithCustomValidation: mocks.NewMockJobWithCustomValidation(mockctrl),
 			}
 			job.MockGenericJob.EXPECT().Object().Return(tc.job).AnyTimes()
-			job.MockJobWithCustomValidation.EXPECT().ValidateOnCreate().Return(tc.customValidationFailure, tc.customValidationError).AnyTimes()
+			job.MockJobWithCustomValidation.EXPECT().ValidateOnCreate(gomock.Any()).Return(tc.customValidationFailure, tc.customValidationError).AnyTimes()
 
 			w := &jobframework.BaseWebhook{
 				FromObject: func(object runtime.Object) jobframework.GenericJob {
@@ -334,7 +334,7 @@ func TestValidateOnUpdate(t *testing.T) {
 				}
 				mj.MockGenericJob.EXPECT().Object().Return(job).AnyTimes()
 				mj.MockGenericJob.EXPECT().IsSuspended().Return(ptr.Deref(job.Spec.Suspend, false)).AnyTimes()
-				mj.MockJobWithCustomValidation.EXPECT().ValidateOnUpdate(gomock.Any()).Return(customValidationFailure, customValidationError).AnyTimes()
+				mj.MockJobWithCustomValidation.EXPECT().ValidateOnUpdate(gomock.Any(), gomock.Any()).Return(customValidationFailure, customValidationError).AnyTimes()
 				return mj
 			}
 
diff --git a/pkg/controller/jobframework/reconciler_test.go b/pkg/controller/jobframework/reconciler_test.go
index aa9bd8fd6..a6d68c01a 100644
--- a/pkg/controller/jobframework/reconciler_test.go
+++ b/pkg/controller/jobframework/reconciler_test.go
@@ -157,8 +157,8 @@ func TestReconcileGenericJob(t *testing.T) {
 			mgj.EXPECT().GVK().Return(testGVK).AnyTimes()
 			mgj.EXPECT().IsSuspended().Return(ptr.Deref(tc.job.Spec.Suspend, false)).AnyTimes()
 			mgj.EXPECT().IsActive().Return(tc.job.Status.Active != 0).AnyTimes()
-			mgj.EXPECT().Finished().Return("", false, false).AnyTimes()
-			mgj.EXPECT().PodSets().Return(tc.podSets, nil).AnyTimes()
+			mgj.EXPECT().Finished(gomock.Any()).Return("", false, false).AnyTimes()
+			mgj.EXPECT().PodSets(gomock.Any()).Return(tc.podSets, nil).AnyTimes()
 
 			cl := utiltesting.NewClientBuilder(batchv1.AddToScheme, kueue.AddToScheme).
 				WithObjects(tc.objs...).
diff --git a/pkg/controller/jobs/appwrapper/appwrapper_controller_test.go b/pkg/controller/jobs/appwrapper/appwrapper_controller_test.go
index 43bf522b4..a310e0ddd 100644
--- a/pkg/controller/jobs/appwrapper/appwrapper_controller_test.go
+++ b/pkg/controller/jobs/appwrapper/appwrapper_controller_test.go
@@ -155,7 +155,8 @@ func TestPodSets(t *testing.T) {
 	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()
+			ctx, _ := utiltesting.ContextWithLog(t)
+			gotPodSets, err := fromObject(tc.job).PodSets(ctx)
 			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 a4a25aaba..a0a71cbf7 100644
--- a/pkg/controller/jobs/job/job_controller_test.go
+++ b/pkg/controller/jobs/job/job_controller_test.go
@@ -164,7 +164,8 @@ func TestPodsReady(t *testing.T) {
 
 	for name, tc := range testcases {
 		t.Run(name, func(t *testing.T) {
-			got := tc.job.PodsReady()
+			ctx, _ := utiltesting.ContextWithLog(t)
+			got := tc.job.PodsReady(ctx)
 			if tc.want != got {
 				t.Errorf("Unexpected response (want: %v, got: %v)", tc.want, got)
 			}
@@ -307,9 +308,10 @@ func TestPodSetsInfo(t *testing.T) {
 	}
 	for name, tc := range testcases {
 		t.Run(name, func(t *testing.T) {
+			ctx, _ := utiltesting.ContextWithLog(t)
 			origSpec := *tc.job.Spec.DeepCopy()
 
-			gotErr := tc.job.RunWithPodSetsInfo(tc.runInfo)
+			gotErr := tc.job.RunWithPodSetsInfo(ctx, tc.runInfo)
 
 			if diff := cmp.Diff(tc.wantRunError, gotErr, cmpopts.EquateErrors()); diff != "" {
 				t.Errorf("node selectors mismatch (-want +got):\n%s", diff)
@@ -504,7 +506,8 @@ func TestPodSets(t *testing.T) {
 	for name, tc := range cases {
 		t.Run(name, func(t *testing.T) {
 			features.SetFeatureGateDuringTest(t, features.TopologyAwareScheduling, tc.enableTopologyAwareScheduling)
-			gotPodSets, err := tc.job.PodSets()
+			ctx, _ := utiltesting.ContextWithLog(t)
+			gotPodSets, err := tc.job.PodSets(ctx)
 			if err != nil {
 				t.Fatalf("unexpected error: %v", err)
 			}
diff --git a/pkg/controller/jobs/job/job_multikueue_adapter_test.go b/pkg/controller/jobs/job/job_multikueue_adapter_test.go
index 896e2b353..7db07fde0 100644
--- a/pkg/controller/jobs/job/job_multikueue_adapter_test.go
+++ b/pkg/controller/jobs/job/job_multikueue_adapter_test.go
@@ -692,7 +692,7 @@ func Test_multiKueueAdapter_SyncJob(t *testing.T) {
 			}
 
 			adapter := &multiKueueAdapter{}
-			ctx := t.Context()
+			ctx, _ := utiltesting.ContextWithLog(t)
 
 			// Function call under test with result (error) assertion.
 			if err := adapter.SyncJob(ctx, tt.args.localClient, tt.args.remoteClient, tt.args.key, tt.args.workloadName, tt.args.origin); (err != nil) != tt.want.err {
diff --git a/pkg/controller/jobs/job/job_webhook_test.go b/pkg/controller/jobs/job/job_webhook_test.go
index 42d66b3cf..2242eaf80 100644
--- a/pkg/controller/jobs/job/job_webhook_test.go
+++ b/pkg/controller/jobs/job/job_webhook_test.go
@@ -393,7 +393,8 @@ func TestValidateCreate(t *testing.T) {
 
 			jw := &JobWebhook{}
 
-			gotValidationErrs, gotErr := jw.validateCreate((*Job)(tc.job))
+			ctx, _ := utiltesting.ContextWithLog(t)
+			gotValidationErrs, gotErr := jw.validateCreate(ctx, (*Job)(tc.job))
 
 			if diff := cmp.Diff(tc.wantErr, gotErr); diff != "" {
 				t.Errorf("validateCreate() error mismatch (-want +got):\n%s", diff)
@@ -688,8 +689,8 @@ func TestValidateUpdate(t *testing.T) {
 	for _, tc := range testcases {
 		t.Run(tc.name, func(t *testing.T) {
 			features.SetFeatureGateDuringTest(t, features.TopologyAwareScheduling, tc.topologyAwareScheduling)
-
-			gotValidationErrs, gotErr := new(JobWebhook).validateUpdate((*Job)(tc.oldJob), (*Job)(tc.newJob))
+			ctx, _ := utiltesting.ContextWithLog(t)
+			gotValidationErrs, gotErr := new(JobWebhook).validateUpdate(ctx, (*Job)(tc.oldJob), (*Job)(tc.newJob))
 			if diff := cmp.Diff(tc.wantErr, gotErr, cmpopts.IgnoreFields(field.Error{})); diff != "" {
 				t.Errorf("validateUpdate() error mismatch (-want +got):\n%s", diff)
 			}
diff --git a/pkg/controller/jobs/jobset/jobset_controller_test.go b/pkg/controller/jobs/jobset/jobset_controller_test.go
index 2ff425505..fb1d1837b 100644
--- a/pkg/controller/jobs/jobset/jobset_controller_test.go
+++ b/pkg/controller/jobs/jobset/jobset_controller_test.go
@@ -106,7 +106,8 @@ func TestPodsReady(t *testing.T) {
 	for name, tc := range testcases {
 		t.Run(name, func(t *testing.T) {
 			jobSet := (JobSet)(tc.jobSet)
-			got := jobSet.PodsReady()
+			ctx, _ := utiltesting.ContextWithLog(t)
+			got := jobSet.PodsReady(ctx)
 			if tc.want != got {
 				t.Errorf("Unexpected response (want: %v, got: %v)", tc.want, got)
 			}
@@ -188,8 +189,9 @@ func TestReclaimablePods(t *testing.T) {
 
 	for name, tc := range testcases {
 		t.Run(name, func(t *testing.T) {
+			ctx, _ := utiltesting.ContextWithLog(t)
 			jobSet := (*JobSet)(tc.jobSet)
-			got, err := jobSet.ReclaimablePods()
+			got, err := jobSet.ReclaimablePods(ctx)
 			if err != nil {
 				t.Fatalf("Unexpected error: %s", err)
 			}
@@ -342,7 +344,8 @@ func TestPodSets(t *testing.T) {
 	for name, tc := range testCases {
 		t.Run(name, func(t *testing.T) {
 			features.SetFeatureGateDuringTest(t, features.TopologyAwareScheduling, tc.enableTopologyAwareScheduling)
-			gotPodSets, err := tc.jobSet.PodSets()
+			ctx, _ := utiltesting.ContextWithLog(t)
+			gotPodSets, err := tc.jobSet.PodSets(ctx)
 			if err != nil {
 				t.Fatalf("unexpected error: %s", err)
 			}
diff --git a/pkg/controller/jobs/jobset/jobset_webhook_test.go b/pkg/controller/jobs/jobset/jobset_webhook_test.go
index 465eea294..95d2cf6bc 100644
--- a/pkg/controller/jobs/jobset/jobset_webhook_test.go
+++ b/pkg/controller/jobs/jobset/jobset_webhook_test.go
@@ -181,8 +181,8 @@ func TestValidateUpdate(t *testing.T) {
 	for _, tc := range testcases {
 		t.Run(tc.name, func(t *testing.T) {
 			features.SetFeatureGateDuringTest(t, features.TopologyAwareScheduling, tc.topologyAwareScheduling)
-
-			gotValidationErrs, gotErr := new(JobSetWebhook).validateUpdate((*JobSet)(tc.oldJob), (*JobSet)(tc.newJob))
+			ctx, _ := utiltesting.ContextWithLog(t)
+			gotValidationErrs, gotErr := new(JobSetWebhook).validateUpdate(ctx, (*JobSet)(tc.oldJob), (*JobSet)(tc.newJob))
 			if diff := cmp.Diff(tc.wantErr, gotErr, cmpopts.IgnoreFields(field.Error{})); diff != "" {
 				t.Errorf("validateUpdate() error mismatch (-want +got):\n%s", diff)
 			}
diff --git a/pkg/controller/jobs/kubeflow/jobs/jaxjob/jaxjob_controller_test.go b/pkg/controller/jobs/kubeflow/jobs/jaxjob/jaxjob_controller_test.go
index 56687b869..535927eeb 100644
--- a/pkg/controller/jobs/kubeflow/jobs/jaxjob/jaxjob_controller_test.go
+++ b/pkg/controller/jobs/kubeflow/jobs/jaxjob/jaxjob_controller_test.go
@@ -164,7 +164,8 @@ func TestPodSets(t *testing.T) {
 			for _, gate := range tc.enableFeatureGates {
 				features.SetFeatureGateDuringTest(t, gate, true)
 			}
-			gotPodSets, err := fromObject(tc.job).PodSets()
+			ctx, _ := utiltesting.ContextWithLog(t)
+			gotPodSets, err := fromObject(tc.job).PodSets(ctx)
 			if err != nil {
 				t.Fatalf("unexpected error: %v", err)
 			}
@@ -260,7 +261,8 @@ func TestValidate(t *testing.T) {
 		t.Run(name, func(t *testing.T) {
 			features.SetFeatureGateDuringTest(t, features.TopologyAwareScheduling, tc.topologyAwareScheduling)
 
-			gotValidationErrs, gotErr := fromObject(tc.job).ValidateOnCreate()
+			ctx, _ := utiltesting.ContextWithLog(t)
+			gotValidationErrs, gotErr := fromObject(tc.job).ValidateOnCreate(ctx)
 			if diff := cmp.Diff(tc.wantErr, gotErr); diff != "" {
 				t.Errorf("validate create error list mismatch (-want +got):\n%s", diff)
 			}
@@ -268,7 +270,7 @@ func TestValidate(t *testing.T) {
 				t.Errorf("validate create validation errors list mismatch (-want +got):\n%s", diff)
 			}
 
-			gotValidationErrs, gotErr = fromObject(tc.job).ValidateOnUpdate(nil)
+			gotValidationErrs, gotErr = fromObject(tc.job).ValidateOnUpdate(ctx, nil)
 			if diff := cmp.Diff(tc.wantErr, gotErr); diff != "" {
 				t.Errorf("validate create error list mismatch (-want +got):\n%s", diff)
 			}
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 e9ae04bd0..194e4d627 100644
--- a/pkg/controller/jobs/kubeflow/jobs/paddlejob/paddlejob_controller_test.go
+++ b/pkg/controller/jobs/kubeflow/jobs/paddlejob/paddlejob_controller_test.go
@@ -337,7 +337,8 @@ func TestPodSets(t *testing.T) {
 	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()
+			ctx, _ := utiltesting.ContextWithLog(t)
+			gotPodSets, err := fromObject(tc.job).PodSets(ctx)
 			if err != nil {
 				t.Fatalf("unexpected error: %v", err)
 			}
@@ -464,7 +465,8 @@ func TestValidate(t *testing.T) {
 		t.Run(name, func(t *testing.T) {
 			features.SetFeatureGateDuringTest(t, features.TopologyAwareScheduling, tc.topologyAwareScheduling)
 
-			gotValidationErrs, gotErr := fromObject(tc.job).ValidateOnCreate()
+			ctx, _ := utiltesting.ContextWithLog(t)
+			gotValidationErrs, gotErr := fromObject(tc.job).ValidateOnCreate(ctx)
 			if diff := cmp.Diff(tc.wantErr, gotErr); diff != "" {
 				t.Errorf("validate create error mismatch (-want +got):\n%s", diff)
 			}
@@ -472,7 +474,7 @@ func TestValidate(t *testing.T) {
 				t.Errorf("validate create validation errors list mismatch (-want +got):\n%s", diff)
 			}
 
-			gotValidationErrs, gotErr = fromObject(tc.job).ValidateOnUpdate(nil)
+			gotValidationErrs, gotErr = fromObject(tc.job).ValidateOnUpdate(ctx, nil)
 			if diff := cmp.Diff(tc.wantErr, gotErr); diff != "" {
 				t.Errorf("validate create error mismatch (-want +got):\n%s", diff)
 			}
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 f8703b87e..f09f3e8d5 100644
--- a/pkg/controller/jobs/kubeflow/jobs/pytorchjob/pytorchjob_controller_test.go
+++ b/pkg/controller/jobs/kubeflow/jobs/pytorchjob/pytorchjob_controller_test.go
@@ -328,7 +328,8 @@ func TestPodSets(t *testing.T) {
 	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()
+			ctx, _ := utiltesting.ContextWithLog(t)
+			gotPodSets, err := fromObject(tc.job).PodSets(ctx)
 			if err != nil {
 				t.Fatalf("unexpected error: %v", err)
 			}
@@ -466,7 +467,8 @@ func TestValidate(t *testing.T) {
 		t.Run(name, func(t *testing.T) {
 			features.SetFeatureGateDuringTest(t, features.TopologyAwareScheduling, tc.topologyAwareScheduling)
 
-			gotValidationErrs, gotErr := fromObject(tc.job).ValidateOnCreate()
+			ctx, _ := utiltesting.ContextWithLog(t)
+			gotValidationErrs, gotErr := fromObject(tc.job).ValidateOnCreate(ctx)
 			if diff := cmp.Diff(tc.wantErr, gotErr); diff != "" {
 				t.Errorf("validate create error mismatch (-want +got):\n%s", diff)
 			}
@@ -474,7 +476,7 @@ func TestValidate(t *testing.T) {
 				t.Errorf("validate create validation errors list mismatch (-want +got):\n%s", diff)
 			}
 
-			gotValidationErrs, gotErr = fromObject(tc.job).ValidateOnUpdate(nil)
+			gotValidationErrs, gotErr = fromObject(tc.job).ValidateOnUpdate(ctx, nil)
 			if diff := cmp.Diff(tc.wantErr, gotErr); diff != "" {
 				t.Errorf("validate create error mismatch (-want +got):\n%s", diff)
 			}
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 f43491e61..633fa2c19 100644
--- a/pkg/controller/jobs/kubeflow/jobs/tfjob/tfjob_controller_test.go
+++ b/pkg/controller/jobs/kubeflow/jobs/tfjob/tfjob_controller_test.go
@@ -353,7 +353,8 @@ func TestPodSets(t *testing.T) {
 	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()
+			ctx, _ := utiltesting.ContextWithLog(t)
+			gotPodSets, err := fromObject(tc.job).PodSets(ctx)
 			if err != nil {
 				t.Fatalf("unexpected error: %v", err)
 			}
@@ -507,7 +508,8 @@ func TestValidate(t *testing.T) {
 		t.Run(name, func(t *testing.T) {
 			features.SetFeatureGateDuringTest(t, features.TopologyAwareScheduling, tc.topologyAwareScheduling)
 
-			gotValidationErrs, gotErr := fromObject(tc.job).ValidateOnCreate()
+			ctx, _ := utiltesting.ContextWithLog(t)
+			gotValidationErrs, gotErr := fromObject(tc.job).ValidateOnCreate(ctx)
 			if diff := cmp.Diff(tc.wantErr, gotErr); diff != "" {
 				t.Errorf("validate create error mismatch (-want +got):\n%s", diff)
 			}
@@ -515,7 +517,7 @@ func TestValidate(t *testing.T) {
 				t.Errorf("validate create validation error list mismatch (-want +got):\n%s", diff)
 			}
 
-			gotValidationErrs, gotErr = fromObject(tc.job).ValidateOnUpdate(nil)
+			gotValidationErrs, gotErr = fromObject(tc.job).ValidateOnUpdate(ctx, nil)
 			if diff := cmp.Diff(tc.wantErr, gotErr); diff != "" {
 				t.Errorf("validate create error mismatch (-want +got):\n%s", diff)
 			}
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 c40127dbd..21aa83ed4 100644
--- a/pkg/controller/jobs/kubeflow/jobs/xgboostjob/xgboostjob_controller_test.go
+++ b/pkg/controller/jobs/kubeflow/jobs/xgboostjob/xgboostjob_controller_test.go
@@ -335,7 +335,8 @@ func TestPodSets(t *testing.T) {
 	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()
+			ctx, _ := utiltesting.ContextWithLog(t)
+			gotPodSets, err := fromObject(tc.job).PodSets(ctx)
 			if err != nil {
 				t.Fatalf("unexpected error: %v", err)
 			}
@@ -473,7 +474,8 @@ func TestValidate(t *testing.T) {
 		t.Run(name, func(t *testing.T) {
 			features.SetFeatureGateDuringTest(t, features.TopologyAwareScheduling, tc.topologyAwareScheduling)
 
-			gotValidationErrs, gotErr := fromObject(tc.job).ValidateOnCreate()
+			ctx, _ := utiltesting.ContextWithLog(t)
+			gotValidationErrs, gotErr := fromObject(tc.job).ValidateOnCreate(ctx)
 			if diff := cmp.Diff(tc.wantErr, gotErr); diff != "" {
 				t.Errorf("validate create error mismatch (-want +got):\n%s", diff)
 			}
@@ -481,7 +483,7 @@ func TestValidate(t *testing.T) {
 				t.Errorf("validate create validation error list mismatch (-want +got):\n%s", diff)
 			}
 
-			gotValidationErrs, gotErr = fromObject(tc.job).ValidateOnUpdate(nil)
+			gotValidationErrs, gotErr = fromObject(tc.job).ValidateOnUpdate(ctx, nil)
 			if diff := cmp.Diff(tc.wantErr, gotErr); diff != "" {
 				t.Errorf("validate create error mismatch (-want +got):\n%s", diff)
 			}
diff --git a/pkg/controller/jobs/mpijob/mpijob_controller_test.go b/pkg/controller/jobs/mpijob/mpijob_controller_test.go
index 8f92ef5ec..368a79818 100644
--- a/pkg/controller/jobs/mpijob/mpijob_controller_test.go
+++ b/pkg/controller/jobs/mpijob/mpijob_controller_test.go
@@ -328,7 +328,8 @@ func TestPodSets(t *testing.T) {
 	for name, tc := range testCases {
 		t.Run(name, func(t *testing.T) {
 			features.SetFeatureGateDuringTest(t, features.TopologyAwareScheduling, tc.enableTopologyAwareScheduling)
-			gotPodSets, err := tc.job.PodSets()
+			ctx, _ := utiltesting.ContextWithLog(t)
+			gotPodSets, err := tc.job.PodSets(ctx)
 			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 252d8b664..ddbe93588 100644
--- a/pkg/controller/jobs/pod/pod_controller_test.go
+++ b/pkg/controller/jobs/pod/pod_controller_test.go
@@ -89,7 +89,8 @@ func TestPodsReady(t *testing.T) {
 	for name, tc := range testCases {
 		t.Run(name, func(t *testing.T) {
 			pod := FromObject(tc.pod)
-			got := pod.PodsReady()
+			ctx, _ := utiltesting.ContextWithLog(t)
+			got := pod.PodsReady(ctx)
 			if tc.want != got {
 				t.Errorf("Unexpected response (want: %v, got: %v)", tc.want, got)
 			}
@@ -212,7 +213,8 @@ func TestPodSets(t *testing.T) {
 	for name, tc := range testCases {
 		t.Run(name, func(t *testing.T) {
 			features.SetFeatureGateDuringTest(t, features.TopologyAwareScheduling, tc.enableTopologyAwareScheduling)
-			gotPodSets, err := tc.pod.PodSets()
+			ctx, _ := utiltesting.ContextWithLog(t)
+			gotPodSets, err := tc.pod.PodSets(ctx)
 			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 68e68eece..b1c16eb19 100644
--- a/pkg/controller/jobs/raycluster/raycluster_controller_test.go
+++ b/pkg/controller/jobs/raycluster/raycluster_controller_test.go
@@ -283,7 +283,8 @@ func TestPodSets(t *testing.T) {
 	for name, tc := range testCases {
 		t.Run(name, func(t *testing.T) {
 			features.SetFeatureGateDuringTest(t, features.TopologyAwareScheduling, tc.enableTopologyAwareScheduling)
-			gotPodSets, err := tc.rayCluster.PodSets()
+			ctx, _ := utiltesting.ContextWithLog(t)
+			gotPodSets, err := tc.rayCluster.PodSets(ctx)
 			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 843598833..36ed162ca 100644
--- a/pkg/controller/jobs/rayjob/rayjob_controller_test.go
+++ b/pkg/controller/jobs/rayjob/rayjob_controller_test.go
@@ -520,7 +520,8 @@ func TestPodSets(t *testing.T) {
 	for name, tc := range testCases {
 		t.Run(name, func(t *testing.T) {
 			features.SetFeatureGateDuringTest(t, features.TopologyAwareScheduling, tc.enableTopologyAwareScheduling)
-			gotPodSets, err := tc.rayJob.PodSets()
+			ctx, _ := utiltesting.ContextWithLog(t)
+			gotPodSets, err := tc.rayJob.PodSets(ctx)
 			if err != nil {
 				t.Fatalf("unexpected error: %v", err)
 			}
@@ -656,8 +657,9 @@ func TestNodeSelectors(t *testing.T) {
 
 	for name, tc := range cases {
 		t.Run(name, func(t *testing.T) {
+			ctx, _ := utiltesting.ContextWithLog(t)
 			genJob := (*RayJob)(tc.job)
-			gotRunError := genJob.RunWithPodSetsInfo(tc.runInfo)
+			gotRunError := genJob.RunWithPodSetsInfo(ctx, tc.runInfo)
 
 			if diff := cmp.Diff(tc.wantRunError, gotRunError, cmpopts.EquateErrors()); diff != "" {
 				t.Errorf("Unexpected run error (-want/+got): %s", diff)
@@ -733,10 +735,11 @@ func Test_RayJobFinished(t *testing.T) {
 
 	for _, testcase := range testcases {
 		t.Run(testcase.name, func(t *testing.T) {
+			ctx, _ := utiltesting.ContextWithLog(t)
 			rayJob := testingrayutil.MakeJob("job", "ns").Obj()
 			rayJob.Status = testcase.status
 
-			_, success, finished := ((*RayJob)(rayJob)).Finished()
+			_, success, finished := ((*RayJob)(rayJob)).Finished(ctx)
 			if success != testcase.expectedSuccess {
 				t.Logf("actual success: %v", success)
 				t.Logf("expected success: %v", testcase.expectedSuccess)
diff --git a/pkg/controller/jobs/trainjob/trainjob_controller_test.go b/pkg/controller/jobs/trainjob/trainjob_controller_test.go
index 56c2011a6..b83e87576 100644
--- a/pkg/controller/jobs/trainjob/trainjob_controller_test.go
+++ b/pkg/controller/jobs/trainjob/trainjob_controller_test.go
@@ -194,7 +194,7 @@ func TestRunWithPodsetsInfo(t *testing.T) {
 			}
 
 			kTrainJob := (*TrainJob)(tc.trainJob)
-			err = kTrainJob.RunWithPodSetsInfo(tc.podsetsInfo)
+			err = kTrainJob.RunWithPodSetsInfo(ctx, tc.podsetsInfo)
 			if err != nil {
 				if !tc.wantErr {
 					t.Errorf("unexpected RunWithPodSetsInfo() error: %v", err)
diff --git a/pkg/controller/jobs/trainjob/trainjob_webhook_test.go b/pkg/controller/jobs/trainjob/trainjob_webhook_test.go
index 7a3334c76..8941cbf1a 100644
--- a/pkg/controller/jobs/trainjob/trainjob_webhook_test.go
+++ b/pkg/controller/jobs/trainjob/trainjob_webhook_test.go
@@ -67,7 +67,8 @@ func TestValidateCreate(t *testing.T) {
 	for _, tc := range testcases {
 		t.Run(tc.name, func(t *testing.T) {
 			webhook := &TrainJobWebhook{}
-			_, gotErr := webhook.ValidateCreate(t.Context(), tc.trainJob)
+			ctx, _ := utiltesting.ContextWithLog(t)
+			_, gotErr := webhook.ValidateCreate(ctx, tc.trainJob)
 
 			if diff := cmp.Diff(tc.wantErr, gotErr); diff != "" {
 				t.Errorf("validateCreate() mismatch (-want +got):\n%s", diff)
diff --git a/pkg/scheduler/scheduler_tas_test.go b/pkg/scheduler/scheduler_tas_test.go
index 49bc07cdf..93c477575 100644
--- a/pkg/scheduler/scheduler_tas_test.go
+++ b/pkg/scheduler/scheduler_tas_test.go
@@ -2442,14 +2442,6 @@ func TestScheduleForTAS(t *testing.T) {
 				}
 			}
 			scheduler := New(qManager, cqCache, cl, recorder)
-			gotScheduled := make([]workload.Reference, 0)
-			var mu sync.Mutex
-			scheduler.patchAdmission = func(ctx context.Context, wOrig, w *kueue.Workload) error {
-				mu.Lock()
-				gotScheduled = append(gotScheduled, workload.Key(w))
-				mu.Unlock()
-				return nil
-			}
 			wg := sync.WaitGroup{}
 			scheduler.setAdmissionRoutineWrapper(routine.NewWrapper(
 				func() { wg.Add(1) },
@@ -2978,20 +2970,13 @@ func TestScheduleForTASPreemption(t *testing.T) {
 				}
 			}
 			scheduler := New(qManager, cqCache, cl, recorder)
-			gotScheduled := make([]workload.Reference, 0)
-			var mu sync.Mutex
-			scheduler.patchAdmission = func(ctx context.Context, wOrig, w *kueue.Workload) error {
-				mu.Lock()
-				gotScheduled = append(gotScheduled, workload.Key(w))
-				mu.Unlock()
-				return nil
-			}
 			wg := sync.WaitGroup{}
 			scheduler.setAdmissionRoutineWrapper(routine.NewWrapper(
 				func() { wg.Add(1) },
 				func() { wg.Done() },
 			))
 
+			var mu sync.Mutex
 			gotPreempted := sets.New[workload.Reference]()
 			scheduler.preemptor.OverrideApply(func(_ context.Context, w *kueue.Workload, _, _ string) error {
 				mu.Lock()
@@ -3950,20 +3935,13 @@ func TestScheduleForTASCohorts(t *testing.T) {
 				}
 			}
 			scheduler := New(qManager, cqCache, cl, recorder)
-			gotScheduled := make([]workload.Reference, 0)
-			var mu sync.Mutex
-			scheduler.patchAdmission = func(ctx context.Context, wOrig, w *kueue.Workload) error {
-				mu.Lock()
-				gotScheduled = append(gotScheduled, workload.Key(w))
-				mu.Unlock()
-				return nil
-			}
 			wg := sync.WaitGroup{}
 			scheduler.setAdmissionRoutineWrapper(routine.NewWrapper(
 				func() { wg.Add(1) },
 				func() { wg.Done() },
 			))
 
+			var mu sync.Mutex
 			gotPreempted := sets.New[workload.Reference]()
 			scheduler.preemptor.OverrideApply(func(_ context.Context, w *kueue.Workload, _, _ string) error {
 				mu.Lock()
diff --git a/pkg/util/client/client_test.go b/pkg/util/client/client_test.go
index 81e291b6a..7cb8e2506 100644
--- a/pkg/util/client/client_test.go
+++ b/pkg/util/client/client_test.go
@@ -164,7 +164,7 @@ func TestPatch(t *testing.T) {
 	}
 	for name, tt := range tests {
 		t.Run(name, func(t *testing.T) {
-			ctx := t.Context()
+			ctx, _ := utiltesting.ContextWithLog(t)
 			clnt := utiltesting.NewClientBuilder().WithObjects(clientObject).Build()
 			if err := Patch(ctx, clnt, tt.args.obj, tt.args.update, tt.args.options...); (err != nil) != tt.want.err {
 				t.Errorf("Patch() error = %v, wantErr %v", err, tt.want.err)
@@ -289,7 +289,7 @@ func TestPatchStatus(t *testing.T) {
 	}
 	for name, tt := range tests {
 		t.Run(name, func(t *testing.T) {
-			ctx := t.Context()
+			ctx, _ := utiltesting.ContextWithLog(t)
 			clnt := utiltesting.NewClientBuilder().WithObjects(clientObject).Build()
 			if err := PatchStatus(ctx, clnt, tt.args.obj, tt.args.update, tt.args.options...); (err != nil) != tt.want.err {
 				t.Errorf("Patch() error = %v, wantErr %v", err, tt.want.err)
diff --git a/test/integration/singlecluster/scheduler/fairsharing/fair_sharing_test.go b/test/integration/singlecluster/scheduler/fairsharing/fair_sharing_test.go
index 852c46227..0e917af26 100644
--- a/test/integration/singlecluster/scheduler/fairsharing/fair_sharing_test.go
+++ b/test/integration/singlecluster/scheduler/fairsharing/fair_sharing_test.go
@@ -839,7 +839,7 @@ var _ = ginkgo.Describe("Scheduler", ginkgo.Ordered, ginkgo.ContinueOnFailure, f
 			for range 10 {
 				createWorkload("a", "100")
 			}
-			util.ExpectReservingActiveWorkloadsMetric(cqA, 10)
+			util.ExpectAdmittedWorkloadsTotalMetric(cqA, "", 10)
 			ginkgo.By("Creating a newer workload in cqB that needs only nominal quota")
 			createWorkload("b", "500")
 			ginkgo.By("Evict the some workloads in cqA and reclaim the nominal quota in cqB")