Repository: kubernetes-sigs/kueue

Issue:
Replace Scheduler stubs with interceptor functions
**What would you like to be cleaned**:
As mention here: https://github.com/kubernetes-sigs/kueue/pull/6962#discussion_r2372511842

Investigation a possibility to drop stubs of Scheduler, that could be possibly replaced with proper interceptor functions of FakeClient.

**Why is this needed**:
Scheduler stubs add extra layer of indirection that potentially could be avoided.


Discussion:
@mbobrovskyi: /assign@mimowo: Awesome, thank you for working on that. This stub pattern is legacy@mbobrovskyi: /reopen@k8s-ci-robot: @mbobrovskyi: Reopened this issue.

<details>

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

>/reopen


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

Let me know if I missed something.@mbobrovskyi: /reopen

Still need to remove preemption stubs:

https://github.com/kubernetes-sigs/kueue/blob/178155c7eadd36393cecb5da349995c946918b78/pkg/scheduler/preemption/preemption.go#L62-L63@k8s-ci-robot: @mbobrovskyi: Reopened this issue.

<details>

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

>/reopen
>
>Still need to remove preemption stubs:
>
>https://github.com/kubernetes-sigs/kueue/blob/178155c7eadd36393cecb5da349995c946918b78/pkg/scheduler/preemption/preemption.go#L62-L63


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

Solution patch:
diff --git a/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)